[源碼解析] 模型並行分佈式訓練Megatron (1) --- 論文 & 基礎

[源碼解析] 模型並行分佈式訓練Megatron (1) --- 論文 & 基礎

0x00 摘要

NVIDIA Megatron 是一個基於 PyTorch 的分佈式訓練框架，用來訓練超大Transformer語言模型，其通過綜合應用了數據並行，Tensor並行和Pipeline並行來複現 GPT3，值得我們深入分析其背後機理。

本系列大概有6～7篇文章，通過論文和源碼和大家一起學習研究。

本文把 Megatron 的兩篇論文/一篇官方PPT 選取部分內容，糅合在一起進行翻譯分析，希望大家可以通過本文對 Megatron 思路有一個基本瞭解。

0x01 Introduction

1.1 問題

在NLP領域之中，大模型可以帶來更精準強大的語義理解和推理能力，所以隨着規模計算的普及和數據集的增大，使得模型的參數數量也以指數級的速度增長。訓練這樣大的模型非常具有挑戰性，具體原因如下：

(a) 對顯存的挑戰。即使是最大的GPU的主內存也不可能適合這些模型的參數，比如一個175B的GPT-3模型需要（175B * 4bytes）就是700GB模型參數空間，從而梯度也是700G，優化器狀態是1400G，一共2.8TB。
(b) 對計算的挑戰。即使我們能夠把模型放進單個GPU中（例如，通過在主機和設備內存之間交換參數），但是其所需的大量計算操作會導致漫長訓練時間（例如，使用單個V100 NVIDIA GPU來訓練1750億個參數的GPT-3需要大約288年）。如何計算可以參見 2104.04473的附錄 FLOATING-POINT OPERATIONS。
(c) 對計算的挑戰。不同並行策略對應的通信模式和通信量不同。
- 數據並行：通信發生在後向傳播的梯度規約all-reduce操作，通信量是每個GPU之上模型的大小。
- 模型並行：我們在下面會詳述。

這就需要採用並行化來加速。使用硬件加速器來橫向擴展（scale out）深度神經網絡訓練主要有兩種模式：數據並行，模型並行。

1.2 數據並行

數據並行模式會在每個worker之上覆制一份模型，這樣每個worker都有一個完整模型的副本。輸入數據集是分片的，一個訓練的小批量數據將在多個worker之間分割；worker定期彙總它們的梯度，以確保所有worker看到一個一致的權重版本。對於無法放進單個worker的大型模型，人們可以在模型之中較小的分片上使用數據並行。

數據並行擴展通常效果很好，但有兩個限制：

a）超過某一個點之後，每個GPU的batch size變得太小，這降低了GPU的利用率，增加了通信成本；
b）可使用的最大設備數就是batch size，着限制了可用於訓練的加速器數量。

1.3 模型並行

人們會使用一些內存管理技術，如激活檢查點（activation checkpointing）來克服數據並行的這種限制，也會使用模型並行來對模型進行分區來解決這兩個挑戰，使得權重及其關聯的優化器狀態不需要同時駐留在處理器上。

模型並行模式會讓一個模型的內存和計算分佈在多個worker之間，以此來解決一個模型在一張卡上無法容納的問題，其解決方法是把模型放到多個設備之上。

模型並行分爲兩種：流水線並行和張量並行，就是把模型切分的方式。

流水線並行（pipeline model parallel）是把模型不同的層放到不同設備之上，比如前面幾層放到一個設備之上，中間幾層放到另外一個設備上，最後幾層放到第三個設備之上。
張量並行則是層內分割，把某一個層做切分，放置到不同設備之上，也可以理解爲把矩陣運算分配到不同的設備之上，比如把某個矩陣乘法切分成爲多個矩陣乘法放到不同設備之上。

具體如下圖，上面是層間並行（流水線並行），縱向切一刀，前面三層給第一個GPU，後面三層給第二個GPU。下面是層內並行（tensor並行），橫向切一刀，每個張量分成兩塊，分到不同GPU之上。

或者從另一個角度看看，兩種切分同時存在，是正交和互補的（orthogonal and complimentary）。

