[原始碼解析] 模型並行分散式訓練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用來訓練這個模型）。