一、Transformer 大語言模型的 SOTA 訓練技術

1、大語言模型的發(fā)展背景與挑戰(zhàn)

首先和大家分享下大語言模型的發(fā)展背景。過去的幾年內(nèi)，無論從數(shù)據(jù)量的維度還是從模型規(guī)模的維度，語言模型規(guī)模都擴大了非常多的數(shù)量級。隨著數(shù)據(jù)量和模型規(guī)模的擴大，也面臨著一些挑戰(zhàn)。

其中涉及的挑戰(zhàn)主要可以分為兩部分。首先是計算方面，這里給了一個來自于Megatron 論文的公式去計算一個模型訓練時需要的計算 FLOPS，我們可以簡單推算一下，GPT-3 175B 模型使用現(xiàn)在比較合理的 1.5T Tokens 數(shù)據(jù)量訓練，大概需要 128 個 DGX A100 節(jié)點(*僅供技術交流使用），共計 1024 張 A100 卡(*僅供技術交流使用），在效率比較高的條件下連續(xù)訓練 120 天。而且這還是一個非常理想的估計值。因為在實際訓練時候不可避免地會遇到一些例如 Checkpoint、Save 與 Load 的時間、節(jié)點崩潰及重啟時間，或者我們需要取 Debug loss 曲線不正常的時間，因此實際的訓練時間會遠超過這個理想估計時間。

其次是運行大語言模型對顯存開銷的挑戰(zhàn)。例如對 GPT-3 175B 模型訓練時的顯存開銷主要可以分為兩部分，第一部分是這個模型的狀態(tài)所占的顯存，包括模型參數(shù)量、梯度和優(yōu)化器所占的顯存。其中，主要的顯存開銷是優(yōu)化器狀態(tài)部分，也就是我們用 Adam 優(yōu)化器會涉及到的 Momentum 和 Variance 等。另一部分是 Activation 所占的顯存，后面會具體地推算這個 Activation 實際會占多大顯存，以及用了我們的一些優(yōu)化手段后， Activation 會占多少顯存。但無論如何，這兩個部分顯存加起來后會是一個巨大的顯存占用量，我們需要對模型進行合理的切分，把它分布到不同的 GPU 上之后才能將模型運行起來。

今天我們也就圍繞上述的這兩點挑戰(zhàn)展開，一個是我們?nèi)绾伟涯Ｐ瓦\行起來，也就是解決顯存開銷的挑戰(zhàn)；另一個是我們?nèi)绾伟涯Ｐ瓦\行得足夠快，也就是如何選擇合理的并行手段和優(yōu)化手段。

這里介紹一下 NeMO 和 Megatron，因為我們接下來介紹的技術都是基于 NeMO 和 Megatron 的。NeMO 和 Megatron 是 NVIDIA 對于大語言模型的兩個框架，NeMO 是偏產(chǎn)品化的框架，有更好的支持；Megatron-LM 是一個比較經(jīng)典的，偏學術研究的項目。NVIDIA 在這個框架做了大量的優(yōu)化來保證實際的模型訓練性能處于 SOTA 狀態(tài)。

2、大語言模型的 SOTA 訓練技術

接下來介紹下本次分享的大語言模型訓練技術的總體概覽。下圖中標記綠色的是去年或今年新出來的一些技術，而標記藍色的是比較經(jīng)典的技術。并行部分的介紹包括 Pipeline 并行、Tensor 并行和 Megatron v3 版本新出的 Sequence 并行。Expert 并行是針對 MOE 模型的一種并行，本文中不做展開。

第二個部分是顯存優(yōu)化，包括之前常聽到的 ZeRO，在 Megatron 和 NeMO 中已經(jīng)實現(xiàn)了 ZeRO-1，即 Distributed Optimizer；和Checkpoint Activations，其中 Selective Activation Checkpoint 是目前 Megatron 中用的比較多的，也是比較高效的一個 Checkpoint 技術；其他的包括我們比較常用的，例如混合精度訓練，Kernel Fusion 和 Flash Attention 等，以及在最新的 Megatron 和 NeMO 上更新的 Pipeline 通信和 Tensor 通信等。

二、以 GPT-175B 為例，最新訓練技術量化分析

1、GPT-3 模型分布式訓練技術：模型并行

首先簡單地介紹下一些經(jīng)典的并行化技術。例如，Tensor 模型并行（TP）和Pipeline 模型并行（PP）。

