Hugging Face Accelerate 的多後端故事：FSDP 與 DeepSpeed
社群中有兩種流行的 ZeRO 冗餘最佳化器 (Zero) 演算法實現，一種來自 DeepSpeed，另一種來自 PyTorch。Hugging Face Accelerate 為終端使用者提供了這兩種框架，用於訓練/微調他們的模型。本部落格重點介紹了透過 Accelerate 使用這兩個後端時的差異。為了讓使用者能夠無縫地在這些後端之間切換，我們向上遊提交了一項與精度相關的更改和一份概念指南。
FSDP 和 DeepSpeed 可以互換嗎？
最近，我們嘗試使用 DeepSpeed 和 PyTorch FSDP 執行一個訓練流水線。我們注意到得到的結果有所不同。具體的模型是 Mistral-7B 基礎模型，並以半精度 (bfloat16) 載入。雖然 DeepSpeed (藍色) 的損失已經很好地收斂，但 FSDP (橙色) 的損失卻沒有下降，如圖 1 所示。
我們假設學習率可能需要按 GPU 數量進行縮放，由於我們使用了 4 個 GPU，於是我們將學習率提高了 4 倍。然後，我們看到了如圖 2 所示的損失行為。
看起來，透過將 FSDP 的學習率乘以 GPU 數量，我們已經達到了預期的效果！然而，當我們嘗試另一個學習率 (1e-5) 而不進行縮放時，我們觀察到兩個框架的損失和梯度範數特徵相似，如圖 3 所示。
精度至關重要
在 DeepSpeed 的程式碼庫中，特別是在 DeepSpeedZeroOptimizer_Stage3 (顧名思義，它負責處理 Stage 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 issue 中提到的，進行混合精度訓練時的一個經驗法則是將可訓練引數保持在
float32。 - 當在大量 GPU 上進行分片時，像
DeepSpeed中那樣的上轉型對記憶體消耗的影響可能微不足道。然而，當在少量 GPU 上使用DeepSpeed時，2 倍的記憶體消耗增長可能非常顯著。 - PyTorch 原生的 FSDP 實現不強制上轉型，允許使用者在低精度下執行 PyTorch 最佳化器。這比
DeepSpeed的原生上轉型提供了更大的靈活性。
在 🤗 Accelerate 中統一 DeepSpeed 和 FSDP
為了在 🤗 Accelerate 中更好地統一 DeepSpeed 和 FSDP，我們可以在啟用混合精度時為 FSDP 自動執行上轉型。我們建立了一個包含此更改的拉取請求，並已包含在 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 架構) 進行吞吐量比較。我們比較了模型浮點運算利用率 (MFU) 和 每秒/每 GPU 的 token 數指標，並展示了 FSDP (完全分片) 和 DeepSpeed (Zero3) 的結果。
我們像之前一樣使用了四塊 A100 GPU，並使用了以下超引數
- 批次大小為 8
- 模型以
torch.bfloat16載入 - 混合精度使用相同的資料型別。
表 3 顯示，FSDP 和 DeepSpeed 的效能預計會相似。
我們計劃後續進行全面的吞吐量比較，並探討提高吞吐量的方法 (例如，使用 packing 的 4D 掩碼、torch.compile、選擇性啟用檢查點)，因為像 InstructLab 和 GLAN 這樣的大規模對齊技術正變得越來越流行。
| 框架 | 每裝置每秒處理的 token 數 | 步長時間 (秒) | 模型浮點運算利用率 (MFU) |
|---|---|---|---|
| FSDP (對齊模式) | 3158.7 | 10.4 | 0.41 |
| DeepSpeed | 3094.5 | 10.6 | 0.40 |
表 3：在四塊 A100 GPU 上 FSDP 和 DeepSpeed 的大致吞吐量比較。
總結
我們提供了一份新的概念指南，以幫助使用者在這兩個框架之間遷移。該指南幫助使用者回答以下問題：
- 我們如何實現等效的分片策略？
- 我們如何執行高效的模型載入？
- 在 FSDP 和 DeepSpeed 中如何管理權重預取？
- 在 DeepSpeed 中與 FSDP 的 wrapping 等效的操作是什麼？
我們考慮了在 🤗 Accelerate 中配置這些框架的各種模式，
- 在
accelerate launch期間從命令列 - 透過 🤗 Accelerate 為 (
DeepSpeed)[https://huggingface.co/docs/accelerate/main/en/package_reference/deepspeed] 和 (FSDP)[https://huggingface.co/docs/accelerate/main/en/package_reference/fsdp] 提供的各種Plugin類
🤗 Accelerate 使得在 FSDP 和 DeepSpeed 之間切換變得幾乎**輕而易舉**，大部分工作只是更改 Accelerate 配置檔案 (請參閱新的概念指南以獲取相關說明)。
除了配置更改之外，其他一些需要考慮的因素 (指南中也已概述) 是檢查點處理方式的差異等。
本部落格中的所有實驗都可以使用原始 🤗 Accelerate issue 中的程式碼進行復現。
我們計劃後續進行大規模的吞吐量比較，並探討如何更好地利用這些 GPU 進行微調和對齊任務，同時保持模型質量。
致謝
這項工作是多個組織中多個團隊共同努力的結果。它始於 IBM Research，特別是 Aldo Pareja 發現了這個問題，以及 Fabian Lim 識別了精度差距並解決了這個問題。Zach Mueller 和 Stas Bekman 在提供反饋和對 accelerate 的修復方面給予了巨大的幫助。來自 Meta PyTorch 團隊的 Less Wright 在 FSDP 引數問題上提供了非常有益的幫助。最後，我們還要感謝 DeepSpeed 團隊對本部落格提供的反饋。