圖來自：GTC 2020: Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

1.3.1 通信

我們接下來看看模型並行的通信狀況。

張量並行：通信發生在每層的前向傳播和後向傳播過程之中，通信類型是all-reduce，不但單次通信數據量大，並且通信頻繁。
流水線並行：通信在流水線階段相鄰的切分點之上，通信類型是P2P通信，單詞通信數據量較少但是比較頻繁，而且因爲流水線的特點，會產生GPU空閒時間，這裏稱爲流水線氣泡（Bubble）。

比如下圖之中，上方是原始流水線，下面是模型並行，中間給出了 Bubble 位置。

因爲張量並行一般都在同一個機器之上，所以通過 NVLink 來進行加速，對於流水線並行，一般通過 Infiniband 交換機進行連接。

圖來自 Megatron 論文。

1.3.2 張量並行

有些工作在張量（層內）模型並行化（ tensor (intra-layer) model parallelism）做出了一些嘗試，即每個transformer 層內的矩陣乘法被分割到多個GPU上，雖然這種方法在NVIDIA DGX A100服務器（有8個80GB-A100 GPU）上對規模不超過200億個參數的模型效果很好，但對更大的模型就會出現問題。因爲較大的模型需要在多個多GPU服務器上分割，這導致了兩個問題。

(a) 張量並行所需的all-reduce通信需要通過服務器間的鏈接，這比多GPU服務器內的高帶寬NVLink要慢；
(b) 高度的模型並行會產生很多小矩陣乘法（GEMMs），這可能會降低GPU的利用率。

1.3.3 流水線並行

流水線模型並行化是另一項支持大型模型訓練的技術。在流水線並行之中，一個模型的各層會在多個GPU上做切分。一個批次（batch）被分割成較小的微批（microbatches），並在這些微批上進行流水線式執行。

通過流水線並行，一個模型的層被分散到多個設備上。當用於具有相同transformer塊重複的模型時，每個設備可以被分配相同數量的transformer層。Megatron不考慮更多的非對稱模型架構，在這種架構下，層的分配到流水線階段是比較困難的。在流水線模型並行中，訓練會在一個設備上執行一組操作，然後將輸出傳遞到流水線中下一個設備，下一個設備將執行另一組不同操作。

原生（naive）流水線會有這樣的問題：一個輸入在後向傳遞中看到的權重更新並不是其前向傳遞中所對應的。所以，流水線方案需要確保輸入在前向和後向傳播中看到一致的權重版本，以實現明確的同步權重更新語義。

模型的層可以用各種方式分配給worker，並且對於輸入的前向計算和後向計算使用不同的schedule。層的分配策略和調度策略導致了不同的性能權衡。無論哪種調度策略，爲了保持嚴格的優化器語義，優化器操作步驟（step）需要跨設備同步，這樣，在每個批次結束時需要進行流水線刷新來完成微批執行操作（同時沒有新的微批被注入）。Megatron引入了定期流水線刷新。

在每個批次的開始和結束時，設備是空閒的。我們把這個空閒時間稱爲流水線bubble，並希望它儘可能的小。根據注入流水線的微批數量，多達50%的時間可能被用於刷新流水線。微批數量與流水線深度（size）的比例越大，流水線刷新所花費的時間就越少。因此，爲了實現高效率，通常需要較大的batch size。

一些方法將參數服務器與流水線並行使用。然而，這些都存在不一致的問題。TensorFlow的GPipe框架通過使用同步梯度下降克服了這種不一致性問題。然而，這種方法需要額外的邏輯來處理這些通信和計算操作流水線，並且會遇到降低效率的流水線氣泡，或者對優化器本身的更改會影響準確性。

某些異步和bounded-staleness方法，如PipeMare、PipeDream和PipeDream-2BW完全取消了刷新，但這樣會放鬆了權重更新語義。Megatron會在未來的工作中考慮這些方案。

1.4 技術組合

用戶可以使用各種技術來訓練他們的大型模型，每種技術都有不同的權衡。此外，這些技術也可以被結合起來使用。然而，結合這些技術會導致複雜的相互作用，對於系統拓撲是個極大的挑戰，不僅要對模型做合理切割（依據算法特點），還需要做軟硬件一體的系統架構設計，需要仔細推理以獲得良好的性能。因此以下問題就特別重要：

