一、背景

本文中我們繼續(xù)介紹另一種非常常見的并行方案——流水線并行（Pipeline Parallelism）。

二、Microsoft Pipelined BP

2.1 摘要

我們在之前的文章中提到過，2012: ImageNet Classification with Deep Convolutional Neural Networks 中使用 Tensor Parallelism 在 2 個 GPU 上訓練 AlexNet。同一年，微軟的研究者開始使用 Pipeline Parallelism 訓練語音識別模型 CD-DNN-HMM（Context-Dependent Deep-Neural-Network HMM）。

首先，作者提出了CD-DNN-HMM 模型，其性能大幅優(yōu)于傳統(tǒng)的基于高斯混合的 HMM。然而，在單個 GPU（NVIDIA Tesla S2090 GPU）訓練這樣的模型需要 59 天。為了解決這個問題，作者提出了近似的 Pipelined BP，通過對模型層進行切分，并使用多個 Mini Batch 同時訓練，可以在單臺機器內(nèi)的 2 個或 4 個 GPU 上實現(xiàn) 1.9x 和 3.3x 的端到端加速，并行化效率達到 0.95 和 0.82，并且不會損失精度。

PS：這里的不損失精度是在作者自己的場景下，換個場景也許會損失精度，這個并不能保證。

2.2 方案

