回顧ChatGPT的發展歷程，我們可以總結出大語言模型（LLM）取得驚艷效果的要點（重要性從高到低排序）：

愿意燒錢，且接受“燒錢 != 好模型”的現實

高質量的訓練語料

高效的分布式訓練框架和充沛優質的硬件資源

算法的迭代創新

在大模型訓練這個系列里，我們將一起探索學習幾種經典的分布式并行范式，包括流水線并行（Pipeline Parallelism），數據并行(Data Parallelism)和張量并行（Tensor Parallesim）。微軟開源的分布式訓練框架FastSpeed，融合了這三種并行范式，開發出3D并行的框架，實現了千億級別模型參數的訓練。

本篇文章將探索流水線并行，經典的流水線并行范式有Google推出的Gpipe，和微軟推出的PipeDream。兩者的推出時間都在2019年左右，大體設計框架一致。主要差別為：在梯度更新上，Gpipe是同步的，PipeDream是異步的。異步方法更進一步降低了GPU的空轉時間比。雖然PipeDream設計更精妙些，但是Gpipe因為其“夠用”和淺顯易懂，更受大眾歡迎（torch的PP接口就基于Gpipe）。因此本文以Gpipe作為流水線并行的范例進行介紹。內容包括：

1、優化目標
2、模型并行
3、流水線并行

切分micro-batch

Re-materialization (active checkpoint)

4、實驗效果

推薦閱讀： ChatGPT技術解析系列之：訓練框架InstructGPT ChatGPT技術解析系列之：GPT1、GPT2與GPT3 ChatGPT技術解析系列之：賦予GPT寫代碼能力的Codex

一、優化目標

當你從單卡窮人變成多卡富翁時，你做分布式訓練的總體目標是什么呢？（雖然手握一張A100怎么能是窮呢）

能訓練更大的模型。理想狀況下，模型的大小和GPU的數量成線性關系。即GPU量提升x倍，模型大小也能提升x倍。

能更快地訓練模型。理想狀況下，訓練的速度和GPU的數量成線性關系。即GPU量提升x倍，訓練速度也能提升x倍。

這是目標，也是難點，難在于：

訓練更大的模型時，每塊GPU里不僅要存模型參數，還要存中間結果（用來做Backward）。而更大的模型意味著需要更多的訓練數據，進一步提高了中間結果的大小。加重了每塊GPU的內存壓力。我們將在下文詳細分析這一點。（對應著GPU中的內存限制）

網絡通訊開銷。數據在卡之間進行傳輸，是需要通訊時間的。不做設計的話，這個通訊時間可能會抹平多卡本身帶來的訓練速度提升。（對應著GPU間的帶寬限制）

明確這兩個訓練目標后，我們來看并行范式的設計者，是如何在現有硬件限制的條件下，完成這兩個目標的。

二、模型并行

當你有一個單卡裝不下的大模型時，一個直接的解決辦法是，把模型隔成不同的層，每一層都放到一塊GPU上，如下圖：

此時，模型做一輪forward和backward的過程如下

其中下標表示GPU編號，例如表示在GPU0上做foward，表示在GPU0上做backward。圖中的橫軸表示timestep。

這張圖的含義是：我在GPU0上做完一次forward，然后將GPU0上最后一層的輸入傳給GPU1，繼續做forward，直到四塊GPU都做完forward后，我再依次做backward。等把四塊GPU上的backward全部做完后，最后一個時刻我統一更新每一層的梯度。

這樣做確實能訓更大的模型了，但也帶來了兩個問題：

（1）GPU利用度不夠。

如圖，陰影部分所表示的時間段里，總有GPU在空轉。在Gpipe中，將陰影部分定義為bubble。我們來計算一下bubble。假設有塊GPU，而單塊GPU上做一次forward和backward的時間為：。則：

圖中灰色長方形的整體面積為：（寬=，長=）

圖中實際在做forward和backward的面積為：

圖中陰影部分的面積為：

圖像陰影部分的占比為：

則我們定義出bubble部分的時間復雜度為：，當K越大，即GPU的數量越多時，空置的比例接近1，即GPU的資源都被浪費掉了。因此這個問題肯定需要解決。

（2）中間結果占據大量內存

在做backward計算梯度的過程中，我們需要用到每一層的中間結果z。假設我們的模型有L層，每一層的寬度為d，則對于每塊GPU，不考慮其參數本身的存儲，額外的空間復雜度為。從這個復雜度可以看出，隨著模型的增大，N，L，d三者的增加可能會平滑掉K增加帶來的GPU內存收益。因此，這也是需要優化的地方。

三、訓練數據與訓練方法

樸素的模型并行存在GPU利用度不足，中間結果消耗內存大的問題。而Gpipe提出的流水線并行，就是用來解決這兩個主要問題的。

3.1 切分micro-batch

流水線并行的核心思想是：在模型并行的基礎上，進一步引入數據并行的辦法，即把原先的數據再劃分成若干個batch，送入GPU進行訓練。未劃分前的數據，叫mini-batch。在mini-batch上再劃分的數據，叫micro-batch。

圖例如下：