應該如何組合並行技術，以便在保留嚴格的優化器語義的同時，在給定的batch size下最大限度地提高大型模型的訓練吞吐量？

Megatron-LM 開發人員展示了一個如何結合流水線、張量和數據並行，名爲PTD-P的技術，這項技術將以良好的計算性能（峯值設備吞吐量的52%）在1000個GPU上訓練大型語言模型。PTD-P利用跨多GPU服務器的流水線並行、多GPU服務器內的張量並行和數據並行的組合，在同一服務器和跨服務器的GPU之間具有高帶寬鏈接的優化集羣環境中訓練具有一萬億個參數的模型，並具有優雅的擴展性。

要實現這種規模化的吞吐量，需要在多個方面進行創新和精心設計：

高效的核（kernel）實現，這使大部分計算操作是計算綁定（compute-bound）而不是內存綁定（memory-bound。
在設備上對計算圖進行智能分割，以減少通過網絡發送的字節數，同時也限制設備的空閒時間。
實施特定領域的通信優化和使用高速硬件（比如最先進的GPU，並且同一服務器內和不同服務器GPU之間使用高帶寬鏈接）。

1.5 指導原則

Megatron 開發者研究了各種組合之間如何影響吞吐量，基於這些研究得出來分佈式訓練的一些指導原則：

不同的並行模式以複雜的方式互相作用：並行化策略影響通信量、核的計算效率，以及worker因流水線刷新（流水線氣泡）而等待的空閒時間。例如，張量模型並行在多GPU服務器中是有效的，但大模型必須採用流水線模型並行。
用於流水線並行的schdule對通信量、流水線氣泡大小和用於存儲激活的內存都有影響。Megatron 提出了一個新的交錯schdule，與以前提出的schdule相比，它可以在稍微提高內存佔用的基礎上提高多達10%的吞吐量。
超參數的值，如microbatch size，對memory footprint、在worker上執行的核效果和流水線bubble大小有影響。
分佈式訓練是通信密集型的。使用較慢的節點間連接或更多的通信密集型分區會阻礙性能。

0x02 張量模型並行（Tensor Model Parallelism）

2.1 原理

我們用 GEMM 來看看如何進行模型並行，這裏要進行的是 XA = Y，對於模型來說，X 是輸入，A是權重，Y是輸出。從數學原理上來看，對於linear層就是把矩陣分塊進行計算，然後把結果合併，對於非linear層則不做額外設計。

2.1.1 行並行（Row Parallelism）

我們先看看Row Parallelism，就是把 A 按照行分割成兩部分。爲了保證運算，同時我們也把 X 按照列來分割爲兩部分，這裏 $X_1$的最後一個維度等於 $A_1$ 最前的一個維度，理論上是：

\[XA = \begin{bmatrix}X_1& X_2\end{bmatrix} \begin{bmatrix}A_1 \\ A_2\end{bmatrix} = X_1 A_1 + X_2 A_2 = Y_1 + Y_2 = Y \]

所以，$X_1 $和 $ A_1 $ 就可以放到第一個 GPU 之上計算，$X_2$ 和 $A_2$ 可以放到第二個 GPU 之上，然後把結果相加。

我們接下來進行計算。第一步是把圖上橫向紅色箭頭和縱向箭頭進行點積，得到Y中的綠色。

第三步，計算出來一個新的綠色。

第四步，計算了輸出的一行。

第五步，繼續執行，得出了一個 $Y_1$。

第六步，得出了藍色的 $Y_2$，此時，可以把 $Y_1, Y_2$ 加起來，得到最終的輸出 Y。

2.1.2 列並行（Column Parallelism）

我們接下來看看另外一種並行方式Column Parallelism，就是把 A按照列來分割。

最終計算結果如下：

圖來自：GTC 2020: Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

2.2 Model Parallel Transformers

這裏Transformer的模型並行，特指層內切分，即 Tensor Model Parallel。

2.2.1 Transformer