其中，TP 是對模型的每個層做了一個層內(nèi)的拆分。目前 TP 只用于 Linear 層，當 Linear 層的 Shape 比較大時，使用 TP 是能達到很好的 GPU 利用率的。但 TP 也有一個缺點，即對每個 Linear 層做拆分，因此它的通信粒度是非常細的。這就意味著，在不做優(yōu)化的條件下，TP 每計算一次層的拆分，就需要有一次通信來做它的 Activation 合并操作，雖然 TP 單次通信量較小，但是它通信頻率頻次都很高，對帶寬的要求也很高。

PP 是流水線并行，也就是一個層間的拆分，把不同的層放到不同的 GPU 上。PP 的切分粒度比 TP 更粗，因此 PP 的通信量遠小于 TP 的通信量，但是 PP 也存在一些缺點。一個完整的 Pipeline 運行起來需要將一個很大的 workload 切分成很小的多個 Workload，也就是需要將一個比較大 Batch size 切分成很多個小 Batch 才能保持流水線并行的高吞吐。并且如果在流水線拆分時候不同的 GPU 上的 Workload 不均勻會導致整個流水線的效率非常低。

下圖是一個 TP 的示意圖。MLP 基于 TP 做了拆分，由兩個 FFN 組成。在做 TP 拆分時做了一次優(yōu)化，也就是對第一個 Linear 層做列拆分，對第二個 Linear 層做行拆分，因此在每次 Forward 或 Backward 時只需要一次通信就能完成兩個 Linear 層的 TP 拆分。但在未做優(yōu)化時，需要兩次通信完成 TP 拆分。具體而言，也就是在 Forward 計算時，會在 g 操作處做一次 All-reduce 操作，實現(xiàn) Activations 的一次平均；相應地，在 Backward 計算時會在f 操作處做一次 All-reduce 操作，實現(xiàn)梯度的平均。

另一個是對 Self-Attention 的 TP 拆分，如下面的示意圖，包括 Multi-head Attention 和 Linear 層兩個部分，對Multi-head Attention 的 TP 拆分是基于不同的 head 實現(xiàn)的，將不同的 head 放到不同的 GPU 上。同時，對 Linear 層做了行拆分。與 MLP 的兩個 Linear 層 TP 拆分類似，在 Forward 計算時在 g 操作處做一次 All-reduce，以及在 Backward 計算時在 f 操作處做一次 All-reduce。

接下來介紹的是 Sequence Parallelism，即序列并行，是對 TP 的優(yōu)化補充。例如，TP 拆分只會對 MLP 和 Self-Attention 的 Linear 層做并行化拆分，但是對其他的層，例如 Dropout 和 Layernorm 等沒有做并行化拆分，因為在一個 TP group 內(nèi)，Layernorm 層的輸入是相同的，所以它的計算是完全相同的，也即在每 8 張 GPU 上的計算是完全相同的，因此這個部分的計算是冗余的。但是我們可以在它的 input 做 Sequence 維度上的拆分，經(jīng)過拆分后，在不同的 GPU 上它的 Layernorm 層的輸入是經(jīng)過拆分的，包括 Activation 相當于也被拆分過了，所以不會有冗余的計算，從而能減少顯存的占用。

如下 Sequence Parallelism 示意圖，具體實現(xiàn) Sequence Parallelism 時，我們實際上是在 Layernorm 層做 Sequence Parallelism 拆分，然后在 Linear 層，又轉換成 TP 拆分，在 g 操作處將 Activation 進行一次 all-gather 通信。執(zhí)行 TP 之后，再將 Activation 做一次 Reduce-scatter，之后又轉換成 input 在 Sequence 維度被拆分的狀態(tài)。對于通信量而言，使用 Sequence Parallelism 后，原先的一次 All-reduce 通信變成了一次 All-gather 通信加上一次 Reduce-scatter 通信，表面上通信的次數(shù)是變多了一倍，但是由于 all-reduce 的通信量其實是 All-gather 或 Reduce-scatter 的近兩倍，所以整體上的通信量差不多是一致的。也就是我們在 Sequence Parallelism 優(yōu)化后，通信量上并沒有改變，但是我們能收獲更少的計算量以及更少的 Activation 顯存占用。

下面介紹一個 4 GPU 的流水線并行的例子。

