并行策略

正如和分布式一樣，如何利用多設(shè)備和多硬件也是很重要的一個環(huán)節(jié)。大模型訓(xùn)練也是如此，如今訓(xùn)練大模型離不開各種分布式并行策略，常用的并行策略包括：

• 數(shù)據(jù)并行（data parallelism, DP）：假設(shè)有N張顯卡，每張顯卡都加載完整的模型，每一次迭代（iteration/step）都將一個批次的訓(xùn)練數(shù)據(jù)據(jù)分割成N份系統(tǒng)大小的小批次（micro-batch），每張顯卡按照自身拿到的小批次數(shù)據(jù)進行獨立的計算梯度，然后調(diào)用AllReduce計算梯度均值，進而達到所有的卡的參數(shù)更新后始終保持一致。

下圖為其中一種算法，將所有的梯度分為五份，然后按照圓圈的方式傳播（而不是廣播數(shù)據(jù)）累加，直到每個GPU都完成一個部分的累加，然后再批量同步數(shù)據(jù)。

經(jīng)過五次循環(huán)之后，每個GPU再將各自負責(zé)的那部分數(shù)據(jù)直接同步給其他顯卡，避免多次廣播運算。

• 模型并行（model parallelism/tensor parallelism, MP/TP）：有的模型tensor/layer很大，單張卡無法加載，因此將tensor切分多塊，一張顯卡存其中一塊。
• 流水并行（pipeline parallelism, PP）：按照神經(jīng)網(wǎng)絡(luò)按層進行切割，劃分成多個小組，一張顯卡存放一組。

由于數(shù)據(jù)并行簡單易實現(xiàn)，應(yīng)用最為廣泛。然而每張顯卡都存儲完整的模型，此時顯存大小成為了瓶頸。例如若存在2張顯卡，那么系統(tǒng)中就存在2份模型參數(shù)，如果有4張卡，那么系統(tǒng)中就存在4份模型參數(shù)，如果有N張卡，系統(tǒng)中就存在N份模型參數(shù)，而其中N-1份其實都是冗余。那么有新的辦法來解決這個問題沒有，ZeRO就是其中的一種方法。

初識ZeRO

零冗余優(yōu)化器（Zero Redundancy Optimizer，簡稱ZeRO）是一種用于大規(guī)模分布式深度學(xué)習(xí)的新型內(nèi)存優(yōu)化技術(shù)。ZeRO可以在當(dāng)前一代GPU集群上訓(xùn)練具有100B參數(shù)的深度學(xué)習(xí)模型，吞吐量是當(dāng)前最佳系統(tǒng)的吞吐量的三到五倍。它還為訓(xùn)練具有數(shù)萬億個參數(shù)的模型提供了一條清晰的路徑，展示了深度學(xué)習(xí)系統(tǒng)技術(shù)的前所未有的飛躍。同時ZeRO也被集成到了DeepSpeed，它是用于加速分布式深度學(xué)習(xí)訓(xùn)練的高性能類庫。

在搞清楚ZeRO之前需要先明白在模型訓(xùn)練時候，顯卡的占用情況。一般而言，模型在訓(xùn)練的時候會不斷地更新參數(shù)，目前采用的都是梯度下降法。而混合精度訓(xùn)練（mixed precision training）和Adam優(yōu)化器基本上已經(jīng)是目前訓(xùn)練模型的標配。

Adam在SGD基礎(chǔ)上，為每個參數(shù)梯度增加了一階動量（momentum）和二階動量（variance）1?；旌暇扔?xùn)練，字如其名，同時存在fp16和fp32兩種格式的數(shù)值，其中模型參數(shù)、模型梯度都是fp16，此外還有fp32的模型參數(shù)，如果優(yōu)化器是Adam，則還有fp32的momentum和variance。

<大白話， Adam 是一種可以替代傳統(tǒng)隨機梯度下降過程的一階優(yōu)化算法，它能基于訓(xùn)練數(shù)據(jù)迭代地更新神經(jīng)網(wǎng)絡(luò)權(quán)重!!>

回過頭來看看顯卡中的數(shù)據(jù)分類，一般而言可以分為兩類：

• 模型狀態(tài)（Model States）由模型參數(shù)（fp16）、模型梯度（fp16）和Adam狀態(tài)(FP32的模型參數(shù)副本，F(xiàn)P32的Momentum和FP32的Variance)。

假設(shè)模型參數(shù)量為Φ ，則共需要2Φ+2Φ+(4Φ+4Φ+4Φ)=4Φ+12Φ=16Φ 字節(jié)存儲空間。從這個統(tǒng)計中可以發(fā)現(xiàn)，其實Adam優(yōu)化器占據(jù)了75%的空間。<若讀者要微調(diào)3B的模型，那么則需要48B的空間！??！>

• 剩余狀態(tài)（Residual States）: 除了模型狀態(tài)之外的顯存占用，包括激活值（activation）、各種臨時緩沖區(qū)（buffer）以及無法使用的顯存碎片（fragmentation）。

ZeRO是為了克服數(shù)據(jù)并行性和模型并行性的局限性，同時實現(xiàn)兩者的優(yōu)點。它通過在數(shù)據(jù)并行進程中對模型狀態(tài)（參數(shù)、梯度和優(yōu)化器狀態(tài)）進行分區(qū)，而不是復(fù)制它們，從而消除了數(shù)據(jù)并行進程之間的內(nèi)存冗余。同時采用動態(tài)通信計劃在分布式訓(xùn)練設(shè)備之間共享必要的狀態(tài)，以保持數(shù)據(jù)并行性的計算粒度和通信量。

如何來解讀上圖，藍色就代表模型參數(shù)數(shù)Φ（若每個參數(shù)類型是FP16，則為2Φ字節(jié)），橙色代表模型梯度數(shù)Φ，K為優(yōu)化器Adam的參數(shù)量系數(shù)，即KΦ。

• 第一行是沒有任何優(yōu)化的情況下，假如訓(xùn)練的模型規(guī)模為7.5B，則單卡需要120G的內(nèi)存。
• ZeRO首先進行分區(qū)的是模型狀態(tài)中的Adam優(yōu)化器，也就是上圖第二行的Pos ，這里os指的是optimizer states。模型參數(shù)（parameters）和梯度（gradients）仍舊是每張卡保持一份，此時，每張卡的模型狀態(tài)所需顯存是4Φ+12Φ/N字節(jié)，當(dāng)N比較大時，趨向于4Φ，也就是原來16Φ的1/4 。這個時候單卡內(nèi)存需求量降低到31.5G。
• 若繼續(xù)對模型梯度進行分區(qū)，也就是第三行的Pos+g ，模型參數(shù)仍舊是每張卡保持一份，此時，每張卡的模型狀態(tài)所需顯存是2Φ+(2Φ+12Φ)/N，當(dāng)N比較大時，趨向于2Φ ，也即是原來16Φ的1/8。這個時候單卡內(nèi)存需求量再次降低到16.66G。
• 最后若繼續(xù)對模型參數(shù)進行分區(qū)，也就是最后一行的Pos+g+p ，此時每張卡的模型狀態(tài)所需顯存是16Φ/N，當(dāng) N比較大時，趨向于0。這個時候單卡需要的內(nèi)存量僅僅需要1.9G。

在DeepSpeed的訓(xùn)練框架中， Pos對應(yīng)ZeRO-1， Pos+g對應(yīng)ZeRO-2， Pos+g+p對應(yīng)ZeRO-3，而一般在生產(chǎn)中ZeRO-1就足夠使用。

本文轉(zhuǎn)載自??魯班模錘??，作者： 龐德公 ????