自從2018年Google的Attention論文推出之後，近年的模型架構都是在 Transformer基礎之上完成，模型有多少層，就意味着模型有多少個Transformer塊，所以語言模型的計算量主要是Transformer的計算，而Transformer本質上就是大量的矩陣計算，適合GPU並行操作。

Transformers層由一個Masked Multi Self Attention和Feed Forward兩部分構成，Feed Forward 部分是一個MLP網絡，由多個全連接層構成，每個全連接層是由矩陣乘操作和GeLU激活層或者Dropout構成。

Megatron 的 Feed Forward 是一個兩層多層感知器（MLP），第一層是從 H變成4H，第二層是從 4H 變回到 H，所以Transformer具體架構如下，紫色塊對應於全連接層。每個藍色塊表示一個被複制N次的transformer層，紅色的 x L 代表此藍色複製 L 次。

2.2.2 切分 Transformer

分佈式張量計算是一種正交且更通用的方法，它將張量操作劃分到多個設備上，以加速計算或增加模型大小。FlexFlow是一個進行這種並行計算的深度學習框架，並且提供了一種選擇最佳並行化策略的方法。最近，Mesh TensorFlow引入了一種語言，用於指定TensorFlow中的一般分佈式張量計算。用戶在語言中指定並行維度，並使用適當的集合原語編譯生成一個計算圖。我們採用了Mesh TensorFlow的相似見解，並利用transformer's attention heads 的計算並行性來並行化Transformer模型。然而，Megatron沒有實現模型並行性的框架和編譯器，而是對現有的PyTorch transformer實現進行了一些有針對性的修改。Megatron的方法很簡單，不需要任何新的編譯器或代碼重寫，只是通過插入一些簡單的原語來完全實現，

Megatron就是要把 Masked Multi Self Attention 和Feed Forward 都進行切分以並行化，利用Transformers網絡的結構，通過添加一些同步原語來創建一個簡單的模型並行實現。

2.2.3 切分MLP

我們從MLP塊開始。MLP 塊的第一部分是GEMM，後面是GeLU：

\[Y = GeLU(XA) \]

並行化GEMM的一個選項是沿行方向分割權重矩陣A，沿列切分輸入X：

\[X = \begin{bmatrix} X_1& X_2 \end{bmatrix}, A = \begin{bmatrix} A_1 \\ A_2 \end{bmatrix} \]

分區的結果就變成 $ Y = GeLU(X_1A_1 + X_2A_2) $，括號之中的兩項，每一個都可以在一個獨立的GPU之上完成，然後通過 all-reduce 操作完成求和操縱。既然 GeLU 是一個非線性函數，那麼就有 $GeLU(X_1A_1 + X_2A_2) \neq GeLU(X_1A_1) + GeLH(X_2A_2)$，所以這種方案需要在 GeLU 函數之前加上一個同步點。這個同步點讓不同GPU之間交換信息。

另一個選項是沿列拆分A，得到 $A=[A_1，A_2]$。該分區允許GeLU非線性獨立應用於每個分區GEMM的輸出：

\[\begin{bmatrix} Y_1& Y_2 \end{bmatrix}= \begin{bmatrix} GeLU(XA_1),GeLU(XA_2) \end{bmatrix} \]

這個方法更好，因爲它刪除了同步點，直接把兩個 GeLU 的輸出拼接在一起就行。因此，我們以這種列並行方式劃分第一個GEMM，並沿其行分割第二個GEMM，以便它直接獲取GeLU層的輸出，而不需要任何其他通信（比如 all-reduce 就不需要了），如圖所示。

上圖第一個是 GeLU 操作，第二個是 Dropout操作，具體邏輯如下：

