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

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

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

一、存儲消耗

1.1 存儲分類

1.2 混合精度訓練

1.3 存儲大小

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

推薦閱讀：

圖解大模型訓練之：流水線并行，以GPipe為例

圖解大模型訓練之：數據并行上篇（DP, DDP）

一、存儲消耗

1.1 存儲分類

首先，我們來看在大模型訓練的過程中，GPU都需要存什么內容。

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

optimizer states：Adam優化算法中的momentum和variance

gradients：模型梯度

parameters：模型參數W

Residual States指并非模型必須的，但在訓練過程中會額外產生的內容，具體包括：

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

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

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

1.2 精度混合訓練

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

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

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

存儲一份fp32的parameter，momentum和variance（統稱model states）

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

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

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

當模型收斂后，fp32的parameter就是最終的參數輸出。

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

1.3 存儲大小

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

因為采用了Adam優化，所以才會出現momentum和variance，當然你也可以選擇別的優化辦法。因此這里為了更通用些，記模型必存的數據大小為。因此最終內存開銷為：

另外，這里暫不將activation納入統計范圍，原因是：

activation不僅與模型參數相關，還與batch size相關

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

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

二、ZeRO-DP

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

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

Adam優化下的optimizer states只在最終做update時才用到

數據并行中，gradients只在最后做AllReduce和updates時才用到

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

諸如此類

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

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

2.1 : Optimizer State Partitioning

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

復習一下，此時W=fp16，G=fp16，O=fp32。此時，整體數據并行的流程如下：

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

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

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

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

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

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

2.2

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

此時，數據并行的整體流程如下：

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

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

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

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

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

2.3

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

數據并行的流程如下：

（1）每塊GPU上存一份完整的參數W。將一個batch的數據分成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上聚合自己維護的那部分梯度，單卡通訊量。聚合操作結束后，立刻把不是自己維護的G拋棄。

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

顯存和通訊量如下：

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

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

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

2.4 ZeRO VS 模型并行

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

其實ZeRO是模型并行的形式，數據并行的實質。

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

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

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

三、ZeRO-R

說完了以上對model states的顯存優化，現在來看對residual states的優化。

3.1 : Partitioned Activation Checkpointing

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

3.2 : Constant Size Buffer

固定大小的內存buffer，它的目的在于：

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

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

3.3 : Memory Defragmentation

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

四、ZeRO-Offload與ZeRO-Infinity

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

ZeRO-Offload的做法是：

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

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

具體切分如下圖：

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