首先我們需要將第一個藍色的，即第一個 mini batch 的 Forward 計算在 Device 1 運算起來，計算完畢后傳到 Device 2，然后一致傳到最后一個 Device，這個流水線就逐漸繁忙起來。當最后一個 Device 計算完 Forward 后，開始計算 Backward，之后可以將這個 Backward 往前傳。這樣的計算模式被稱為 One Forward One Backward(1F1B)狀態(tài)。最后會由最后一個 Device 將最后一個 mini batch 的 Backward 往前傳。所以如果只看最后一個 Device 的話，整個流水線是會先空幾個格，然后一直繁忙，然后在最后再空幾個格。

在 1F1B 計算狀態(tài)，前面的幾個空格被稱為 Warm up 階段，最后的幾個空格階段被稱為 Cool down 階段。如果我們只看最后一個 Device，整個流水線的空閑時間被稱為 Bubble 時間，也就是它的 Warm up 時間加上 Cool down 時間。

如何將流水線的空閑時間減少呢？根據(jù)示意圖給出的公式，需要看 Bubble 時間除以理論時間得到的比例。我們需要盡量減小這個比例，例如可以盡量增大 mini batch 數(shù)量，也即需要拆分出來更多的 mini batch。但是并不是拆分更多的 min batch 就一定有收益。如果 mini batch 特別小，那么 GPU 計算的 Overhead 會比較大，單次的計算時間會增加，因此需要有一定的 Trade-off。所以我們對 PP 進行了優(yōu)化，采用交錯流水線并行策略。

交錯流水線并行如下邊的示意圖。與默認流水線并行的區(qū)別在于，默認的流水線并行是對模型層進行了一個連續(xù)的拆分，交錯流水線并行是對模型層進行一個跳躍式的拆分。例如 Device 1 需要計算 Layer 0 至 3。根據(jù)前面介紹的 1F1B 計算狀態(tài)，Device 4 需要在前面有一個 Warm up 時間才能進入計算狀態(tài)。如何減少這個 Warm up 時間呢？其中一個方法是拆分更小的 mini batch，另一個方法是讓每個 Device 計算更少的層就能獲取到第一個 Activation。例如示意圖的下半部分，Device 1 需要計算的模型層是 Layer 0、1 和 Layer 8、9。在不考慮通信的條件下，流水線的 bubble 時間就縮小了一半。這個也是 Megatron 和 Nemo 對流水線并行的一個優(yōu)化。

接下來我們將三個并行優(yōu)化耦合到一起，看下它的通信模式是什么樣。下圖介紹的是一個 20 層的模型，包括了 DP、PP 和 TP 三個維度的并行優(yōu)化。

首先在 DP 維度，我們設置了兩組 DP。由于每組 DP 獲取了不同數(shù)據(jù)，所以產(chǎn)生的梯度是不同的，因此在完成了 Backward 計算后需要對它的梯度做一次更新，在梯度更新前需要執(zhí)行一次梯度的 All-reduce 操作同步。在每組 DP 內(nèi)部，20 層的模型被拆分為四個流水線并行階段 stage，每個 stage 包含了模型的 5 個層。對每個 stage 而言，需要計算完包含的 5 個層，并將對應的 Activation 傳到后續(xù)的 stage，以此類推，這個是流水線并行的通信模式。

在流水線并行的 stage 內(nèi)部是 TP 拆分，由四個 GPU 維護了這五層模型的計算，也即每個 GPU 只計算了其中的 1/4。這四個 GPU 在計算完第一層后需要做一次 All-reduce，然后才能計算第二層，以此類推，這個是 TP 的通信模式。由于 TP 的通信量是最大的，而且通信頻次很高，同時 TP 大O量的通信是必須和計算順序執(zhí)行的，沒辦法做 Overlap 優(yōu)化。因此，我們通常會將TP 設置成最里面的維度，盡量使用 NVLINK 高速帶寬通信。DP 的通信量比 PP 的通信量大很多，所以 DP 優(yōu)化通常設置成第二個維度，也是可以用 NVLINK 通信。當模型比較大時，我們將TP 設置為 8 后，相應的 DP 通信會跨節(jié)點通信。

PP 對通信的需求是最小的，且通信的頻次比較低，通信量也比較小，所以一般會將 PP 優(yōu)化設置成最外面的那個維度。下邊的示意圖存在部分差異是為了方便介紹。通常情況下，集群節(jié)點間的通信帶寬比較低時，可以盡量在節(jié)點間使用流水線并行。