MLP的整個輸入 X 通過 f 放置到每一塊 GPU 之上。
對於第一個全連接層：
1. 使用列分割，把權重矩陣切分到兩塊 GPU 之上，得到 $A_1, A_2$。
2. 在每一塊 GPU 之上進行矩陣乘法得到第一個全連接層的輸出 $Y_1$ 和 $Y_2$。
對於第二個全連接層：
1. 使用行切分，把權重矩陣切分到兩個 GPU 之上，得到 $B_1, B_2$。
2. 前面輸出 $Y_1$ 和 $Y_2$ 正好滿足需求，直接可以和 B 的相關部分（$B_1, B_2$）做相關計算，不需要通信或者其他操作，就得到了 $Z_1, Z_2$。分別位於兩個GPU之上。
$Z_1, Z_2$ 通過 g 做 all-reduce（這是一個同步點），再通過 dropout 得到了最終的輸出 Z。

然後在GPU之上，第二個GEMM的輸出在傳遞到dropout層之前進行規約。這種方法將MLP塊中的兩個GEMM跨GPU進行拆分，並且只需要在前向過程中進行一次 all-reduce 操作（g 操作符）和在後向過程中進行一次 all-reduce 操作（f 操作符）。這兩個操作符是彼此共軛體，只需幾行代碼就可以在PyTorch中實現。作爲示例，f 運算符的實現如下所示：

f算子的實現。g類似於f，在後向函數中使用identity，在前向函數中使用all-reduce。

2.2.4 切分self attention

如下圖所示。

首先，對於自我注意力塊，Megatron 利用了多頭注意力操作中固有的並行性，以列並行方式對與鍵（K）、查詢（Q）和值（V）相關聯的GEMM進行分區，從而在一個GPU上本地完成與每個注意力頭對應的矩陣乘法。這使我們能夠在GPU中分割每個attention head參數和工作負載，每個GPU得到了部分輸出。
其次，對於後續的全連接層，因爲每個GPU之上有了部分輸出，所以對於權重矩陣B就按行切分，與輸入的 $Y_1, Y_2$ 進行直接計算，然後通過 g 之中的 all-reduce 操作和Dropout 得到最終結果 Z。

圖：具有模型並行性的transformer塊。f和g是共軛的。f在前向傳播中使用一個identity運算符，在後向傳播之中使用了all reduce，而g在前向傳播之中使用了all reduce，在後向傳播中使用了identity運算符。

2.2.5 通信

來自線性層（在 self attention 層之後）輸出的後續GEMM會沿着其行實施並行化，並直接獲取並行注意力層的輸出，而不需要GPU之間的通信。這種用於MLP和自我注意層的方法融合了兩個GEMM組，消除了中間的同步點，並導致更好的伸縮性。這使我們能夠在一個簡單的transformer層中執行所有GEMM，只需在正向路徑中使用兩個all-reduce，在反向路徑中使用兩個all-reduce（見下圖）。

圖：transformer層中的通信操作。在一個單模型並行transformer層的正向和反向傳播中總共有4個通信操作。

Transformer語言模型輸出了一個嵌入，其維數爲隱藏大小（H）乘以詞彙量大小（v）。由於現代語言模型的詞彙量約爲數萬個（例如，GPT-2使用的詞彙量爲50257），因此將嵌入GEMM的輸出並行化是非常有益的。然而，在transformer語言模型中，想讓輸出嵌入層與輸入嵌入層共享權重，需要對兩者進行修改。

我們沿着詞彙表維度 $E=[E_1，E_2]$（按列）對輸入嵌入權重矩陣$E_{H×v}$進行並行化。因爲每個分區現在只包含嵌入表的一部分，所以在輸入嵌入之後需要一個all-reduce（g操作符）。對於輸出嵌入，一種方法是執行並行 $GEMM[Y_1，Y_2]=[XE_1，XE_2]$ 以獲得logit，然後添加一個all-gather $Y=all-gather([Y_1，Y_2])$，並將結果發送到交叉熵損失函數。但是，在這種情況下，由於詞彙表的很大，all-gather 將傳遞$b×s×v$ 個元素（b是batch size，s是序列長度）。爲了減小通信規模，我們將並行$GEMM[Y_1，Y_2]$的輸出與交叉熵損失進行融合，從而將維數降低到$b×s$。

2.2.6 小結

我們的模型並行方法旨在減少通信和控制GPU計算範圍的。我們不是讓一個GPU計算dropout、layer normalization或 residual connection，並將結果廣播給其他GPU，而是選擇跨GPU複製計算。

