ZeRO 和分片優化器
ZeRO（零冗餘優化器）透過跨 GPU 分片優化器狀態、梯度和權重，消除了資料並行性中浪費的記憶體重複。
概述
It lets you train enormous models with the simplicity of data parallelism but a fraction of the per-GPU memory.
深入探討
在普通資料並行中，每個 GPU 都會儲存優化器狀態、梯度和參數的冗餘完整副本，這是極大的浪費，特別是對於 Adam 來說，其中優化器狀態可能是模型本身大小的幾倍。 DeepSpeed 中的 Microsoft 引入的 ZeRO 透過在 GPU 上對這些張量進行分區來消除這種冗餘，這樣每個設備只擁有一個切片。 ZeRO 分為三個漸進階段：第 1 階段將優化器狀態分片，第 2 階段新增梯度分片，第 3 階段將參數本身分片。根據需要，GPU 透過通訊、計算收集丟失的切片，然後釋放它們。結果是每個 GPU 的記憶體顯著降低，支援數十億到兆參數的訓練，同時保持資料並行的簡單程式設計模型。
技術洞察
ZeRO 用額外的通信來節省記憶體。在第 3 階段，在某個層的前向傳遞之前，全收集將該層的完整參數收集到每個 GPU 上；之後，非擁有的切片將被丟棄以回收記憶體。梯度是減少分散的，因此每個 GPU 僅保留與其擁有的參數相符的梯度切片。 PyTorch 的 FSDP（完全分片資料並行）原生實現了相同的想法，將模組包裝為動態分片和重新分片。
戰略影響
成本與預算
多年來，架構決策決定著效能和營運成本。
更明確的決策
技術教育幫助團隊選擇正確的堆疊，而不僅僅是最新的堆疊。
品質管控
更好的工程選擇可以減少生產中的可靠性事故。
ZeRO 和分片優化器的未來
分片正在成為大規模訓練的預設選擇，而不是一種奇特的選擇。期望與卸載進行更深入的整合（透過 ZeRO-Infinity 將切片推送到 CPU 或 NVMe），透過計算更好地重疊全收集和減少分散以隱藏其成本，以及與張量和管道並行性的組合。隨著模型不斷增長，記憶體高效的分片優化器對於將模型適應實際的硬體預算至關重要。
現實世界的實施
使用 DeepSpeed ZeRO Stage 2 微調數十億參數的語言模型，否則模型會溢出 GPU 記憶體。
使用 PyTorch FSDP 進行訓練，跨 GPU 分片參數、梯度和優化器狀態，並根據需要每層收集它們。
應用 ZeRO-Offload 將優化器狀態推送到 CPU 內存，讓單一 GPU 訓練比其 VRAM 大很多倍的模型。
當 GPU 和 CPU 記憶體耗盡時，透過從 NVMe 儲存串流傳輸參數分片，使用 ZeRO-Infinity 擴展萬億參數模型。
風險與防護欄
優化一項基準測試可以隱藏更廣泛的系統弱點。
基礎設施和維護成本常常被低估。
隨著系統變得更加複雜，安全性和可觀察性差距可能會擴大。
實施路線圖
在實施之前定義延遲、品質和成本目標。
在實際負載和資料條件下進行基準測試。
儀器監控錯誤、漂移和使用者影響。
在擴展之前準備回滾和事件回應路徑。
不斷探索
Free newsletter
Get the daily AI briefing
Three verified AI stories every weekday morning, written in plain English. Free forever, no ads.
One email each weekday. Unsubscribe in one click. We never sell or share your address.
Test yourself
Take the ZeRO and Sharded Optimizers quiz
Instant feedback on every answer, and a shareable certificate with a verifiable ID once you pass a course.
Support free AI education. AI Understanding is a 501(c)(3) nonprofit — no ads, no paywall, ever. Make a donation
常見問題
What is ZeRO and Sharded Optimizers?
ZeRO（零冗餘優化器）透過跨 GPU 分片優化器狀態、梯度和權重，消除了資料並行性中浪費的記憶體重複。它可以讓您透過簡單的資料並行性來訓練巨大的模型，但只需要一小部分 GPU 記憶體。
與普通資料並行相比，ZeRO 消除了哪些冗餘？
標準資料並行性在每個 GPU 上儲存優化器狀態、梯度和權重的完整副本； ZeRO 對這些進行分片，因此每個 GPU 僅保留一個切片。
為什麼優化器狀態通常是 Adam 佔用最大記憶體的地方？
Adam 維護每個參數的運行估計，例如第一和第二時刻，與 fp32 主權重相結合可以使模型本身的大小相形見絀。
ZeRO 第 3 階段分片有哪些第 1 期和第 2 期沒有的分片？
階段 1 對優化器狀態進行分片，第 2 階段新增梯度，而第 3 階段則進一步將模型參數分片到 GPU 上。
在 ZeRO Stage 3 中，GPU 如何取得層前向傳遞所需的完整參數？
在計算某個層之前，全收集會在每個 GPU 上組裝其完整參數；一旦完成，非擁有的切片將被釋放以回收記憶體。
哪個 PyTorch 功能本身實現了 ZeRO 風格的分片？
PyTorch 的完全分片資料並行 (FSDP) 對參數、梯度和最佳化器狀態進行分片，動態收集和重新分片，鏡像 ZeRO。