2、GPT-3 模型分布式訓練技術：Activations 顯存優(yōu)化

Checkpointing activation 是比較常用的顯存優(yōu)化技術。簡單地講，Checkpointing activation 就是做重新計算。例如在 Forward 計算時得到的 Activation 不再保留，做 Backward 計算時，對這些 Activation 進行重新計算，這樣可以極大程度上減少對 GPU 顯存的開銷。

Checkpointing activation 的實現(xiàn)方式有下邊示意圖列出來的幾種。其中一種比較 na?ve 的實現(xiàn)是 Full checkpointing，也即對 Transformer 的每個層都進行重新計算。例如在最后一層做 Backward 計算時，需要這個 Transformer layer 在 Backward 計算之前重新執(zhí)行一次 Forward 計算，當重新計算 Forward 后再開始這個層的 Backward 計算。Full checkpointing 對每個 Transformer layer 都打了一個重算點。

所以 Full checkpointing 的好處在于將顯存開銷降低到 O(n)的復雜度，而不足在于對每個 Transformer layer 都要重新計算一遍，從而帶來了近 36% 的額外計算開銷。

另一個優(yōu)化 Activation 的方式是 Sequence Parallelism 加上 Selective checkpointing，將這個重算的開銷從 36% 降低到 4%。

在前面的介紹中已經(jīng)說明了 Sequence Parallelism 主要優(yōu)化了 Dropout 和Layernorm 層。相對于 na?ve 方式實現(xiàn)的 checkpointing，Selective checkpointing 會選擇一些重算性價比高的OP，對一些計算時間比較小但產(chǎn)生Activation 占用的顯存很大的 OP 進行重算。例如下邊示例圖左邊的 Self-attention 模塊，通過對比分析后得出，對 Self-attention這塊做重算的收益是非常高的，因為它的計算量相對會少一點，但它的一些中間結果輸出占用的顯存開銷非常大。因此我們就可以只對這塊做重算。對其他的層，例如 Linear 和 Layernorm 層，可以采用其他的優(yōu)化方法對 Activation 進行優(yōu)化。

Selective checkpointing 的核心思想是對一些性價比高的 OP 做重算，并與其他的并行優(yōu)化方法聯(lián)合使用，達到 1+1>2 的效果。

3、GPT-3 模型分布式訓練技術：分布式優(yōu)化器

接下來介紹下 Distributed Optimizer，即 Zero-1。在模型訓練時，Optimizer 狀態(tài)是模型固定占用顯存的主要部分。假如模型參數(shù)量是m，則 Optimizer 的顯存開銷是 16 * m 字節(jié)。Zero-1 的思想是在不同的 DP rank 上對 Optimizer 狀態(tài)做拆分，所以 Distributed Optimizer（Zero-1）在每個 DP rank 占用的顯存是 16 * m 除以 DP size，明顯降低了顯存開銷，同時也帶來了通信模式的改變。

例如，我們可以直接對梯度做 Reduce-scatter，之后再對自己的那一部分參數(shù)做Optimizer 相關的狀態(tài)更新，Optimizer 狀態(tài)更新完畢后再用一個 All-gather 操作收集模型權重。這樣由 DP 并行的一個 All-reduce 變成了一個 Reduce-scatter 加上一個 All-gather。根據(jù)前面的介紹，這個 All-reduce 的通信量是 Reduce-scatter 或 All-gather 的兩倍，所以整體的通信量是沒有變化的，但通信的次數(shù)變多了。

其次，使用比較多的是 Zero-2 和Zero-3，但目前在 Megatron 和 Nemo 上還沒有做支持。因為在實際使用過程中，Zero-2 和 Zero-3 需要每次進行額外的 forward 和 backward 計算，并且每次都需要做通信，但是流水線并行會將一個大的 batch 拆分成很多個小的 forward 和 backward，也就會造成大的通信量。所以當同時使用 Zero-2 或 Zero-3 加上流水線并行時，通信量會大幅上升，因此這里不推薦大家同時使用流水線并行和 Zero-2 或 Zero-3，但流水線并行可以和 Zero-1 同時使用。

三、大模型性能調(diào)優(yōu)過程的顯存開銷、通信開銷和計算開銷分析

1、顯存開銷分析