模型並行性與數據並行性是正交的，因此我們可以同時使用二者在來訓練大型模型。下圖顯示了一組用於混合模型並行和數據並行性的GPU。

一個模型需要佔據8張卡，模型被複制了64分，一共啓動了512個即成。
模型並行。同一服務器內的多個GPU形成模型並行組（model parallel group），例如圖中的GPU 1到8，幷包含分佈在這些GPU上的模型實例。其餘的GPU可能位於同一臺服務器內，也可能位於其他服務器中，它們運行其他模型並行組。每個模型並行組內的GPU執行組內所有GPU之間的all-reduce。
數據並行。在每個模型並行組中具有相同位置的GPU（例如圖中的GPU 1，9，…，505）形成數據並行組（data parallel group），即，具有相同模型參數的進程被分配到同一個數據並行組之中。對於數據並行，每個all-reduce操作在每個模型並行組中一個GPU之上執行。
所有通信都是通過pytorch調用NCCL來實現的。

在反向傳播過程中，我們並行運行多個梯度all-reduce操作，以規約每個不同數據並行組中的權重梯度。所需GPU的總數是模型和數據並行組數量的乘積。

混合模型和數據並行的GPU分組，8路模型並行和64路數據並行。

0x03 並行配置

我們接着看如何混合使用各種並行。

3.1 符號說明

以下是本文餘下使用的符號說明。

3.2 Tensor and Pipeline Model Parallelism

張量和流水線模型並行性都可以用於在多個GPU上劃分模型的參數。如前所述，將流水線並行性與週期性刷新一起使用會產生大小爲 $(𝑝 − 1)/𝑚$的流水線氣泡。讓我們假設𝑑 = 1（數據並行大小），因此 𝑡 · 𝑝 = 𝑛。在此情況下，流水線氣泡大小是：

\[\frac{p-1}{m} = \frac{n/t -1}{m} \]

