社群中有兩個流行的 零冗餘最佳化器（Zero Redundancy Optimizer，ZeRO） 演算法實現，一個來自 DeepSpeed，另一個來自 PyTorch。Hugging Face Accelerate 對這兩者都進行了整合並透過介面暴露出來，以供終端使用者在訓練/微調模型時自主選擇其中之一。本文重點介紹了 Accelerate 對外暴露的這兩個後端之間的差異。為了讓使用者能夠在這兩個後端之間無縫切換，我們在 Accelerate 中合併了 一個精度相關的 PR 及 一個新的概念指南。
FSDP 和 DeepSpeed 可以互換嗎？
最近，我們嘗試分別使用 DeepSpeed 和 PyTorch FSDP 進行訓練，發現兩者表現有所不同。我們使用的是 Mistral-7B 基礎模型，並以半精度（bfloat16）載入。可以看到 DeepSpeed（藍色）損失函式收斂良好，但 FSDP（橙色）損失函式沒有收斂，如圖 1 所示。
我們猜想可能需要根據 GPU 數量對學習率進行縮放，且由於我們使用了 4 個 GPU，於是我們將學習率提高了 4 倍。然後，損失表現如圖 2 所示。
看起來，透過按 GPU 數量縮放 FSDP 學習率，已經達到了預期！然而，當我們在不進行縮放的情況下嘗試其他學習率（1e-5）時，我們卻又觀察到這兩個框架的損失和梯度範數特徵又是趨近一致的，如圖 3 所示。
精度很重要
在 DeepSpeed 程式碼庫的 DeepSpeedZeroOptimizer_Stage3（顧名思義，處理第 3 階段最佳化器分片）實現程式碼中，我們注意到 trainable_param_groups（可訓引數組）被傳入一個內部函式 _setup_for_real_optimizer，該函式會呼叫另一個名為 _create_fp32_partitions 的函式。正如其名稱中的 fp32 所示，DeepSpeed 內部執行了精度上轉，並在設計上始終將主權重保持為 fp32 精度。而上轉至全精度意味著：同一個學習率，上轉後的最佳化器可以收斂，而原始低精度下的最佳化器則可能不會收斂。前述現象就是這種精度差異的產物。
在 FSDP 中，在把模型和最佳化器引數分片到各 GPU 上之前，這些引數首先會被“展平”為一維張量。FSDP 和 DeepSpeed 對這些“展平”引數使用了不同的 dtype，這會影響 PyTorch 最佳化器的表現。表 1 概述了兩個框架各自的處理流程，“本地？”列說明了當前步驟是否是由各 GPU 本地執行的，如果是這樣的話，那麼上轉的記憶體開銷就可以分攤到各個 GPU。
| 流程 | 本地？ | 框架 | 詳情 |
|---|---|---|---|
模型載入（如 AutoModel.from_pretrained(..., torch_dtype=torch_dtype)） |
❌ | ||
| 準備，如建立“展平引數” | ✅ | FSDP DeepSpeed |
使用 torch_dtype不管 torch_dtype，直接建立為 float32 |
| 最佳化器初始化 | ✅ | FSDP DeepSpeed |
用 torch_dtype 建立引數用 float32 建立引數 |
| 訓練步（前向、後向、歸約） | ❌ | FSDP DeepSpeed |
遵循 fsdp.MixedPrecision 遵循 deepspeed_config_file 中的混合精度設定 |
| 最佳化器（準備階段） | ✅ | FSDP DeepSpeed |
按需上轉至 torch_dtype所有均上轉至 float32 |
| 最佳化器（實際執行階段） | ✅ | FSDP DeepSpeed |
以 torch_dtype 精度進行以 float32 精度進行 |
表 1： FSDP 與 DeepSpeed 混合精度處理異同
幾個要點：
- 正如 🤗 Accelerate 上的這一問題所述，混合精度訓練的經驗法則是將可訓引數精度保持為
float32。 - 當在大量 GPU 上進行分片時，上轉（如
DeepSpeed中所做的那樣）對記憶體消耗的影響可能可以忽略不計。然而，當在少量 GPU 上使用DeepSpeed時，記憶體消耗會顯著增加，高達 2 倍。 - FSDP 的 PyTorch 原生實現不會強制上轉，其支援使用者以低精度操作 PyTorch 最佳化器，因此相比
DeepSpeed提供了更大的靈活性。
在 🤗 Accelerate 中對齊 DeepSpeed 和 FSDP 的行為
為了在🤗 Accelerate 中更好地對齊 DeepSpeed 和 FSDP 的行為，我們可以在啟用混合精度時自動對 FSDP 執行上轉。我們為此做了一個 PR，該 PR 現已包含在 0.30.0 版本中了。
有了這個 PR，FSDP 就能以兩種模式執行：
- 與 DeepSpeed 一致的
混合精度模式 - 針對記憶體受限場景的低精度模式，如圖 4 所示。
表 2 總結了兩種新的 FSDP 模式，並與 DeepSpeed 進行了比較。
| 框架 | 模型載入 (torch_dtype) |
混合精度 | 準備（本地） | 訓練 | 最佳化器（本地） |
|---|---|---|---|---|---|
| FSDP（低精度模式） | bf16 |
預設（無） | bf16 |
bf16 |
bf16 |
| FSDP（混合精度模式） | bf16 |
bf16 |
fp32 |
bf16 |
fp32 |
| DeepSpeed | bf16 |
bf16 |
fp32 |
bf16 |
fp32 |
表 2：兩種新 FSDP 模式總結及與 DeepSpeed 的對比
吞吐量測試結果
我們使用 IBM Granite 7B 模型（其架構為 Meta Llama2）進行吞吐量比較。我們比較了模型的浮點算力利用率 (Model Flops Utilization，MFU) 和每 GPU 每秒詞元數這兩個指標，並針對 FSDP（完全分片）和 DeepSpeed（ZeRO3）兩個場景進行了測量。
如上文，我們使用 4 張 A100 GPU，超參如下：
- batch size 為 8
- 模型載入為
torch.bfloat16 - 使用
torch.bfloat16混合精度
表 3 表明 FSDP 和 DeepSpeed 的表現類似，這與我們的預期相符。
隨著大規模對齊技術（如 InstructLab 及 GLAN）的流行，我們計劃對結合各種提高吞吐量的方法（如，序列組裝 + 4D 掩碼、torch.compile、選擇性 checkpointing）進行全面的吞吐量對比基準測試。
| 框架 | 每 GPU 每秒詞元數 | 每步耗時（s） | 浮點算力利用率（MFU） |
|---|---|---|---|
| FSDP（混合精度模式） | 3158.7 | 10.4 | 0.41 |
| DeepSpeed | 3094.5 | 10.6 | 0.40 |
表 3：四張 A100 GPU 上 FSDP 和 DeepSpeed 之間的大致吞吐量比較。
最後的話
我們提供了新的 概念指南 以幫助使用者在兩個框架之間遷移。該指南可以幫助使用者釐清以下問題：
- 如何實現等效的分片策略？
- 如何進行高效的模型載入？
- FSDP 和 DeepSpeed 中如何管理權重預取？
- 與 DeepSpeed 對等的 FSDP 封裝是什麼？
我們在 🤗 Accelerate 中考慮了配置這些框架的各種方式:
- 使用
accelerate launch從命令列配置 - 從🤗 Accelerate 提供給
DeepSpeed和FSDP的各種Plugin類中配置
🤗 Accelerate 使得在 FSDP 和 DeepSpeed 之間切換非常絲滑，大部分工作都只涉及更改 Accelerate 配置檔案（有關這方面的說明，請參閱新的概念指南）。
除了配置變更之外，還有一些如檢查點處理方式的差異等，我們一併在指南中進行了說明。
本文中的所有實驗都可以使用 原始 🤗 Accelerate 問題 中的程式碼重現。
我們計劃後續在更大規模 GPU 上進行吞吐量比較，並對各種不同技術進行比較，以在保持模型質量的前提下更好地利用更多的 GPU 進行微調和對齊。
致謝
本工作凝聚了來自多個組織的多個團隊的共同努力。始於 IBM 研究中心，特別是發現該問題的 Aldo Pareja 和發現精度差距並解決該問題的 Fabian Lim。Zach Mueller 和 Stas Bekman 在提供反饋和修復 accelerate 的問題上表現出色。Meta PyTorch 團隊的 Less Wright 對有關 FSDP 引數的問題非常有幫助。最後，我們還要感謝 DeepSpeed 團隊對本文提供的反饋。
英文原文: https://hf.co/blog/deepspeed-to-fsdp-and-back
原文作者: Yu Chin Fabian, aldo pareja, Zachary Mueller, Stas Bekman
譯者: Matrix Yao (姚偉峰)，英特爾深度學習工程師，工作方向為 transformer-family 模型在各模態資料上的應用及大規模模型的訓練推理。