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

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

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

一、存儲(chǔ)消耗

1.1 存儲(chǔ)分類

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

1.3 存儲(chǔ)大小

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

推薦閱讀：

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

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

一、存儲(chǔ)消耗

1.1 存儲(chǔ)分類

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

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

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

gradients：模型梯度

parameters：模型參數(shù)W

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

activation：激活值。在流水線并行中我們?cè)敿?xì)介紹過(guò)。在backward過(guò)程中使用鏈?zhǔn)椒▌t計(jì)算梯度時(shí)會(huì)用到。有了它算梯度會(huì)更快，但它不是必須存儲(chǔ)的，因?yàn)榭梢酝ㄟ^(guò)重新做Forward來(lái)算它。

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

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

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

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

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

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

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

在forward開(kāi)始之前，額外開(kāi)辟一塊存儲(chǔ)空間，將fp32 parameter減半到fp16 parameter。

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

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

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

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

1.3 存儲(chǔ)大小

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

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

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

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

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

因?yàn)閍ctivation的這種靈活性，納入它后不方便衡量系統(tǒng)性能隨模型增大的真實(shí)變動(dòng)情況。因此在這里不考慮它，在后面會(huì)單開(kāi)一塊說(shuō)明對(duì)activation的優(yōu)化。

二、ZeRO-DP

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

注意到，在整個(gè)訓(xùn)練中，有很多states并不會(huì)每時(shí)每刻都用到，舉例來(lái)說(shuō)；

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

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

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

諸如此類

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

沿著這個(gè)思路，我們逐一來(lái)看ZeRO是如何遞進(jìn)做存儲(chǔ)優(yōu)化的。

2.1 : Optimizer State Partitioning

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

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

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

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

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

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

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

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

2.2

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

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

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

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

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

再次比對(duì)下顯存和通訊量：

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

2.3

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

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

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

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

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

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

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

顯存和通訊量如下：

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

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

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

2.4 ZeRO VS 模型并行

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

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

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

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

因?yàn)橄乱黄獙?xiě)模型并行Megatron-LM，因此現(xiàn)在這里羅列一下兩者的對(duì)比。

三、ZeRO-R

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

3.1 : Partitioned Activation Checkpointing

前面說(shuō)過(guò)，對(duì)activation的存儲(chǔ)是靈活的。不像optimizer states，gradients和parameters對(duì)模型更新是必須的，activation只是起到加速梯度計(jì)算的作用。因此，在哪幾層保存activation，保存哪些activation都是可以靈活設(shè)置的。同樣，我們也可以仿照以上切割方式，每塊GPU上只維護(hù)部分的activation，需要時(shí)再?gòu)膭e的地方聚合過(guò)來(lái)就行。需要注意的是，activation對(duì)顯存的占用一般會(huì)遠(yuǎn)高于模型本身，通訊量也是巨大的，所以這塊要靈活、有效地實(shí)驗(yàn)設(shè)計(jì)。

3.2 : Constant Size Buffer

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

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

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

3.3 : Memory Defragmentation

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

四、ZeRO-Offload與ZeRO-Infinity

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

ZeRO-Offload的做法是：

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

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

具體切分如下圖：

ZeRO-infinity也是同理，它們?cè)诮鉀Q的事情都是：找個(gè)除GPU之外的地方，存數(shù)據(jù)。感興趣的朋友可以深入研究，這里就不展開(kāi)了。