接下來我們具體對顯存開銷、通信開銷和計算開銷做分析。首先是顯存開銷，以GPT-3 為例子，主要分為三個部分，第一部分是模型占用的顯存，即模型狀態(tài)的顯存，包括參數(shù)、梯度和 Optimizer 狀態(tài)；第二部分是 Activation 占用的顯存，也就是模型的 OP 產(chǎn)生的中間值，包括模型 Transformer layer 層和 Embedding 層的輸出，這里我們簡化一下問題，只考慮 Transformer layer；第三部分是可能需要的一些臨時 Buffer，這里不做過多的展開，我們主要關注前面兩個比較重要的部分。

先看下模型參數(shù)顯存開銷，即模型做了 TP 和 PP 拆分后每個 GPU 上會有多大的實際參數(shù)。我們先看下 GPT-3 模型的整體構成，模型包含了 96 層 Decoder，在每層 Decoder 內(nèi)包含一個 QKV 權重、Projection 權重以及兩個 FC 的權重，這樣的權重會被重復 96 次，另外包含了 Token Embedding 和 Position Encoding 的權重。當進行TP 拆分時，相當于 QKV 權重和 Projection 權重會除以 t，再擠上兩層的 FC 權重也除以 t，t 是TP 拆分的數(shù)量。然后進行 PP 拆分，也就是原始的 n 層除以 p，p 是 PP 拆分的 stage 數(shù)量，再加上 Token Embedding 和 Position Encoding 也做了 TP 維度的拆分。所以當 TP 等 8，PP 等于 16 時，每個 GPU 上需要存儲 1.37B 的參數(shù)量。

對于 1.37B 參數(shù)量實際占用的顯存，當我們使用混合精度時，模型的參數(shù)和梯度占用的顯存都是 2.74GB。由于 Optimizer 狀態(tài)占用的顯存是 16 倍的參數(shù)量，所以 Optimizer 占用的顯存是 21.9GB，可以看到主要是 Optimizer 狀態(tài)占用很大比例的顯存開銷。

Optimizer 狀態(tài)的顯存開銷可以通過前面介紹的 Distributed optimizer 或 Zero 技術進行優(yōu)化。其中，Zero-1技術可以將 Optimizer 狀態(tài)做一個 DP 維度的拆分，從而降低顯存占用。例如，當DP size 設置為 8 時，Optimizer 狀態(tài)的顯存占用可以從 21.9GB 降低到 2.7GB。

若在 Zero-1 基礎上再做 Zero-2，可以將梯度也進行 DP 維度的拆分，梯度的顯存占用由 2.74GB 降低到 0.34GB。若再做 Zero-3，可以將參數(shù)也做相應的拆分，將其顯存占用降低為 0.34GB。

可以看到，當 DP size 等于8 時，使用 Zero-1 對 Optimizer 狀態(tài)做拆分就可以帶來明顯的顯存開銷降低。相對而言，使用 Zero-2、Zero-3 對梯度和參數(shù)做拆分時的顯存開銷優(yōu)化不明顯，同時會增加非常多的通信 overhead。所以建議如果不是顯存特別緊張的條件下，盡量不要用 Zero-3。一方面它帶來很大的計算開銷和通信開銷，另一方面它的優(yōu)化也會比較難調(diào)試。

從另一個角度來說，當模型參數(shù)被 TP 和 PP 拆分后，每個 GPU 上的參數(shù)和梯度自然地被拆分了，所以它占用的顯存也比較小，使用 Zero-1 就能達到一個很好的效果。另外，值得注意的是，我們并不推薦使用 PP 后再疊加 ZeRO-2/3，主要原因是一方面它對顯存的進一步優(yōu)化不是很明顯；另一方面是它會帶來一定的通信沖突，所以建議大家使用 TP 和 PP 后用 Zero-1 就可以了。

接下來介紹 Activation 顯存開銷。為簡化問題，我們只分析單個 Transformer layer 的 Activation 顯存占用。具體的計算方法是考慮每個 OP 計算輸出的 shape 大小，將各個 OP 輸出的 shape 加起來即可得到預計的顯存開銷大小。如下邊示例圖中的公式，從計算公式可以發(fā)現(xiàn)一個很大的問題，公式包含了一個s 即 Sequence 長度的平方項，當 s 比較大時，這個顯存開銷會變得非常恐怖，有必要進行優(yōu)化。