在 Pipelined BP 中，每個 GPU 負責網(wǎng)絡(luò)的一部分層，數(shù)據(jù)從一個 GPU 流向下一個 GPU，每個 GPU 同時處理其負責的層。由于數(shù)據(jù)在 GPU 之間流動需要時間，因此模型權(quán)重的更新將使用稍微延遲的數(shù)據(jù)。這也就意味著，當一個 GPU 接收到數(shù)據(jù)并計算梯度時，這些梯度可能是基于更早之前的模型參數(shù)，作者稱這種方式為延遲更新（Delayed Update）。延遲更新在模型訓練中是一個比較有挑戰(zhàn)的工作，但作者證明，適當?shù)恼{(diào)整模型參數(shù)（比如減小 Mini Batch）可以有效的緩解延遲更新對性能的影響。（PS：在 Microsoft 的 PipeDream 中也有類似的問題，我們后續(xù)會具體介紹）

如下圖 Figure 1 所示，作者發(fā)現(xiàn)使用更小的 Mini Batch Size 并不會影響模型精度，反而會增加精度。但是，更小的 Mini Batch Size 會導致 GPU 上每個 Kernel 的計算量不足，反而降低模型訓練的速度。因此，綜合考慮速度和精度，作者將 Mini Batch Size 256 到 1024 稱為 “Goldilocks zone”。

此外，作者也進一步驗證，針對這個任務(wù)，在 Goldilocks Zone 里 Delayed Update 并不會影響模型效果：

如下圖 Table 4 所示，作者使用不同配置（切分）驗證了本文方案（共 8 層，其中 7 層 Hidden Layer）的有效性，可以看出，使用 Pipeline + Striped top Layer 可以獲得 59/18=3.3x 的加速：

三、CMU STRADS

3.1 摘要

卡內(nèi)基梅隆大學等作者在 2016: STRADS: A Distributed Framework for Scheduled Model Parallel Machine Learning 中介紹了一種名為 Scheduled Model Parallelism (SchMP) 的編程方法，旨在通過有效調(diào)度機器學習（Machine Learning）算法的參數(shù)更新來提高算法的收斂速度。SchMP 充分考慮了參數(shù)之間的依賴性和不同參數(shù)收斂速度的不均衡性。

為了支持大規(guī)模的 SchMP，作者開發(fā)了一個分布式框架 STRADS，它優(yōu)化了 SchMP 程序的吞吐量，并在四個常見的 ML 應(yīng)用上進行了基準測試：LDA topic 建模、矩陣分解、稀疏最小二乘（Lasso）回歸和稀疏邏輯回歸。通過 SchMP 和 STARDS 可以有效提高 ML 中迭代的吞吐量。比如，SchMP LDA 和 SchMP Lasso 分別比當時成熟的基線快 10x 和 5x。

3.2 方案概覽

如下圖 Figure 1 所示，STRADS 系統(tǒng)主要包含 3 個部分：

• 用戶實現(xiàn)的 SchMP 指令。
• 服務(wù)（Scheduler、Job Executors 和 Parameter Manager）。
• 服務(wù)的實現(xiàn)（Static Engine 和 Dynamic Engine）。

用戶通過實現(xiàn) SchMP 指令來創(chuàng)建 SchMP 程序，而 STRADS 系統(tǒng)自動管理底層的機器/通信協(xié)調(diào)問題。

3.3 Dynamic Engine

Dynamic Engine 是 STARADS 中的關(guān)鍵部分，專門針對需要動態(tài)調(diào)度的模型并行（Model Parallelism）機器學習算法進行優(yōu)化。動態(tài)調(diào)度算法根據(jù)模型參數(shù)的當前狀態(tài)來決定更新的順序和優(yōu)先級。

然而，動態(tài)調(diào)度算法也面臨一些挑戰(zhàn)，比如可能需要更多的計算來確定更新順序，或者可能需要生成更小的任務(wù)，這也可能導致網(wǎng)絡(luò)通信的延遲相對于計算時間變得更為顯著。為了解決這些問題，Dynamic Engine 采用了迭代流水線（Pipelining）的方式來有效隱藏網(wǎng)絡(luò)通信的延遲，當一個迭代的更新正在通過網(wǎng)絡(luò)通信時，其他迭代的計算可以并行的執(zhí)行。此外，流水線的深度是可以配置的，以找到最佳的收斂速度和迭代進度之間的平衡。如下圖 Figure 3 所示為其 Dynamic Engine 的 Pipelining 方案，右圖中通過通信和計算的 overlap 可以有效提升訓練速度：

四、Microsoft PipeDream

4.1 摘要

在上述的 STRADS 中作者已經(jīng)提出了使用 Pipelining 的方式實現(xiàn)通信和計算的 Overlap，進而提升訓練速度的方案，然而，其針對的還是 ML（Maching Learning）場景，并且在 CPU 上運行。在 [1806.03377] PipeDream: Fast and Efficient Pipeline Parallel DNN Training 中，微軟提出了 PipeDream，它是一個針對 GPU 上訓練 DNN（Deep Neural Network）的訓練系統(tǒng)，也是第一個通用的、自動化的融合 Pipeline Parallelism 和 Data Parallelism 的訓練系統(tǒng)。 

PipeDream 在多臺機器之間以流水線的方式并行執(zhí)行計算，這種 Pipeline Parallelism 避免了 Data Parallelism 訓練在面對大模型和/或有限網(wǎng)絡(luò)帶寬時導致的訓練變慢問題（通信占比很高）。具體來說，PipeDream 通過減少通信量（相對大型 DNN 模型的 DP 訓練方式減少高達 95% 通信量），并允許計算和通信的完美折疊，使得所有 GPU 都保持高速運轉(zhuǎn)。

PS：需要指出的是，PipeDream 中使用的還是 Parameter Server 架構(gòu)，有專門的參數(shù)服務(wù)器。此外，PipeDream 最初也不是基于 Google GPipe 改進的方案，PipeDream 比 GPipe 還早幾個月，SOSP 2019 - PipeDream: Generalized Pipeline Parallelism for DNN Training 中增加了與 GPipe 的對比。

4.2 方案概覽

PipeDream 結(jié)合了模型并行和流水線并行，通過流水線的方式處理不同的 Mini Batch，使得不同的 Worker 在任何時刻處理不同的輸入，具體來說，PipeDream 通過以下方式實現(xiàn)流水線并行訓練：

1. 根據(jù)模型架構(gòu)和硬件配置將 DNN 分層自動切分到不同的階段（Stage）。
2. 使用 DP 并行處理某些 Stage，平衡工作負載。
3. 交錯處理每個 Worker 的 Forward 和 Backward，確保所有 Worker 盡可能忙碌，同時防止過多正在訓練的 Mini Batch，確保模型收斂。
4. 維護每個執(zhí)行中的 Mini Batch 的中間狀態(tài)，以確保等價訓練，不會損失效果。

4.3 模型自動切分

PipeDream 提出的時候 Transformer 還沒有火起來，作者針對的主要是傳統(tǒng)的 DNN 模型，其在不同 Layer 的結(jié)構(gòu)、計算量是不同的，模型的切分相對也就沒那么簡單。為此，作者提出了自動切分機制，如下圖 Figure 7 所示，首先使用一些輸入數(shù)據(jù)運行模型，并分析各層的計算時間，然后按照設(shè)備數(shù)盡可能的均勻切分。

如下圖 Figure 9 所示，PipeDream 也有專門的參數(shù)服務(wù)器（Parameter Server）以存儲和更新模型參數(shù)：

4.4 1F1B

如下圖 Figure 3 所示為使用 4 臺機器進行 Model Parallelism（這里模型并行不是 Tensor Parallelism，而是 Pipeline Parallelism） 訓練的執(zhí)行過程。其每一行代表一個機器，藍色表示 Forward，綠色表示 Backward，F(xiàn)orward 和 Backward 中的數(shù)字指的是 Mini Batch 的 ID。由于是按層切分，并且同一時間只有 1 個 Mini Batch，每臺機器都必須等待之前的機器執(zhí)行完才能執(zhí)行對應(yīng)的 Stage，導致存在大量的 Bubble。

實際上當 Machine 1 執(zhí)行完 Mini Batch 1 的 Forward 之后就可以開始 Mini Batch 2 的 Forward，以此類推。在調(diào)度的過程中，系統(tǒng)中的每個設(shè)備都會有兩個選擇：

• 對某個 Mini Batch 執(zhí)行 Forward，進而可以將 Mini Batch 傳遞到下一個設(shè)備。
• 對另一個 Mini Batch 執(zhí)行 Backward，進而確保學習向前推進。

如果始終優(yōu)先考慮 Forward，則會導致阻塞 Backward，模型也就無法學習和更新，因為只有 Backward 后才能執(zhí)行權(quán)重更新；同樣，如果只考慮 Backward 優(yōu)先調(diào)度，則會導致計算資源閑置，無法充分發(fā)揮算力。

為了避免上述問題，作者提出了 1F1B（1次 Forward，1次 Backward）的調(diào)度機制，如下圖 Figure 8 所示，4 個設(shè)備，分成 4 個 Stage。在起始階段允許執(zhí)行多個 Mini Batch 的 Forward，穩(wěn)定后就保持 Forward 和 Backward 的交替執(zhí)行，這樣可以保證 GPU 在穩(wěn)定狀態(tài)下沒有空閑，并且始終繼續(xù)學習。

上述的 1F1B 過程并不需要 Forward 和 Backward 一樣長，實際上，Backward 總是大于 Forward（大約 2 倍），此時 1F1B 依然是有效的調(diào)度機制。

4.5 有效學習

仔細考慮會發(fā)現(xiàn)上述 1F1B 的 Pipelining 方式存在一個問題，在 Backward 后立即更新模型參數(shù)會引入兩種類型的參數(shù)不一致：

• 同一個 Mini Batch 的 Forward 和 Backward 可能使用不同的參數(shù)進行計算。比如，以 Figure 8 的 Machine 1 為例，F(xiàn)orward 5使用的是 Mini Batch 1 更新后的參數(shù)；而Backward 5使用的是 Mini Batch 1,2,3,4 共同更新后的參數(shù)。
• 同一個 Mini Batch 在不同 Stage 使用的參數(shù)版本不一致。還是以 Mini Batch 5 為例，Machine 1 上的Forward 5使用的是 Mini Batch 1 更新后的參數(shù)；Machine 2 上的Forward 5使用的是 Mini Batch 1 和 Mini Batch 2 更新后的參數(shù)。

為了解決參數(shù)不一致的問題，PipeDream 引入了 Weight Stashing 和 Vertical Sync，如下圖 Figure 9 所示：

• Weight Stashing：為每個正在計算的 Mini Batch 都保存一份參數(shù)。Forward 計算時，每個設(shè)備（Stage）都使用最新的權(quán)重參數(shù)計算輸入的 Mini Batch，并將這個參數(shù)保存，直到當前設(shè)備上對應(yīng)的 Backward 計算完成。這樣可以解決上述的第一個參數(shù)不一致問題，但無法解決第二個。
• Vertical Sync：每個 Mini Batch 開始計算時都使用最新版本的權(quán)重參數(shù)，并且參數(shù)的版本號會伴隨該 Mini Batch 數(shù)據(jù)的整個生命周期，在所有 Stage 都使用同一版本的參數(shù)，從而實現(xiàn) Stage 間的參數(shù)一致性。這樣可以解決上述的第二個參數(shù)不一致問題。?

假設(shè)模型按照 Pipeline Parallelism 切分后不同 Stage 的參數(shù)為 w1,w2 等，t 表示第 t 個 Mini Batch 的更新，則原始的 SGD 對應(yīng)的權(quán)重更新可以表示為：

而使用了 Weight Stashing 后的權(quán)重更新可以表示如下，可以看出已經(jīng)不再等價，并且使用的版本不一致：

進一步引入 Vertical Sync 后對應(yīng)的權(quán)重參數(shù)更新可以表示如下，雖然權(quán)重的版本一致了，都是 t–n+1，但相比 t 來說中間間隔了 n-1 個：

在 PipeDream 中默認采用 Weight Stashing，而不使用 Vertiacl Sync，這也就導致 PipeDream 相當于介于正常的 mini batched SGD 和 DP（使用 BSP 同步）之間的一種方案。

4.6 結(jié)果

如下圖 Table 1 所示，作者在兩個不同的集群上對幾種不同的 DNN 模型進行訓練?？梢钥闯?，本文的 PipeDream 相比傳統(tǒng)的 BSP（DP 訓練）最多可以加速 5x。隨著訓練機器的增加，PipeDream 獲得了更高的加速比，近似線性加速。

五、Google GPipe

5.1 摘要

與微軟的 PipeDream 不太一樣，Google 在 [1811.06965] GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism 提出了類似的 Pipeline Parallelism 方案。然而，GPipe 中作者采用 Mini Batch 拆分方案，也就是將 Mini Batch 拆分為多個 Micro Batch，并保證 Mini Batch 中的 Micro Batch 都 Backward 完才更新模型參數(shù)，以保證與 Mini Batch 更新的等價。此外，作者也引入了 Re-Compute 機制，可以有效降低內(nèi)存占用。

使用 GPipe 可以實現(xiàn)幾乎線性的加速?；?GPipe，作者訓練了 557M 的 AmoebaNet 模型（CV）以及 6B 參數(shù)的 Transformer 模型（NLP），分別實現(xiàn) SOTA 性能。

PS：需要說明的是，GPipe 發(fā)表于 2018 年 11 月，和 Mesh-TensorFlow 同期，是在 PipeDream（2018 年 6 月） 之后。此外，GPipe 在 2019 年 7 月的 v5 版本中才加入 Transformer 模型的實驗，獲得更高的加速比。

5.2 方案

如下圖 Figure 2 所示為本文的主要方案，假設(shè) Mini Batch 的大小為 N，有 K 個設(shè)備：

• Mini Batch 會被切分為 M 個 Micro Batch，每個 Micro Batch 的大小為 N/M。
• M 個 Micro Batch 以 Pipeline 的方式在 K 個設(shè)備上依次執(zhí)行 Forward和 Backward。
• 等 M 個 Micro Batch 都計算完后使用 N 個 Sample 對應(yīng)的梯度統(tǒng)一進行權(quán)重 Update 操作（Optimizer Step）。?

?

此外，作者也引入了 [1604.06174] Training Deep Nets with Sublinear Memory Cost 中的 Re-Compute 機制（也叫 Activation Checkpointing 或 Activation Recomputing）。具體來說，如下圖所示，在每個 Device 中都保留該 Stage 的輸入，而不保留中間的 Activation，當 Backward 階段需要相應(yīng)的 Activation 時，使用對應(yīng)的輸入重新計算 Activation：

當然，由于要在 M 個 Micro Batch 后同步 Update，因此就額外引入了更多的 Bubble，如下圖所示。對于 M 個 Micro Batch，K 個 Device，其 Bubble 率為 O((K-1)/(M+K-1))。作者也通過實驗驗證，當 M>=4K 時可以有效緩解 Bubble 問題。

5.3 結(jié)果

如下圖所示，作者在多個 TPU 設(shè)備上驗證了提出方案的有效性：

• AmoebaNet 模型（CNN）：在各層的參數(shù)量、計算量不同，Pipeline Parallelism 導致很難均勻的按層切分，所以只能實現(xiàn)亞線性加速。比如，8 個 TPU，32 個 Micro Batch 可以獲得 3.48x 加速。
• Transformer 模型（NLP）：每層的參數(shù)量、計算量相同，切分更加均勻，可以實現(xiàn)近似線性加速比。比如，8 個 TPU，32 個 Micro Batch 可以獲得 6.3x 加速。?

六、Microsoft PipeDream-Flush

6.1 摘要

在上述的 Google GPipe 中，作者通過 Activation Recomputing 來減少中間 Activation 占用過多顯存空間，然而其也額外引入了更多的計算量。微軟 PipeDream 的作者在 [2006.09503] Memory-Efficient Pipeline-Parallel DNN Training 中提出了相應(yīng)的改進方案 PipeDream-Flush，通過調(diào)整 Micro Batch Forward 和 Backward 的順序來減少內(nèi)存占用。

6.2 方案

如下圖所示，PipeDream 作者參考 GPipe 的實現(xiàn)引入了 PipeDream flush，進而保證模型的更新和非 Pipeline Parallelism 方式是等價的，只不過重新調(diào)整了 Forward 和 Backward 的順序。

• (a)：GPipe中把 M 個 Micro Batch Forward 計算完之后才會開始 Backward，因此其Micro Batch 1（簡寫 M1） 的Backward 執(zhí)行時內(nèi)存中要存儲 M1、M2、M3 和 M4 的中間 Activation（不考慮 Recomputing）。
• (b)：PipeDream-Flush中，Worker 2 里 M1 的 Backward 計算完成之后 Worker 1 可以馬上開始M1 的 Backward，此時內(nèi)存中只有 M1 和 M2 的中間 Activation。并且計算完之后可以馬上釋放 M1 的中間激活，也就是 M2 的 Backward 計算是內(nèi)存中只有 M2 和 M3 的中間 Activation。?

從上也可以看出，PipeDream-Flush 和 GPipe 是數(shù)學等價的，它們都能保證和非 Pipeline Parallelism 方式 Mini Batch 的訓練完全等價。

6.3 結(jié)果

如下圖 Figure 6 所示，作者對比了不同切分方式/方案的性能，可以看出，本文的 PipeDream-Flush 方案吞吐明顯優(yōu)于 GPipe：

七、NVIDIA Megatron-LM

7.1 摘要

我們在之前的文章中介紹過 NVIDIA 提出的 Megatron-LM（[1909.08053] Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism），其充分分析了 Transformer 模型的 Tensor Parallelism（TP）切分方案，以及與 Data Parallelism（DP）的混合訓練。

在 [2104.04473] Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM 中，NVIDIA 的作者進一步對 Megatron-LM 進行擴展，新增了 Pipeline Parallelism（PP），并提出 Interleaved-1F1B 方案；最后，作者進一步將 PP 與 DP 和 TP 結(jié)合來訓練超大 LLM?；谔岢龅姆桨?，作者使用 3072 個 GPU 以 502 pFLOP/s 的速度訓練了 1T 規(guī)模的 LLM，其 MFU 達到 52%（163 TFlops/GPU，A100 BF16 算力 312 TFlops）。

PS：本文的一作 Deepak Narayanan 也是上述 Microsoft PipeDream-Flush 的一作，并且都是實習期間的工作。

7.2 Interleaved 1F-1B

如下圖 Figure 4 所示為本文作者基于 PipeDream Flush 的 1F1B 提出的 Interleaved 1F1B 調(diào)度方案。

• 上圖為PipeDream Flush 的 1F1B：對應(yīng) K=4 個 Device，Micro Batch 個數(shù)為 M=8，也就是每 M=8 個 Micro Batch 進行一次同步梯度更新。
• 下圖為本文的 Interleaved 1F1B：與標準 1F1B 的主要不同是層的切分方式。假設(shè)模型有 16 層，每個 Device 有 4層。

• 1F1B：模型被分為4 個 Stage，Device 1 包含 Layer (0,1,2,3)，Device 2 包含 Layer (4,5,6,7)，Device 3 包含 Layer(8,9,10,11)，Device 4 包含 Layer(12,13,14,15)。
• Interleaved 1F1B：模型被分為8 個 Stage，Device 1 包含 Layer (0,1,8,9)，Device 2 包含 Layer (2,3,10,11)，Device 3 包含 Layer(4,5,12,13)，Device 4 包含 Layer(6,7,14,15)?？梢钥闯?，相當于將模型切分為 8 個 Stage，但是交替放在 4 個 Device 上，下圖中深色代表前 4 個 Stage（Layer 0-7），淺色代表后 4 個 Stage（Layer 8-15）。以此就可以實現(xiàn)更細力度的調(diào)度，減少 Bubble。?

如下圖所示為 1 個 Mini Batch（8 個 Micro Batch）詳細的調(diào)度過程，其中紅色數(shù)字表示 Layer ID，白色箭頭表示調(diào)度順序?？梢钥闯觯?/p>

• Time-Efficient：在執(zhí)行 M1 Layer(8,9) 的 Backward 之前，M5、M6 和 M7 部分 Stage 的 Forward 可以插入 Bubble 中計算；而在標準 1F1B 中，M5、M6 和 M7 的 Forward 始終是在 M1 的 Backward 之后的，以此可以實現(xiàn)更少的 Bubble。
• Memory-Efficient：在執(zhí)行 M1 Layer(8,9) 的 Backward 之前，Device 1 中只有 M1/M2/M3/M4 Layer(0,1,8,9) 以及 M5/M6/M7 Layer(0,1) 對應(yīng)的中間 Activation；而標準 1F1B 中，Device 1 中只有 M1/M2/M3/M4 Layer(0,1,2,3) 對應(yīng)的中間激活，因此Interleaved-1F1B 會比標準 1F1B 占用更多內(nèi)存，但依然低于 GPipe。
• Memory-Imbalanced：在執(zhí)行 M1 Layer(14,15) 的 Backward 之前，Device 4 只用緩存 M1/M2/M3/M4 Layer(0,1) 以及 M1 Layer(8,9) 對應(yīng)的中間 Activation，比 Device 1 少很多，這也就導致 Memory 不均衡的問題。
• Communication-Inefficient：Interleaved-1F1B 帶來的另外一個問題 Communication 的增加，1F1B 時模型分為 4 個 Stage，在 Forward 中只需要 3 個 P2P 通信，而 Interleaved-1F1B 將其劃分為 8 個 Stage，F(xiàn)orward 中就需要 8 個 P2P 通信；Backward 的通信也會類似增加，整體來說通信量明顯增加。當然，通過充分的 Overlap 也可以避免增加的通信量對訓練效率的影響。?

綜上，Interleaved-1F1B 相比 1F1B 可以降低 Bubble 率，雖然會多占用一些內(nèi)存，但依然優(yōu)于 GPipe；然而，Interleaved-1F1B 也會加劇 Memory 不均衡的問題，并且會增加更多的通信量。

PS：Interleaved-1F1B 也有個約束條件，Mini Batch 中 Micro Batch 的個數(shù) M 需要是 Device 數(shù) K 的整數(shù)倍，比如這里 4 個 Device，得有 4/8/12/16/… 個 Micro Batch。

7.3 結(jié)果

如下圖 Figure 11 所示，作者首先驗證了標準 1F1B 的擴展能力，其中模型包含 24 個 Transformer Layer，共 12.1B 參數(shù)量?？梢钥闯觯?Mini Batch Size 比較?。?）的時候隨著 PP Stage 個數(shù)增加，性能損失比較嚴重；當 Mini Batch Size 比較大（128）的時候隨著 PP Stage 個數(shù)增加模型性能并沒有受到太大影響。

此外，作者也使用 175B 模型（96 層） 96 個 GPU 對比了原始 1F1B 和 Interleaved-1F1B 的性能差異。如下圖 Figure 12 所示，提出的 Interleaved-1F1B 有比較明顯優(yōu)勢，但是隨著 Mini Batch Size 的增加（Micro Batch 個數(shù)增加），這種差異也在逐漸縮小，但是兩者的性能也都有所提升。

最后，作者也使用 162.2B 模型在 64 個 GPU 對比了 TP 和 PP 組合的方案。如下圖 Figure 13 所示，在 TP=8，PP=8 時性能最優(yōu)，兩頭時最差，這也證明了單獨使用 TP 或 PP 可能不是最優(yōu)方案：

八、Sea Zero-Bubble

8.1 摘要

Sea AI-Lab 團隊在 [2401.10241] Zero Bubble Pipeline Parallelism 中進一步提出了 Zero Bubble 方案，可以進一步降低 Bubble 率，作者也聲稱這是第一個 Zero Bubble 的方案。

Zero Bubble 的核心思路是將 Backward 分為兩個部分，一部分計算輸入的梯度，另一部分計算參數(shù)的梯度?；谶@一想法，可以定制更加高效的調(diào)度策略，為此，作者提出根據(jù)特定的模型配置和內(nèi)存約束自動找到最優(yōu)調(diào)度方案，進而可以實現(xiàn) Zero Bubble。此外，作者也提出了在 Optimizer 階段繞過同步的方案。

實驗結(jié)果表明，在類似的內(nèi)存約束下，Zero Bubble 可以比 1F1B 吞吐提高 23%。進一步放寬內(nèi)存約束時，提升可以達到 31%。

8.2 方案

如下圖 Figure 1 所示，作者將 Backward 分成兩個部分，一部分計算輸入的梯度，一部分是計算權(quán)重的梯度。這里計算輸入的梯度有明確的依賴關(guān)系，也是鏈式法則不斷傳遞的基礎(chǔ)；而計算權(quán)重的梯度卻沒有明確的依賴，甚至可以滯后很多。此外，三個紅色部分計算量相當，這也就是為什么之前 1F1B 或者 Interleaved-1F1B 中 Backward 的長度為 Forward 的 2 倍。

如下圖所示為 1F1B 和 本文的 ZB-H1、ZB-H2 的區(qū)別，其中 1F1B 中 Backward 沒有拆分為兩個部分，所以長度時 Forward 2 倍；本文的 ZB 里 Backward 分成了 B 和 W，因此 F、B、W 的長度相同，表示計算 Latency 相當。

• ZB-H1：B 的計算 Latency 更短，也就可以前置，W 的計算沒有明顯的依賴關(guān)系，可以滯后，這樣也就提供了更小化 Bubble 的可能。
• ZB-H2：可以進一步的將 W 的計算滯后，并使用其他 Micro Batch 的 F 和 B 前置來填充 Bubble。只要 Device 里面的所有 Micro Batch 的 W 完成就可以立即開始 Optimizer Step 來更新模型參數(shù)，并且Device 之間不用同步更新。然而，滯后 W 也就意味著相應(yīng)的 Activation 不能釋放，因此ZB-H2 需要占用更多的內(nèi)存空間。?

8.3 結(jié)果

如下圖 Table 4 所示，作者使用不同的模型和 GPU 數(shù)目對比了不同的方案，其中 1F1B-I 指的是 Interleaved-1F1B，ZB-2p 和 ZB-1p 為本文提出的方案。可以看出，在幾乎不增加內(nèi)存占用的情況下 ZB-1p 比 1F1B 有一定的優(yōu)勢；如果允許使用更多內(nèi)存，ZB-2p 的優(yōu)勢會更加明顯。

九、LLaMA 3.1 Pipeline Parallelism

我們前面已經(jīng)提到，PP 有幾個約束：

• 對Mini Batch 大小的約束，比如，需要是 Stage 整數(shù)倍。
• 存在Memory 負載不均衡的問題。
• 也可能存在計算不均衡的問題，比如，在最后的一個 Stage，需要計算 LM head 以及 loss，導致這個 Stage 的 Latency 比較高，成為瓶頸（[2210.02414] GLM-130B: An Open Bilingual Pre-trained Model等工作中都提出了類似的方案）。

為了解決上述幾個問題，Meta 作者訓練 LLaMA 3 模型時修改了 PP 調(diào)度（[2407.21783] The Llama 3 Herd of Models），如下圖 Figure 6 所示，允許靈活的設(shè)置 N，比如這里 N=5，這樣在每個 Batch 中都可以執(zhí)行任意的 Micro Batch。這樣有兩個好處：

• 在大規(guī)模訓練中，可以使用比 PP Stage 更少的 Micro Batch，以便滿足 Batch Size 的要求。
• 也可以使用更多的 Micro Batch，以便隱藏 P2P 通信。?

為了實現(xiàn)更好的負載均衡，作者在 PP 的第一個和最后一個 Stage 減少了一個 Transformer Layer。

為了減少 PP Bubble，作者采用了 Interleaved-1F1B。此外，采用異步 P2P 通信，也打開 TORCH_NCCL_AVOID_RECORD_STREAMS 來減少異步 P2P 通信中的內(nèi)存開銷。

為了進一步減少內(nèi)存開銷，作者也進行了詳細的內(nèi)存分配 Profiling，主動釋放后續(xù)不會使用的 Tensor，比如每個 PP Stage 的輸入和輸出 Tensor。

十、Bloom DP + PP + TP

如下圖 Figure 6 所示為 [2211.05100] BLOOM: A 176B-Parameter Open-Access Multilingual Language Model 訓練（176B）采用的分布式并行方案：8DP 12PP 4TP，比較清晰的說明了怎么結(jié)合幾種分布式并行方案。

十一、參考鏈接

1. ??https://proceedings.neurips.cc/paper_files/paper/2012/file/c399862d3b9d6b76c8436e924a68c45b-Paper.pdf??
2. ??https://dl.acm.org/doi/pdf/10.1145/2901318.2901331??
3. ??https://arxiv.org/abs/1806.03377??
4. ??https://www.pdl.cmu.edu/PDL-FTP/BigLearning/sosp19-final271.pdf??
5. ??https://arxiv.org/abs/1811.06965??
6. ??https://arxiv.org/abs/1604.06174??
7. ??https://arxiv.org/abs/2006.09503??
8. ??https://arxiv.org/abs/1909.08053??
9. ??https://arxiv.org/abs/2104.04473??
10. ??https://arxiv.org/abs/2401.10241??
11. ??https://arxiv.org/abs/2210.02414??
12. ??https://arxiv.org/abs/2407.21783??
13. ??https://arxiv.org/abs/2211.05100??

本文轉(zhuǎn)載自??AI閑談??，作者： AI閑談 ????