其中，第一個下標表示GPU編號，第二個下標表示micro-batch編號。假設我們將mini-batch劃分為M個，則流水線并行下，bubble的時間復雜度為(推導過程略，可參照第二部分的bubble推導流程)。Gpipe通過實驗證明，當$M>=4K時，bubble產生的空轉時間占比對最終訓練時長影響是微小的，可以忽略不計。

將batch切好，并逐一送入GPU的過程，就像一個流水生產線一樣（類似于CPU里的流水線），因此也被稱為Pipeline Parallelism。

3.2 re-materialization（active checkpoint）

解決了GPU的空置問題，提升了GPU計算的整體效率。接下來，就要解決GPU的內存問題了。

前文說過，隨著模型的增加，每塊GPU中存儲的中間結果也會越大。對此，Gpipe采用了一種非常簡單粗暴但有效的辦法：用時間換空間，在論文里，這種方法被命名為re-materalization，后人也稱其為active checkpoint。

具體來說，就是幾乎不存中間結果，等到backward的時候，再重新算一遍forward，圖例如下：

每塊GPU上，我們只保存來自上一塊的最后一層輸入z，其余的中間結果我們算完就廢。等到backward的時候再由保存下來的z重新進行forward來算出。

現在我們來計算每塊GPU峰值時刻的內存：

每塊GPU峰值時刻存儲大小 = 每塊GPU上的輸入數據大小 + 每塊GPU在forward過程中的中間結果大小

每塊GPU上固定需要保存它的起始輸入，我們記起始輸入為（即mini-batch的大小）。

每個micro-batch是流水線形式進來的，算完一個micro-batch才算下一個。在計算一個micro-batch的過程中，我們會產生中間變量，它的大小為（其中M為micro-batch個數）。

因此，每塊GPU峰值時刻的空間復雜度為

將其與樸素模型并行中的GPU空間復雜度比較，可以發現，由于采用了micro-batch的方法，當L變大時，流水線并行相比于樸素模型并行，對GPU內存的壓力顯著減小。

如果你使用Pytorch提供的PP接口，其中有一個參數叫checkpoint，就是用來做這一項的。

最后，再提一點，在micro-batch的劃分下，我們在計算Batch Normalization時會有影響。Gpipe的方法是，在訓練時計算和運用的是micro-batch里的均值和方差，但同時持續追蹤全部mini-batch的移動平均和方差，以便在測試階段進行使用。Layer Normalization則不受影響。

四、實驗效果

回顧第二部分的兩個目標，Gpipe真的實現了嗎？如果實現不了，又是因為什么原因呢？我們來看下實驗效果。

4.1 GPU數量 VS 模型大小

Gpipe分別在AmoebaNet（圖像）和Transformer（自然語言）兩個大模型上做了實驗。

Naive表示單卡

Pipeline-N表示re-materalization + N卡。

AmeobaNet-D和Trasformer-L一行表示超參數的量

of Model Parameter表示模型的參數量

Total Model Parameter Memory表示模型參數所占內存大小

Peak Activation Memory表示峰值時中間結果大小。可以發現，中間結果占據的內存大小是相當可觀的。

從實驗結果里，我們可以發現：

在Transformer上，Gpipe基本實現了模型大小（參數量）和GPU個數之間的線性關系。例如從32卡增到128卡時，模型的大小也從21.08B增加到82.9B，約擴4倍

對AmoebaNet而言，卻沒有完全實現線性增長。例如從4卡到8卡，模型大小從1.05B到1.8B，不滿足2倍的關系。本質原因是AmoebaNet模型在切割時，沒有辦法像Transformer一樣切得勻稱，保證每一塊GPU上的內存使用率是差不多的。因此對于AmoebaNet，當GPU個數上升時，某一塊GPU可能成為木桶的短板。

GPU數量 VS 訓練速度

（1）關掉NVlinks

為了驗證Gpipe框架帶來的收益，實驗中關掉了NVlinks（GPU間快速通信的橋梁。估計是通過強迫GPU先連CPU然后再連別的GPU做到的）。關掉的意義在于說明，不靠硬件本身的高效通訊帶來的收益，Gpipe一樣能做的很好。實驗效果如下：

M=32表示micro-batch的數量為32，K表示GPU數量。從實驗結果可知，在關掉NVlinks的情況下，Gpipe一樣也能實現隨著GPU數量的增加，訓練速度也增加的效果。雖然這兩者間不是線性的。同樣，因為模型切割不均的原因，AmoebaNet的表現不如Transformer。

（2）開啟NVlinks，并尋找最佳M

當重新開啟NVlinks后，我們來看M的大?。戳魉€的核心）對訓練速度的影響。

當M=1的時候，如前文所說，GPU的空置率太高，因此兩個模型都沒有實現訓練速度和GPU個數間的線性關系

當M=4時，表現明顯好轉。

當M=32時，表現最佳，且Transformer基本實現了訓練速度和GPU個數的線性關系。

4.3 Gpipe下時間消耗分布

對每塊GPU來說，約2/3的時間，是真正花在計算上的。

其余1/3的時間，大部分花在re-materalization策略下的重計算上。因為采用流水線的方法，bubble的時間也被壓縮到很短，可以忽略不計。