使用 DeepSpeed 加速大型模型訓練

Hugging Face 已將 DeepSpeed 的 Zero Redundancy Optimizer（ZeRO）整合到 accelerate 函式庫中，讓使用者透過減少 GPU 之間的記憶體冗餘，以更有效率地訓練大規模模型。此整合使開發者能避免記憶體不足（OOM）錯誤，並在不需要大量程式碼變更的情況下提升批次大小。

了解 ZeRO 資料平行化

ZeRO（Zero Redundancy Optimizer）透過將訓練狀態切分至可用的 GPU 來最佳化資料平行化時的記憶體使用。此框架分為多個階段，以減少模型訓練的記憶體占用：

  • Stage 1：將 optimizer 狀態切分至資料平行工作者 / GPU 上。
  • Stage 2：同時將 optimizer 狀態與梯度切分至工作者 / GPU 上。
  • Stage 3：將 optimizer 狀態、梯度與模型參數皆切分至工作者 / GPU 上。
  • Optimizer Offload：在 Stage 2 基礎上，將梯度與 optimizer 狀態卸載至 CPU 或磁碟。
  • Param Offload：在 Stage 3 基礎上，將模型參數卸載至 CPU 或磁碟。

零程式碼整合（透過 Accelerate）

使用者只需透過 accelerate config 指令與 Accelerate DeepSpeed Plugin，即可在不修改訓練程式碼的情況下使用 DeepSpeed ZeRO Stage‑2。

在使用兩張 24GB NVIDIA Titan RTX GPU 微調 DeBERTa‑v2‑xlarge‑mnli（900M 參數）時，DeepSpeed ZeRO Stage‑2 明顯優於 Distributed Data Parallel（DDP）：

方法 最大批次大小 每 epoch 訓練時間 (秒) 每 epoch 評估時間 (秒) F1 分數 正確率
DDP 8 103.57 2.04 0.931 0.904
DeepSpeed ZeRO Stage 2 40 28.98 1.79 0.936 0.912

DeepSpeed 使最大批次大小提升 5 倍，總訓練時間加速約 3.5 倍，且性能指標未出現退化。

使用 DeepSpeed 設定檔的進階配置

若需更細緻的控制，使用者可以透過 accelerate config 提供 DeepSpeed 設定 JSON 檔。當設定檔中定義了 optimizer 與 scheduler 時，只需將標準的 PyTorch optimizer 與 scheduler 替換為 accelerate.utils.DummyOptim 與 accelerate.utils.DummyScheduler，即可完成最小程式碼調整。

在 MuDoConv 資料集上測試 BlenderBot‑400M‑distill 模型時，DeepSpeed ZeRO Stage‑2 允許的最大批次大小為 200，而 DDP 僅為 100。此結果帶來 1.44 倍 的訓練加速與 1.23 倍 的評估加速。

使用 CPU/磁碟 卸載訓練巨型模型

DeepSpeed ZeRO Stage‑3 搭配 CPU 卸載，使得即使批次大小為 1 仍無法容納於 GPU 記憶體的模型得以訓練。

在單張 24GB NVIDIA Titan RTX GPU 上訓練 GPT‑XL（1.5B 參數）時，DDP 立即觸發 OOM 錯誤；相較之下，DeepSpeed ZeRO Stage‑3 透過將 optimizer 狀態、梯度與參數卸載至 CPU，成功以 批次大小 16 完成訓練，一個 epoch 耗時 6608.35 秒。

方法 最大批次大小 每 epoch 訓練時間 (秒) 備註
DDP - - OOM Error
DeepSpeed ZeRO Stage 3 16 6608.35 -

Sources