在數(shù)據(jù)并行上篇中，我們介紹了樸素數(shù)據(jù)并行（DP）與分布式數(shù)據(jù)并行（DDP）。兩者的總通訊量雖然相同，但DP存在負載不均的情況，大部分的通訊壓力集中在Server上，而Server的通訊量與GPU數(shù)量呈線性關(guān)系，導(dǎo)致DP一般適用于單機多卡場景。而DDP通過采用Ring-AllReduce這一NCCL操作，使得通訊量均衡分布到每塊GPU上，且該通訊量為一固定常量，不受GPU個數(shù)影響，因此可實現(xiàn)跨機器的訓(xùn)練。

在上篇介紹中，通訊負載不均的優(yōu)化我們解釋過了，但還遺留了一個顯存開銷問題：數(shù)據(jù)并行中，每個GPU上都復(fù)制了一份完整模型，當(dāng)模型變大時，很容易打爆GPU的顯存，那要怎么辦呢？

今天這篇文章，我們將介紹由微軟開發(fā)的ZeRO（零冗余優(yōu)化），它是DeepSpeed這一分布式訓(xùn)練框架的核心，被用來解決大模型訓(xùn)練中的顯存開銷問題。ZeRO的思想就是用通訊換顯存。如果初讀ZeRO，覺得它邏輯跳躍，晦澀難懂，那么這篇文章或許可以幫到你～全文結(jié)構(gòu)如下：

一、存儲消耗

1.1 存儲分類

1.2 混合精度訓(xùn)練

1.3 存儲大小

二、ZeRO-DP（），ZeRO與模型并行
三、ZeRO-R
四、ZeRO-offload與ZeRO-Infinity
五、參考

推薦閱讀：

圖解大模型訓(xùn)練之：流水線并行，以GPipe為例

圖解大模型訓(xùn)練之：數(shù)據(jù)并行上篇（DP, DDP）

一、存儲消耗

1.1 存儲分類

首先，我們來看在大模型訓(xùn)練的過程中，GPU都需要存什么內(nèi)容。

存儲主要分為兩大塊：Model States和Residual StatesModel States指和模型本身息息相關(guān)的，必須存儲的內(nèi)容，具體包括：

optimizer states：Adam優(yōu)化算法中的momentum和variance

gradients：模型梯度

parameters：模型參數(shù)W

Residual States指并非模型必須的，但在訓(xùn)練過程中會額外產(chǎn)生的內(nèi)容，具體包括：

activation：激活值。在流水線并行中我們曾詳細介紹過。在backward過程中使用鏈?zhǔn)椒▌t計算梯度時會用到。有了它算梯度會更快，但它不是必須存儲的，因為可以通過重新做Forward來算它。

temporary buffers: 臨時存儲。例如把梯度發(fā)送到某塊GPU上做加總聚合時產(chǎn)生的存儲。

unusable fragment memory：碎片化的存儲空間。雖然總存儲空間是夠的，但是如果取不到連續(xù)的存儲空間，相關(guān)的請求也會被fail掉。對這類空間浪費可以通過內(nèi)存整理來解決。

1.2 精度混合訓(xùn)練

知道了存儲分類，進一步，我們想知道，假設(shè)模型的參數(shù)W大小是，那么每一類存儲具體占了多大的空間呢？

在分析這個問題前，我們需要來了解精度混合訓(xùn)練。

對于模型，我們肯定希望其參數(shù)越精準(zhǔn)越好，也即我們用fp32（單精度浮點數(shù)，存儲占4byte）來表示參數(shù)W。但是在forward和backward的過程中，fp32的計算開銷也是龐大的。那么能否在計算的過程中，引入fp16或bf16（半精度浮點數(shù)，存儲占2byte），來減輕計算壓力呢？于是，混合精度訓(xùn)練就產(chǎn)生了，它的步驟如下圖：

存儲一份fp32的parameter，momentum和variance（統(tǒng)稱model states）

在forward開始之前，額外開辟一塊存儲空間，將fp32 parameter減半到fp16 parameter。

正常做forward和backward，在此之間產(chǎn)生的activation和gradients，都用fp16進行存儲。

用fp16 gradients去更新fp32下的model states。

當(dāng)模型收斂后，fp32的parameter就是最終的參數(shù)輸出。

通過這種方式，混合精度訓(xùn)練在計算開銷和模型精度上做了權(quán)衡。如果不了解fp32，fp16和bf16的細節(jié)也沒關(guān)系，不影響下文的閱讀。只要記住它們所占的存儲空間和精度表達上的差異即可。

1.3 存儲大小

現(xiàn)在，我們可以來計算模型在訓(xùn)練時需要的存儲大小了，假設(shè)模型的參數(shù)W大小是，以byte為單位，存儲如下：

因為采用了Adam優(yōu)化，所以才會出現(xiàn)momentum和variance，當(dāng)然你也可以選擇別的優(yōu)化辦法。因此這里為了更通用些，記模型必存的數(shù)據(jù)大小為。因此最終內(nèi)存開銷為：

