ZeRO 優化（ZeRO Optimization）是什麼？

Microsoft 提出的分布式訓練優化技術，通過將梯度、優化器狀態和模型參數分片存儲在多個 GPU 上，大幅降低記憶體占用，支持訓練超大規模模型。

核心概念

ZeRO 優化基於一個關鍵觀察：在數據並行訓練中，每個 GPU 都存儲模型參數、梯度和優化器狀態的完整副本。這造成巨大冗餘，記憶體占用是必要的數倍。ZeRO 的思想是：不需要每個 GPU 都有完整副本，通過分片和通信可以達到相同的效果，同時大幅節省記憶體。

ZeRO 分為三個級別：

• ZeRO-1：梯度分片，每個 GPU 只存儲部分梯度
• ZeRO-2：梯度 + 優化器狀態分片
• ZeRO-3：完全分片，包括模型參數

更高級別的優化帶來更多記憶體節省（可達 N 倍，N 是 GPU 數量），但通信開銷也增加。

運作原理

ZeRO-1 的梯度分片：

1. 數據並行訓練中，每個 GPU 計算全部梯度
2. ZeRO-1 將梯度按參數分片：GPU 0 存儲參數 0-k 的梯度，GPU 1 存儲參數 k+1-2k 的梯度，等等
3. 優化器步驟中，每個 GPU 只更新自己分片的參數
4. 然後通過全聚集（all-gather）操作讓所有 GPU 獲得最新參數

ZeRO-2 的梯度和優化器狀態分片：

1. 在 ZeRO-1 基礎上，優化器狀態（如 Adam 的動量項）也分片存儲
2. 優化器只在自己的 GPU 上維護分配的參數的狀態
3. 進一步節省記憶體

ZeRO-3 的完全分片：

1. 模型參數也分片存儲
2. 前向傳播時，通過全聚集獲得需要的參數
3. 反向傳播後，丟棄不需要的參數副本
4. 最激進的優化，記憶體節省最大，但通信開銷也最大

實際應用

ZeRO 在超大規模模型訓練中發揮了關鍵作用。OpenAI 的 GPT-3（1750 億參數）、Microsoft 的 Turing-NLG（170 億參數）等模型都使用了 ZeRO 進行訓練。DeepSpeed 框架（集成了 ZeRO）成為訓練大規模語言模型的標準工具。

在自然語言處理中，ZeRO 使得在有限的硬體上訓練超大規模預訓練模型成為可能。許多研究機構使用 ZeRO 訓練自己的大語言模型，而無需購買超大規模計算集群。

ZeRO 也應用於微調場景。使用 ZeRO-2 或 ZeRO-3 能在消費級 GPU 上微調超大預訓練模型（如 GPT-3 規模），這在實踐中很有價值。

在計算機視覺中，ZeRO 支持訓練超大規模視覺模型。Vision Transformer 的大規模版本、多模態模型（CLIP、BLIP 等）都可以使用 ZeRO 加速訓練。

ZeRO 也用於強化學習，特別是大規模策略網路（如遊戲 AI）的訓練。

常見誤區

誤區一：認為 ZeRO-3 總是最好選擇。ZeRO-3 記憶體節省最大，但通信開銷也最大。在高延遲網路或節點較少的情況下，ZeRO-1 或 ZeRO-2 可能更快。應根據硬體配置選擇。

誤區二：認為 ZeRO 不需要調整其他超參數。實際上，使用 ZeRO 時可能需要調整學習率、批量大小、通信頻率等，因為訓練動力學有所改變。

誤區三：假設 ZeRO 能解決所有分布式訓練問題。ZeRO 優化的是記憶體和參數冗餘，不直接解決通信瓶頸、同步延遲等問題。需要配合其他優化（如梯度壓縮、非同步更新）。

與相關技術的比較

• ZeRO vs 模型分片（Pipeline Parallelism）：模型分片沿層次維度分割模型，各 GPU 處理不同層。ZeRO 沿參數維度分片，各 GPU 處理相同層的不同參數。兩者應用於不同的並行維度。

• ZeRO vs 張量並行：張量並行在多維上分割張量計算，ZeRO 分片存儲。兩者可結合使用實現混合並行。

• ZeRO vs 梯度累積：梯度累積延遲更新，ZeRO 通過分片降低記憶體。兩者互補，在超大模型訓練中常結合使用。

• ZeRO vs 量化：量化降低精度節省空間，ZeRO 通過分片節省空間。兩者可結合，量化進一步減少每個分片的記憶體需求。