假如我們固定𝐵, 𝑏, 和𝑑 (𝑚 = 𝐵/(𝑏 · 𝑑) 也固定下來），當 𝑡 增加時，流水線氣泡會相應減小。

不同GPU之間通信量也受𝑝 和𝑡 的影響。管道模型並行具有更便宜的點對點通信。另一方面，張量模型並行性使用更消耗帶寬的all-reduce通信（向前和向後傳遞中各有兩個all-reduce操作）。

使用流水線並行，在每對連續設備（向前或向後傳播）之間爲每個微批次執行的通信總量爲𝑏𝑠h，𝑠 是序列長度，h是隱藏大小（hidden size）。
使用張量模型並行，每個層前向傳播和後向傳播中，總大小𝑏𝑠h的張量需要在 𝑡 個模型副本之中 all-reduce 兩次。

因此，我們看到張量模型並行性增加了設備之間的通信量。因此，當 𝑡 大於單個節點中的GPU數量時，在較慢的節點間鏈路上執行張量模型並行是不合算的。

因此得到：

結論#1：當考慮不同形式的模型並行時，當使用𝑔-GPU服務器，通常應該把張量模型並行度控制在 𝑔 之內，然後使用流水線並行來跨服務器擴展到更大的模型。

3.3 Data and Model Parallelism

然後考慮數據並行和模型並行。

3.3.1 Pipeline Model Parallelism.

我們給定 𝑡 = 1 (tensor-model-parallel size)，那麼每個流水線的微批次數目是 $𝑚 = 𝐵/(𝑑 · 𝑏) = 𝑏′/𝑑,$，這裏 $𝑏′ := 𝐵/𝑏$。給定 GPU 數目爲 n，流水線階段的數目是 𝑝 = 𝑛/(𝑡 · 𝑑) = 𝑛/𝑑，流水線氣泡大小是：

\[\frac{p-1}{m} = \frac{n/d -1}{b'/d} = \frac{n-d}{b'} \]

當 𝑑 變大，𝑛 − 𝑑 變小，因此流水線氣泡變小。因爲模型訓練需要的內存佔用可能大於單個加速器的內存容量，所以不可能增加𝑑 一直到𝑛。而數據並行性所需的all-reduce通信不會隨着更高的數據並行度而增加。

我們還可以分析 batch size 𝐵 增加帶來的影響。對於給定的並行配置，如批大小𝐵 增加，𝑏′ = 𝐵/𝑏 增加，(𝑛 − 𝑑)/𝑏′ 會相應減少，從而增加吞吐量。數據並行所需的all-reduce也變得更少，從而進一步提高了吞吐量。

3.3.2 Data and Tensor Model Parallelism.

使用張量模型並行，每個微批次都需要執行all-reduce通信。這在多GPU服務器之間可能非常昂貴。另一方面，數據並行性對於每個批次只需執行一次 all-reduce。此外，使用張量模型並行，每個模型並行rank在每個模型層中只執行計算的子集，因此對於不夠大的層，現代GPU可能無法以最高效率執行這些子矩陣計算。

結論#2：當使用數據和模型並行時，總的模型並行大小應該爲𝑀 = 𝑡 · 𝑝 ，這樣模型參數和中間元數據可以放入GPU內存。數據並行性可用於將訓練擴展到更多GPU。

3.4 Microbatch Size

微批尺寸 𝑏 的選擇也影響到模型訓練的吞吐量。例如，在單個GPU上，如果微批尺寸較大，每個GPU的吞吐量最多可增加1.3倍。現在，在定並行配置（𝑝,𝑡,𝑑）和批量大小𝐵下，我們想確定最佳微批尺寸𝑏。

無論微批大小如何，數據並行通信量將是相同的。鑑於函數 $t_f(b)$和 $t_b(b)$將微批大小映射到單個微批的前向和後向計算時間，在忽略通信成本的條件下，計算一個batch的總時間爲（如前，定義𝑏′爲𝐵/𝑑）。

\[（b'/b + p - 1).(t_f(b) + t_b(b)) \]

因此，微批的大小既影響操作的算術強度，也影響管道 bubble 大小（通過影響𝑚）。

經驗之談#3: 最佳微批尺寸𝑏取決於模型的吞吐量和內存佔用特性，以及管道深度𝑝、數據並行尺寸𝑑和批尺寸𝐵。

3.5 對比

我們接下來看看各種並行機制的對比。

3.5.1 Tensor versus Pipeline Parallelism.

我們觀察到，張量模型的並行性在節點（DGX A100服務器）內是最好的，因爲它會減少通信量。另一方面，流水線模型並行使用更便宜的點對點通信，可以跨節點執行，而不會限制整個計算。然而，流水線並行性會在流水線氣泡中花費大量時間，因此，應限制流水線級的總數，以便流水線中的microbatches數量是流水線深度的合理倍數。因此，當張量並行大小等於單個節點中的GPU數量（8個，DGX A100個節點）時會達到峯值性能。這一結果表明，單獨使用張量模型並行性（Megatron V1）和流水線模型並行性（PipeDream）都無法與這兩種技術結合使用的性能相匹配。

3.5.2 Pipeline versus Data Parallelism.

通過實驗發現，對於每個batch size，吞吐量隨着流水線並行規模的增加而降低。流水線模型並行應該主要用於支持不適合單個 worker 的大型模型訓練，數據並行應該用於擴大訓練規模。

3.5.3 Tensor versus Data Parallelism.

接下來看看數據和張量模型的並行性對性能的影響。在較大的批處理量和微批處理量爲1的情況下，數據並行通信並不頻繁；張量模型並行需要對批處理中的每個微批進行all-to-all通信。這種all-to-all的通信與張量模型並行主義的通信主導了端到端的訓練時間，特別是當通信需要在多GPU節點上進行時。此外，隨着張量模型並行規模的增加，我們在每個GPU上執行較小的矩陣乘法，降低了每個GPU的利用率。

我們應該注意到，儘管數據並行可以帶來高效的擴展，但我們不能單獨使用數據並行來處理訓練批量有限的大型模型，因爲a）內存容量不足，b）數據並行的擴展限制（例如，GPT-3的訓練批量爲1536。因此，數據並行性只支持並行到1536個GPU；然而，大約有10000個GPU用來訓練這個模型）。