另外，這里暫不將activation納入統(tǒng)計范圍，原因是：

activation不僅與模型參數(shù)相關(guān)，還與batch size相關(guān)

activation的存儲不是必須的。存儲activation只是為了在用鏈?zhǔn)椒▌t做backward的過程中，計算梯度更快一些。但你永遠可以通過只保留最初的輸入X，重新做forward來得到每一層的activation（雖然實際中并不會這么極端）。

因為activation的這種靈活性，納入它后不方便衡量系統(tǒng)性能隨模型增大的真實變動情況。因此在這里不考慮它，在后面會單開一塊說明對activation的優(yōu)化。

二、ZeRO-DP

知道了什么東西會占存儲，以及它們占了多大的存儲之后，我們就可以來談如何優(yōu)化存儲了。

注意到，在整個訓(xùn)練中，有很多states并不會每時每刻都用到，舉例來說；

Adam優(yōu)化下的optimizer states只在最終做update時才用到

數(shù)據(jù)并行中，gradients只在最后做AllReduce和updates時才用到

參數(shù)W只在做forward和backward的那一刻才用到

諸如此類

所以，ZeRO想了一個簡單粗暴的辦法：如果數(shù)據(jù)算完即廢，等需要的時候，我再想辦法從個什么地方拿回來，那不就省了一筆存儲空間嗎？

沿著這個思路，我們逐一來看ZeRO是如何遞進做存儲優(yōu)化的。

2.1 : Optimizer State Partitioning

首先，從 optimizer state開始優(yōu)化。將optimizer state分成若干份，每塊GPU上各自維護一份。這樣就減少了相當(dāng)一部分的顯存開銷。如下圖：

復(fù)習(xí)一下，此時W=fp16，G=fp16，O=fp32。此時，整體數(shù)據(jù)并行的流程如下：

（1）每塊GPU上存一份完整的參數(shù)W。將一個batch的數(shù)據(jù)分成3份，每塊GPU各吃一份，做完一輪foward和backward后，各得一份梯度。

（2）對梯度做一次AllReduce，得到完整的梯度G，產(chǎn)生單卡通訊量。為了表達簡明，這里通訊量我們就不再換算成byte了，而直接根據(jù)參數(shù)量來計算。對AllReduce（reduce-scatter + all-gather）不熟悉的朋友，可以先去看上一篇文章。

（3）得到完整梯度G，就可以對W做更新。我們知道W的更新由optimizer states和梯度共同決定。由于每塊GPU上只保管部分optimizer states，因此只能將相應(yīng)的W（藍色部分）進行更新。（2）和（3）可以用下圖表示：

（4）此時，每塊GPU上都有部分W沒有完成更新（圖中白色部分）。所以我們需要對W做一次All-Gather，從別的GPU上把更新好的部分W取回來。產(chǎn)生單卡通訊量。

做完后，設(shè)GPU個數(shù)為，顯存和通訊量的情況如下：

假設(shè)各變量大小如表格第二列所示，那么在增加1.5倍單卡通訊開銷的基礎(chǔ)上，將單卡存儲降低了4倍。看起來是個還不錯的trade-off，那么還能做得更好嗎？

2.2

現(xiàn)在，更近一步，我們把梯度也拆開，每個GPU格子維護一塊梯度。

此時，數(shù)據(jù)并行的整體流程如下：

（1）每塊GPU上存一份完整的參數(shù)W。將一個batch的數(shù)據(jù)分成3份，每塊GPU各吃一份，做完一輪foward和backward后，算得一份完整的梯度（下圖中綠色+白色）。

（2）對梯度做一次Reduce-Scatter，保證每個GPU上所維持的那塊梯度是聚合梯度。例如對GPU1，它負責(zé)維護G1，因此其他的GPU只需要把G1對應(yīng)位置的梯度發(fā)給GPU1做加總就可。匯總完畢后，白色塊對GPU無用，可以從顯存中移除。單卡通訊量。（1）和（2）見下圖：

（3）每塊GPU用自己對應(yīng)的O和G去更新相應(yīng)的W。更新完畢后，每塊GPU維持了一塊更新完畢的W。同理，對W做一次All-Gather，將別的GPU算好的W同步到自己這來。單卡通訊量。

再次比對下顯存和通訊量：

和樸素DP相比，存儲降了8倍，單卡通訊量持平，好像更牛皮了呢！那么，還可以優(yōu)化嗎？

2.3

看到這里，也許你有點感覺了，ZeRO的思想就是：萬物皆可切，萬物皆可拋。所以現(xiàn)在，我們把參數(shù)也切開。每塊GPU置維持對應(yīng)的optimizer states，gradients和parameters（即W）。

數(shù)據(jù)并行的流程如下：

（1）每塊GPU上存一份完整的參數(shù)W。將一個batch的數(shù)據(jù)分成3份，每塊GPU各吃一份。

