以Gpipe作為流水線并行的范例進(jìn)行介紹

回顧C(jī)hatGPT的發(fā)展歷程，我們可以總結(jié)出大語言模型（LLM）取得驚艷效果的要點(diǎn)（重要性從高到低排序）：

愿意燒錢，且接受“燒錢 != 好模型”的現(xiàn)實(shí)

高質(zhì)量的訓(xùn)練語料

高效的分布式訓(xùn)練框架和充沛優(yōu)質(zhì)的硬件資源

算法的迭代創(chuàng)新

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

本篇文章將探索流水線并行，經(jīng)典的流水線并行范式有Google推出的Gpipe，和微軟推出的PipeDream。兩者的推出時間都在2019年左右，大體設(shè)計(jì)框架一致。主要差別為：在梯度更新上，Gpipe是同步的，PipeDream是異步的。異步方法更進(jìn)一步降低了GPU的空轉(zhuǎn)時間比。雖然PipeDream設(shè)計(jì)更精妙些，但是Gpipe因?yàn)槠洹皦蛴谩焙蜏\顯易懂，更受大眾歡迎（torch的PP接口就基于Gpipe）。因此本文以Gpipe作為流水線并行的范例進(jìn)行介紹。內(nèi)容包括：

1、優(yōu)化目標(biāo)
2、模型并行
3、流水線并行

切分micro-batch

Re-materialization (active checkpoint)

4、實(shí)驗(yàn)效果

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

一、優(yōu)化目標(biāo)

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

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

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

這是目標(biāo)，也是難點(diǎn)，難在于：

訓(xùn)練更大的模型時，每塊GPU里不僅要存模型參數(shù)，還要存中間結(jié)果（用來做Backward）。而更大的模型意味著需要更多的訓(xùn)練數(shù)據(jù)，進(jìn)一步提高了中間結(jié)果的大小。加重了每塊GPU的內(nèi)存壓力。我們將在下文詳細(xì)分析這一點(diǎn)。（對應(yīng)著GPU中的內(nèi)存限制）

網(wǎng)絡(luò)通訊開銷。數(shù)據(jù)在卡之間進(jìn)行傳輸，是需要通訊時間的。不做設(shè)計(jì)的話，這個通訊時間可能會抹平多卡本身帶來的訓(xùn)練速度提升。（對應(yīng)著GPU間的帶寬限制）

明確這兩個訓(xùn)練目標(biāo)后，我們來看并行范式的設(shè)計(jì)者，是如何在現(xiàn)有硬件限制的條件下，完成這兩個目標(biāo)的。

二、模型并行

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

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

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

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

這樣做確實(shí)能訓(xùn)更大的模型了，但也帶來了兩個問題：

（1）GPU利用度不夠。

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

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

圖中實(shí)際在做forward和backward的面積為：

圖中陰影部分的面積為：

圖像陰影部分的占比為：

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

（2）中間結(jié)果占據(jù)大量內(nèi)存

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

三、訓(xùn)練數(shù)據(jù)與訓(xùn)練方法

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

3.1 切分micro-batch

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

圖例如下：

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

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

3.2 re-materialization（active checkpoint）

解決了GPU的空置問題，提升了GPU計(jì)算的整體效率。接下來，就要解決GPU的內(nèi)存問題了。

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

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

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

現(xiàn)在我們來計(jì)算每塊GPU峰值時刻的內(nèi)存：

每塊GPU峰值時刻存儲大小 = 每塊GPU上的輸入數(shù)據(jù)大小 + 每塊GPU在forward過程中的中間結(jié)果大小

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

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

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

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

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

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

四、實(shí)驗(yàn)效果

回顧第二部分的兩個目標(biāo)，Gpipe真的實(shí)現(xiàn)了嗎？如果實(shí)現(xiàn)不了，又是因?yàn)槭裁丛蚰?？我們來看下?shí)驗(yàn)效果。

4.1 GPU數(shù)量 VS 模型大小

Gpipe分別在AmoebaNet（圖像）和Transformer（自然語言）兩個大模型上做了實(shí)驗(yàn)。

Naive表示單卡

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

AmeobaNet-D和Trasformer-L一行表示超參數(shù)的量

of Model Parameter表示模型的參數(shù)量

Total Model Parameter Memory表示模型參數(shù)所占內(nèi)存大小

Peak Activation Memory表示峰值時中間結(jié)果大小?？梢园l(fā)現(xiàn)，中間結(jié)果占據(jù)的內(nèi)存大小是相當(dāng)可觀的。

從實(shí)驗(yàn)結(jié)果里，我們可以發(fā)現(xiàn)：

在Transformer上，Gpipe基本實(shí)現(xiàn)了模型大?。▍?shù)量）和GPU個數(shù)之間的線性關(guān)系。例如從32卡增到128卡時，模型的大小也從21.08B增加到82.9B，約擴(kuò)4倍

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

GPU數(shù)量 VS 訓(xùn)練速度

（1）關(guān)掉NVlinks

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

M=32表示micro-batch的數(shù)量為32，K表示GPU數(shù)量。從實(shí)驗(yàn)結(jié)果可知，在關(guān)掉NVlinks的情況下，Gpipe一樣也能實(shí)現(xiàn)隨著GPU數(shù)量的增加，訓(xùn)練速度也增加的效果。雖然這兩者間不是線性的。同樣，因?yàn)槟Ｐ颓懈畈痪脑颍珹moebaNet的表現(xiàn)不如Transformer。

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

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

當(dāng)M=1的時候，如前文所說，GPU的空置率太高，因此兩個模型都沒有實(shí)現(xiàn)訓(xùn)練速度和GPU個數(shù)間的線性關(guān)系

當(dāng)M=4時，表現(xiàn)明顯好轉(zhuǎn)。

當(dāng)M=32時，表現(xiàn)最佳，且Transformer基本實(shí)現(xiàn)了訓(xùn)練速度和GPU個數(shù)的線性關(guān)系。

4.3 Gpipe下時間消耗分布

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

其余1/3的時間，大部分花在re-materalization策略下的重計(jì)算上。因?yàn)椴捎昧魉€的方法，bubble的時間也被壓縮到很短，可以忽略不計(jì)。