使用前面介紹的 Full checkpointing 優(yōu)化方法可以有效減少 Activation 顯存開銷。在 Full checkpointing 中，僅存儲每個 Transformer layer 的輸出，且中間計算結果在 Backward 時直接重算，從而使得 Activation 顯存開銷計算變?yōu)?2 * b * s * h。

使用 TP 拆分，實際上就是對 Self-attention 和 MLP 進行了 TP 維度的拆分，將相應 QKV 和 FC 計算輸出的 shape 除以 t，如果在這個基礎上，再加上Sequence Parallelism 的話，會進一步對兩個 Layernorm 輸出的 shape 除以 t。在這個基礎上繼續(xù)使用 Selective checkpointing 時，就可以將這個 s 的平方給優(yōu)化掉，因為我們會把這個 Self-attention 的輸出 Activation 直接 Drop 掉。

因此，整體的重算開銷從 Full checkpointing 時的 36% 降低為TP+SP+Selective checkpointing 時的 4%，相應的顯存開銷從原先的 2.86GB 降低為 106MB。

2、通信開銷分析

通信開銷也分成 DP、TP 和 PP 三個部分。首先看下 DP 的通信開銷，如下圖所示，DP 的通信開銷等于梯度通信的大小除以總線帶寬，再乘以一個 All-reduce 系數(shù)，其中總線帶寬需要根據(jù)實際的消息大小去測試出一個總線帶寬的估計。

TP 和 TP+SP 兩個優(yōu)化方法的通信量大小是一致的。具體的 TP 通信開銷計算如下邊示例圖中的公式，分成每次 TP 通信量的大小除以實際獲得的總線帶寬，再乘以一個 All-reduce 系數(shù)，再乘以 Self-attention 和 MLP 各做三次通信次數(shù)，再乘以 Transformer layer 的層數(shù)以及 mini batch 的數(shù)量。

流水線并行的通信開銷主要看下邊的示例圖的綠色部分，具體計算是首先獲取單次通信的時間再乘以一個通信的次數(shù)。其中，通信次數(shù)包括了 1F1B 次數(shù)，再加上 Warm up 和 Cool down 的次數(shù)。

3、計算開銷分析

接下來介紹下計算開銷，具體的計算公式推導見下邊的示例圖公式。示例圖左邊的是對 Self-attention、FFN 和 Logit 層的計算拆解，主要計算一些 Linear 層的 FLOPS。將這些 FLOPS 加起來就得到單個 Transformer layer 的 FLOPS，然后再乘上 Batch 數(shù)量以及層數(shù)就能得到整體的模型 FLOPS。通過模型 FLOPS 可以衡量整個訓練系統(tǒng)的效率是否足夠好。例如針對 GPT-3 175B 而言，如果它的模型 FLOPS 在 A100 (*僅供技術交流使用）上小于 150 TFLOPS，說明訓練可能是有一點問題的，如果小于 120TFLOPS 說明訓練肯定有一些比較嚴重的異常問題，造成速度下降。

四、總結

最后來進行一下總結。對于 GPT-3 175B 的訓練而言，我們該怎么去組合這些優(yōu)化？首先有些優(yōu)化項是默認開啟的，比如我們通常推薦大家使用混合精度、Flash attention 以及 Megatron 上大量的默認優(yōu)化 OP。其次是半精度訓練選擇，比如要選擇用 BF16 還是 FP16，對于大模型訓練推薦優(yōu)先考慮 BF16，因為 FP16 會有很多潛在的問題和坑，使用 BF16 會比較穩(wěn)定，尤其是對 20B 以上的模型，推薦大家只用 BF16。

接下來是如何將這些優(yōu)化進行組合，例如，如果訓練運行起來有顯存開銷問題，可以依次打開這些優(yōu)化，首先可以嘗試使用 Selective activation checkpointing，然后開啟 Distributed Optimizer，之后逐漸開啟 TP。我們不需要一開始就將 TP 設置成 8，這里有個準則，需要評判 Hidden size 除以 TP 后不能小于 1024，最好要大于 2048，這樣會有一個比較好的收益。之后逐漸地去開啟 PP。不要一開始就將 PP 設置得很大，然后可以用 Full activation checkpointing。

如果不會遇到 OOM，并且 GBS 能開的足夠大，那么盡量使用 DP 去擴展即可。