大型 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 ZeRO++ 並讓 AI 社區中的每個人都可以使用它。請訪問我們的 GitHub 頁面以獲取 LLM 訓練教程。用於 DeepSpeed-Chat 的 ZeRO++ 將在未來幾周內發佈。有關 ZeRO++ 的更多技術細節，請查看我們的 arxiv 論文。
DeepSpeed-ZeRO++ 是 DeepSpeed 生態系統的一部分。要了解更多資訊，請訪問我們的網站，在那裡您可以找到詳細的部落格文章、教程和有用的文件。
您還可以在我們的英文 Twitter、日文 Twitter 和中文知乎 上獲取最新的 DeepSpeed 新聞。
DeepSpeed 歡迎您的貢獻！我們鼓勵您在 DeepSpeed GitHub 頁面上報告問題、貢獻 PR 並加入討論。有關更多詳細資訊，請參閱我們的貢獻指南。我們對與大學、研究實驗室和公司的合作持開放態度。對於此類請求（以及其他不適合 GitHub 的請求），請直接發送電子郵件至 deepspeed-info@microsoft.com。
貢獻者:
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組。