并行策略

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

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

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

經過五次循環之后，每個GPU再將各自負責的那部分數據直接同步給其他顯卡，避免多次廣播運算。

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

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

初識ZeRO

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

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

Adam在SGD基礎上，為每個參數梯度增加了一階動量（momentum）和二階動量（variance）1。混合精度訓練，字如其名，同時存在fp16和fp32兩種格式的數值，其中模型參數、模型梯度都是fp16，此外還有fp32的模型參數，如果優化器是Adam，則還有fp32的momentum和variance。

<大白話， Adam 是一種可以替代傳統隨機梯度下降過程的一階優化算法，它能基于訓練數據迭代地更新神經網絡權重!!>

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

• 模型狀態（Model States）由模型參數（fp16）、模型梯度（fp16）和Adam狀態(FP32的模型參數副本，FP32的Momentum和FP32的Variance)。

假設模型參數量為Φ ，則共需要2Φ+2Φ+(4Φ+4Φ+4Φ)=4Φ+12Φ=16Φ 字節存儲空間。從這個統計中可以發現，其實Adam優化器占據了75%的空間。<若讀者要微調3B的模型，那么則需要48B的空間！！！>

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

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

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

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

在DeepSpeed的訓練框架中， Pos對應ZeRO-1， Pos+g對應ZeRO-2， Pos+g+p對應ZeRO-3，而一般在生產中ZeRO-1就足夠使用。

本文轉載自??魯班模錘??，作者： 龐德公 ????