機器之心轉載
來源：知乎
作者：微軟DeepSpeed
大型 AI 模型正在改變數字世界。基于大型語言模型 (LLM) 的 Turing-NLG、ChatGPT 和 GPT-4 等生成語言模型用途廣泛，能夠執行摘要、代碼生成和翻譯等任務。同樣，DALL・E、Microsoft Designer 和 Bing Image Creator 等大型多模态生成模型可以生成藝術、建築、視訊和其他數字資産，使内容創作者、建築師和工程師能夠探索全新的創意生産力。
然而，訓練這些大型模型需要在數百甚至數千個 GPU 裝置上使用大量記憶體和計算資源。例如，訓練 Megatron-Turing NLG 530B 模型需要使用超過 4,000 個 NVidia A100 GPU。有效地利用這些資源需要一個複雜的優化系統，以将模型合理配置設定到各個裝置的記憶體中，并有效地并行化這些裝置上的計算。同時，為了使深度學習社群能夠輕松進行大型模型訓練，這些優化必須易于使用。
DeepSpeed 的 ZeRO 優化系列為這些挑戰提供了強大的解決方案，并已廣泛用于大型深度學習模型例如 TNLG-17B、Bloom-176B、MPT-7B、Jurrasic-1 的訓練中 。盡管它具有變革性的能力 ，在一些關鍵場景中，ZeRO 會在 GPU 之間産生大量資料傳輸開銷，這降低了訓練效率。這種情況特别發生在以下場景中：a) 全局 batch size 較小，而 GPU 數量多，這導緻每個 GPU 上 batch size 較小，需要頻繁通信；或者 b) 在低端叢集上進行訓練，其中跨節點網絡帶寬有限，導緻高通信延遲。在這些情況下，ZeRO 的訓練效率會受到限制。
為了解決這些限制，我們釋出了 ZeRO++ 。ZeRO++ 相比 ZeRO 将總通信量減少了 4 倍，而不會影響模型品質。這有兩個關鍵意義：
1. ZeRO++ 加速大型模型預訓練和微調
每個 GPU 上 batch size 較小時：無論是在數千個 GPU 上預訓練大型模型，還是在數百個甚至數十個 GPU 上對其進行微調，當每個 GPU 的 batch size 較小時，ZeRO++ 提供比 ZeRO 高 2.2 倍的吞吐量，直接減少訓練時間和成本。
低帶寬計算叢集: ZeRO++ 使低帶寬叢集能夠實作與帶寬高 4 倍的高端叢集類似的吞吐量。是以，ZeRO++ 可以跨更廣泛的叢集進行高效的大型模型訓練。
2. ZeRO++ 加速 ChatGPT 類的 RLHF 訓練
雖然 ZeRO++ 主要是為訓練而設計的，但它的優化也自動适用于 ZeRO-Inference，因為通信開銷對于 ZeRO 的訓練和推理同樣适用。是以，ZeRO++ 可以提高人類回報強化學習 (RLHF) 等算法的效率，因為 RLHF 結合了訓練和推理。
通過與 DeepSpeed-Chat 的內建，與原始 ZeRO 相比，ZeRO++ 可以将 RLHF 訓練的生成階段效率提高多達 2 倍，強化學習訓練階段效率提高多達 1.3 倍。
接下來，我們将更深入地解釋 ZeRO 及其通信開銷，并讨論 ZeRO++ 中為解決這些問題而進行的關鍵優化。然後我們将展示 ZeRO++ 對不同模型大小、批量大小和帶寬限制的訓練吞吐量的影響。我們還将讨論 ZeRO++ 如何應用于 DeepSpeed-Chat，以加速使用 RLHF 的對話模型的訓練。
ZeRO++ 詳解
圖2：ZeRO optimizer 工作流程圖（此為部分展示，完整流程請看知乎原文）
ZeRO 是資料并行 (Data Parallelism) 的一種記憶體高效版本，其中模型狀态會被分割儲存在所有 GPU 上，而不需要在訓練期間使用基于 gather/broadcas 的通信進行複制和重建。這使 ZeRO 能夠有效地利用所有裝置的聚合 GPU 記憶體和計算力，同時提供簡單易用的資料并行訓練。
假設模型大小為 M。在前向傳播過程中，ZeRO 執行全收集 / 廣播 (all-gather/broadcast) 操作以在需要之時為每個模型層收集參數（總共大小為 M）。在向後傳遞中，ZeRO 對每一層的參數采用類似的通信模式來計算其局部梯度（總大小為 M）。此外，ZeRO 在對每個局部梯度計算完畢後會立刻使用 reduce 或 reduce-scatter 通信進行平均和分割儲存（總大小為 M）。是以，ZeRO 總共有 3M 的通信量，平均分布在兩個全收集 / 廣播 (all-gather/broadcast) 和一個減少分散 / 減少 (reduce-scatter/reduce) 操作中。
為了減少這些通信開銷，ZeRO++ 進行了三組通信優化，分别針對上述三個通信集合：
圖 3:qwZ 的分區量化圖例
ZeRO 通信過程中的權重量化 (qwZ)
首先，為了減少 all-gather 期間的參數通信量，我們采用權重量化在通信前将每個模型參數從 FP16（兩個位元組）動态縮小為 INT8（一個位元組）資料類型，并在通信後對權重進行反量化。然而，簡單地對權重進行量化會降低模型訓練的準确性。為了保持良好的模型訓練精度，我們采用分區量化，即對模型參數的每個子集進行獨立量化。目前尚且沒有針對分區量化的高性能現有實作。是以，我們自行從頭開始實作了一套高度優化的量化 CUDA 核心，與基本量化相比，精度提高 3 倍，速度提高 5 倍。
圖 4: 權重的分層分割存儲 (hpZ)
ZeRO 模型權重的分層分割存儲 (hpZ)
其次，為了減少向後傳遞期間全收集 (all-gather) 權重的通信開銷，我們用 GPU 記憶體進行通信。更具體地說，我們不像在 ZeRO 中那樣将整個模型權重分布在所有機器上，而是在每台機器中維護一個完整的模型副本。以更高的記憶體開銷為代價，這允許我們用機器内的模型權重全收集 / 廣播 (all-gather/broadcast) 代替昂貴的跨機器全收集 / 廣播 (all-gather/broadcast)，由于機器内通信帶寬更高，這使得通信速度大幅提升。
圖 5: qgZ 端到端的工作流程
ZeRO 通信過程中梯度量化 (qgZ)
第三，要降低梯度的 reduce-scatter 通信成本更具挑戰性。因為直接應用量化來減少通信量是不可行的。即使我們使用分區量化來降低量化誤差，梯度 reduce 也會累積并放大量化誤差。為了解決這個問題，我們隻在通信之前量化梯度，但在任何 reduce 操作之前将它們反量化到原有精度。為了有效地做到這一點，我們發明了一種名為 qgZ 的基于 all-to-all 的新型量化梯度通信範式，它在功能上等同于壓縮的歸約 - 分散 (reduce-scatter) 操作。
qgZ 旨在解決兩個挑戰：i) 如果我們簡單地在 INT4/INT8 中實施 reduce-scatter 會導緻顯著精度損失，以及 ii) 在傳統 tree 或 ring-based reduce-scatter 中使用量化需要一長串量化和反量化步驟，這直接導緻誤差積累和顯著的延遲，即使我們在全精度上進行 reduce。為了解決這兩個挑戰，qgZ 不使用 tree 或 ring-based reduce-scatter 算法，而是基于一種新穎的分層 all-to-all 方法。
qgZ 中有三個主要步驟：
- 梯度切片重新排序；
- 節點内通信和 reduce；
- 節點間通信和 reduce。
首先，在任何通信發生之前，我們對梯度進行切片并對張量切片重新排序，以保證通信結束時每個 GPU 上的最終梯度位置（即圖 5 中的綠色塊）是正确的。其次，我們量化重新排序的梯度切片，在每個節點内進行 all-to-all 通信，從 all-to-all 中對接收到的梯度切片進行反量化，并進行局部 reduce。第三，我們再次量化局部 reduce 後的梯度，進行節點間的 all-to-all 通信，再次對接收到的梯度進行反量化，并計算最終的高精度梯度 reduce，得到圖 5 中綠色塊的結果。
這種分層方法的原因是為了減少跨節點通信量。更準确地說，給定每個節點 N 個 GPU、M 的模型大小和 Z 的量化比率，單跳 all-to-all 将生成 M*N/Z 跨節點流量。相比之下，通過這種分層方法，我們将每個 GPU 的跨節點流量從 M/Z 減少到 M/(Z*N)。是以，總通信量從 M*N/Z 減少到 M*N/(Z*N) = M/Z。我們通過重疊節點内和節點間通信以及融合 CUDA 核心來進一步優化 qgZ 的端到端延遲（張量切片重新排序 (Tensor Slice Reordering)+ 節點内量化 (Intra-node quantization)）和（節點内反量化 (Intra-node Dequantization) + 節點内梯度整合 (Intra-node Reduction) + 節點間量化 (inter-node quantization)）。
通信總量優化
通過結合以上所有三個元件，我們将跨節點通信量從 3M 減少到 0.75M。更具體地說，我們使用 qwZ 将模型權重的前向全收集 / 廣播從 M 減少到 0.5M。我們使用 hpZ 消除了反向傳播期間的跨節點 all-gather，将通信從 M 減少到 0。最後，我們使用 qgZ 将反向傳播期間的跨節點 reduce-scatter 通信從 M 減少到 0.25M。
ZeRO++ 加速大型語言模型訓練
在這裡，我們展示了 ZeRO++ 在 384 個 Nvidia V100 GPU 上的真實 LLM 訓練場景的測試結果。
圖 6: 在 384 個 V100 GPU 上的各種模型大小下 ZeRO++ 與 ZeRO 的吞吐量，節點間使用 4 個 Infiniband (IB) 進行互連，每個以 100 Gbps 運作。
在 GPU 小 batch size 情況下 ZeRO++ 實作更高的訓練效率
高帶寬叢集：如圖 6 所示，我們首先展示了 ZeRO++ 相對于 ZeRO 的吞吐量改進，針對不同的模型大小和微批量 (micro-batch size) 大小，測試使用 4x Infiniband (IB) 以實作 400Gbps 跨節點互連帶寬，每個以 100Gbps 運作。在 micro-batch size 為每 GPU 1k tokens 時，ZeRO++ 比 ZeRO-3 的吞吐量提高了 28% 到 36%。對于 2k tokens micro-batch size 大小，ZeRO++ 比 ZeRO-3 實作了 24% 到 29% 的吞吐量增益。
圖 7: 在 384 個 V00 GPU 上 100Gbps 跨節點帶寬時各種 LLM 的吞吐量
低帶寬叢集：在 100Gbps 等低帶寬網絡環境中，ZeRO++ 的性能明顯優于 ZeRO-3。如圖 7 所示，與 ZeRO-3 相比，ZeRO++ 在端到端吞吐量方面實作了高達 2.2 倍的加速。平均而言，ZeRO++ 比 ZeRO-3 基線實作了大約 2 倍的加速。
圖 8: ZeRO++ 以顯著降低的帶寬實作高帶寬叢集性能
實作高帶寬 ZeRO 和低帶寬 ZeRO++ 叢集之間的模型訓練效率等效
此外，與 ZeRO 在高得多的帶寬環境下相比，ZeRO ++ 可以在低帶寬叢集中實作相當的系統吞吐量。如圖 8 所示，對于 18B 和 138B 模型大小，具有 200Gbps 跨節點帶寬的 ZeRO++ 可以達到與 800Gbps 跨節點帶寬的 ZeRO-3 相似的 TFLOP。
鑒于 ZeRO++ 出色的可擴充性，我們将 ZeRO++ 視為用于訓練大型 AI 模型的下一代 ZeRO。
DeepSpeed-Chat 與 ZeRO++ 結合用于 RLHF 訓練
RLHF 訓練簡介
ChatGPT 類模型由 LLM 提供支援，并使用 RLHF 進行微調。RLHF 由生成（推理）階段和訓練階段組成。在生成階段，演員 (actor) 模型将部分對話作為輸入，并使用一系列前向傳遞生成響應。然後在訓練階段，評論 (critic) 模型根據品質對生成的響應進行排名，為演員模型提供強化信号。使用這些排名對參與者模型進行微調，使其能夠在後續疊代中生成更準确和适當的響應。
RLHF 訓練帶來了巨大的記憶體壓力，因為它使用了四種模型（演員、參考、評論、獎勵）。常見的解決方案是采用低秩自适應訓練 (LoRA) 來解決 RLHF 的記憶體壓力。LoRA 當機了預訓練模型的權重，并将可訓練的秩分解矩陣注入到 Transformer 架構的每一層中，顯著減少了可訓練參數的數量。LoRA 通過減少記憶體使用來加速 RLHF，允許更大的批處理 (batch) 大小，進而大大提高吞吐量。
DeepSpeed-Chat with ZeRO++ 用于 RLHF 訓練
圖 9: ZeRO++ 加速了 RLHF 訓練的生成和訓練階段
ZeRO++ 在 RLHF + LoRA 的場景下有着獨特的應用，因為大多數模型權重都被當機了。這意味着 ZeRO++ 可以将這些當機的權重量化儲存到 INT4/8 中，而不是将它們存儲在 fp16 中并在每次通信操作之前對其進行量化。通信後的反量化仍然是為了讓權重為計算做好準備，但反量化後的權重在計算後被簡單地丢棄。
以這種方式使用 ZeRO++ 進行 RLHF 訓練可以減少記憶體使用和通信量。這意味着通過減少通信以及由于減少記憶體使用而啟用更大的批處理大小來提高訓練吞吐量。在生成階段，ZeRO++ 使用 hpZ 将所有權重通信保持在每個節點内，以利用更高的節點内通信帶寬，減少通信量，進一步提高生成吞吐量。
ZeRO++ 已內建到 DeepSpeed-Chat 中，以支援 ChatGPT 類模型的 RLHF 訓練。在圖 9 中，我們比較了不同大小的 actor 模型的 RLHF 生成吞吐量。測試配置為 32 個 V100 GPU ，actor 模型大小為 30B 和 66B 以測試 ZeRO 和 ZeRO++ 性能。結果表明，ZeRO++ 的 RLHF 生成吞吐量比 ZeRO 高出 2.25 倍。我們還展示了在 16 個 V100 GPU 上訓練階段的加速，其中 ZeRO++ 實作了比 ZeRO 高 1.26 倍的吞吐量，這是由于 ZeRO++ 支援的更低通信量和更大批量大小。
DeepSpeed ZeRO++ 現已釋出！
貢獻者:
DeepSpeed 團隊的以下人員的貢獻使該項目成為可能：
Guanhua Wang, Heyang Qin, Sam Ade Jacobs, Connor Holmes, Samyam Rajbhandari, Olatunji Ruwase, Ammar Ahmad Awan, Jeff Rasley, Michael Wyatt, Yuxiong He (team lead)
本文轉載自微軟DeepSpeed組。