（2）做forward時，對W做一次All-Gather，取回分布在別的GPU上的W，得到一份完整的W，單卡通訊量。forward做完，立刻把不是自己維護的W拋棄。

（3）做backward時，對W做一次All-Gather，取回完整的W，單卡通訊量。backward做完，立刻把不是自己維護的W拋棄。

（4）做完backward，算得一份完整的梯度G，對G做一次Reduce-Scatter，從別的GPU上聚合自己維護的那部分梯度，單卡通訊量。聚合操作結(jié)束后，立刻把不是自己維護的G拋棄。

（5）用自己維護的O和G，更新W。由于只維護部分W，因此無需再對W做任何AllReduce操作。

顯存和通訊量如下：

到這一步，我們用1.5倍的通訊開銷，換回近120倍的顯存。只要梯度計算和異步更新做的好，通訊時間大部分可以被計算時間隱藏，因此這樣的額外通訊開銷，也是劃算的。

到這里，我們可以放出原始論文中的說明圖了，經(jīng)過以上分析，這張說明圖是不是瞬間就能看懂了。不得不吐槽下，雖然ZeRO的設(shè)計不復(fù)雜，但對應(yīng)論文寫得真是邏輯跳躍，晦澀難懂...

仔細一想，ZeRO其實掌握了降本增效的精髓：用完即棄，需要再補。反正我補一個和你差不多的，也不會花費很多通（找）訊（人）時間，還大大降低了我的成本。模型的每一層多算（造）幾（輪）遍（子）有啥關(guān)系呢，反正在我的預(yù)算里每個人都一刻不停地干活，就行啦！

2.4 ZeRO VS 模型并行

知道模型并行的朋友，可能會想，既然ZeRO都把參數(shù)W給切了，那它應(yīng)該是個模型并行呀？為什么要歸到數(shù)據(jù)并行里呢？

其實ZeRO是模型并行的形式，數(shù)據(jù)并行的實質(zhì)。

模型并行，是指在forward和backward的過程中，我只需要用自己維護的那塊W來計算就行。即同樣的輸入X，每塊GPU上各算模型的一部分，最后通過某些方式聚合結(jié)果。

但對ZeRO來說，它做forward和backward的時候，是需要把各GPU上維護的W聚合起來的，即本質(zhì)上還是用完整的W進行計算。它是不同的輸入X，完整的參數(shù)W，最終再做聚合。

因為下一篇要寫模型并行Megatron-LM，因此現(xiàn)在這里羅列一下兩者的對比。

三、ZeRO-R

說完了以上對model states的顯存優(yōu)化，現(xiàn)在來看對residual states的優(yōu)化。

3.1 : Partitioned Activation Checkpointing

前面說過，對activation的存儲是靈活的。不像optimizer states，gradients和parameters對模型更新是必須的，activation只是起到加速梯度計算的作用。因此，在哪幾層保存activation，保存哪些activation都是可以靈活設(shè)置的。同樣，我們也可以仿照以上切割方式，每塊GPU上只維護部分的activation，需要時再從別的地方聚合過來就行。需要注意的是，activation對顯存的占用一般會遠高于模型本身，通訊量也是巨大的，所以這塊要靈活、有效地實驗設(shè)計。

3.2 : Constant Size Buffer

固定大小的內(nèi)存buffer，它的目的在于：

提升帶寬利用率。當(dāng)GPU數(shù)量上升，GPU間的通訊次數(shù)也上升，每次的通訊量可能下降（但總通訊量不會變）。數(shù)據(jù)切片小了，就不能很好利用帶寬了。所以這個buffer起到了積攢數(shù)據(jù)的作用：等數(shù)據(jù)積攢到一定大小，再進行通訊。

使得存儲大小可控。在每次通訊前，積攢的存儲大小是常量，是已知可控的。更方便使用者對訓(xùn)練中的存儲消耗和通訊時間進行預(yù)估。

3.3 : Memory Defragmentation

在前文提過，設(shè)置機制，對碎片化的存儲空間進行重新整合，整出連續(xù)的存儲空間。防止出現(xiàn)總存儲足夠，但連續(xù)存儲不夠而引起的存儲請求fail。

四、ZeRO-Offload與ZeRO-Infinity

最后，簡單介紹一下ZeRO-Offload。它的核心思想是：顯存不夠，內(nèi)存來湊。如果我把要存儲的大頭卸載(offload)到GPU上，而把計算部分放到GPU上，這樣比起跨機，是不是能既降顯存，也能減少一些通訊壓力呢？

ZeRO-Offload的做法是：

forward和backward計算量高，因此和它們相關(guān)的部分，例如參數(shù)W（fp16），activation，就全放入GPU。

update的部分計算量低，因此和它相關(guān)的部分，全部放入CPU中。例如W(fp32)，optimizer states（fp32）和gradients(fp16)等。

具體切分如下圖：

ZeRO-infinity也是同理，它們在解決的事情都是：找個除GPU之外的地方，存數(shù)據(jù)。感興趣的朋友可以深入研究，這里就不展開了。