社區(qū)中有兩個流行的零冗余優(yōu)化器 (Zero Redundancy Optimizer，ZeRO)算法實現(xiàn)，一個來自DeepSpeed，另一個來自PyTorch。Hugging FaceAccelerate對這兩者都進行了集成并通過接口暴露出來，以供最終用戶在訓(xùn)練/微調(diào)模型時自主選擇其中之一。

本文重點介紹了 Accelerate 對外暴露的這兩個后端之間的差異。為了讓用戶能夠在這兩個后端之間無縫切換，我們在 Accelerate 中合并了一個精度相關(guān)的 PR及一個新的概念指南。

• 零冗余優(yōu)化器 (Zero Redundancy Optimizer，ZeRO)https://arxiv.org/abs/1910.02054

• DeepSpeedhttps://github.com/microsoft/DeepSpeed

• PyTorchhttps://pytorch.org/docs/stable/fsdp.html

• Acceleratehttps://hf.co/docs/accelerate/en/index

• 一個精度相關(guān)的 PRhttps://github.com/huggingface/accelerate/issues/2624

• 一個新的概念指南https://hf.co/docs/accelerate/concept_guides/fsdp_and_deepspeed

FSDP 和 DeepSpeed 可以互換嗎？

最近，我們嘗試分別使用 DeepSpeed 和 PyTorch FSDP 進行訓(xùn)練，發(fā)現(xiàn)兩者表現(xiàn)有所不同。我們使用的是 Mistral-7B 基礎(chǔ)模型，并以半精度 (bfloat16) 加載。可以看到 DeepSpeed (藍色) 損失函數(shù)收斂良好，但 FSDP (橙色) 損失函數(shù)沒有收斂，如圖 1 所示。

我們猜想可能需要根據(jù) GPU 數(shù)量對學(xué)習(xí)率進行縮放，且由于我們使用了 4 個 GPU，于是我們將學(xué)習(xí)率提高了 4 倍。然后，損失表現(xiàn)如圖 2 所示。

看起來，通過按 GPU 數(shù)量縮放 FSDP 學(xué)習(xí)率，已經(jīng)達到了預(yù)期！然而，當(dāng)我們在不進行縮放的情況下嘗試其他學(xué)習(xí)率 (1e-5) 時，我們卻又觀察到這兩個框架的損失和梯度范數(shù)特征又是趨近一致的，如圖 3 所示。

精度很重要

在 DeepSpeed 代碼庫的 DeepSpeedZeroOptimizer_Stage3 (顧名思義，處理第 3 階段優(yōu)化器分片) 實現(xiàn)代碼中，我們注意到 trainable_param_groups (可訓(xùn)參數(shù)組) 被傳入一個內(nèi)部函數(shù) _setup_for_real_optimizer，該函數(shù)會調(diào)用另一個名為 _create_fp32_partitions 的函數(shù)。正如其名稱中的 fp32 所示，DeepSpeed 內(nèi)部執(zhí)行了精度上轉(zhuǎn)，并在設(shè)計上始終將主權(quán)重保持為 fp32 精度。而上轉(zhuǎn)至全精度意味著：同一個學(xué)習(xí)率，上轉(zhuǎn)后的優(yōu)化器可以收斂，而原始低精度下的優(yōu)化器則可能不會收斂。前述現(xiàn)象就是這種精度差異的產(chǎn)物。

在 FSDP 中，在把模型和優(yōu)化器參數(shù)分片到各 GPU 上之前，這些參數(shù)首先會被“展平”為一維張量。FSDP 和 DeepSpeed 對這些“展平”參數(shù)使用了不同的 dtype，這會影響 PyTorch 優(yōu)化器的表現(xiàn)。表 1 概述了兩個框架各自的處理流程，“本地？”列說明了當(dāng)前步驟是否是由各 GPU 本地執(zhí)行的，如果是這樣的話，那么上轉(zhuǎn)的內(nèi)存開銷就可以分?jǐn)偟礁鱾€ GPU。

流程	本地？	框架	詳情
模型加載 (如 AutoModel.from_pretrained(..., torch_dtype=torch_dtype))	?
準(zhǔn)備，如創(chuàng)建“展平參數(shù)”	?	FSDP DeepSpeed	使用 torch_dtype 不管 torch_dtype，直接創(chuàng)建為 float32
優(yōu)化器初始化	?	FSDP DeepSpeed	用 torch_dtype 創(chuàng)建參數(shù) 用 float32 創(chuàng)建參數(shù)
訓(xùn)練步 (前向、后向、歸約)	?	FSDP DeepSpeed	遵循fsdp.MixedPrecision 遵循 deepspeed_config_file 中的混合精度設(shè)置
優(yōu)化器 (準(zhǔn)備階段)	?	FSDP DeepSpeed	按需上轉(zhuǎn)至 torch_dtype 所有均上轉(zhuǎn)至 float32
優(yōu)化器 (實際執(zhí)行階段)	?	FSDP DeepSpeed	以 torch_dtype 精度進行 以 float32 精度進行

表 1：FSDP 與 DeepSpeed 混合精度處理異同

• fsdp.MixedPrecisionhttps://pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.MixedPrecision

幾個要點：

• 正如 🤗 Accelerate 上的這一問題所述，混合精度訓(xùn)練的經(jīng)驗法則是將可訓(xùn)參數(shù)精度保持為 float32。

• 當(dāng)在大量 GPU 上進行分片時，上轉(zhuǎn) (如 DeepSpeed 中所做的那樣) 對內(nèi)存消耗的影響可能可以忽略不計。然而，當(dāng)在少量 GPU 上使用 DeepSpeed 時，內(nèi)存消耗會顯著增加，高達 2 倍。

• FSDP 的 PyTorch 原生實現(xiàn)不會強制上轉(zhuǎn)，其支持用戶以低精度操作 PyTorch 優(yōu)化器，因此相比 DeepSpeed 提供了更大的靈活性。

• 這一問題https://github.com/huggingface/accelerate/issues/2624#issuecomment-2058402753

在 🤗 Accelerate 中對齊 DeepSpeed 和 FSDP 的行為

為了在🤗 Accelerate 中更好地對齊 DeepSpeed 和 FSDP 的行為，我們可以在啟用混合精度時自動對 FSDP 執(zhí)行上轉(zhuǎn)。我們?yōu)榇俗隽艘粋€ PR，該 PR 現(xiàn)已包含在0.30.0 版本中了。

• 0.30.0 版本https://github.com/huggingface/accelerate/releases/tag/v0.30.0

有了這個 PR，FSDP 就能以兩種模式運行：

• 與 DeepSpeed 一致的混合精度模式

• 針對內(nèi)存受限場景的低精度模式，如圖 4 所示。

表 2 總結(jié)了兩種新的 FSDP 模式，并與 DeepSpeed 進行了比較。

框架	模型加載 (torch_dtype)	混合精度	準(zhǔn)備 (本地)	訓(xùn)練	優(yōu)化器 (本地)
FSDP (低精度模式)	bf16	缺省 (無)	bf16	bf16	bf16
FSDP (混合精度模式)	bf16	bf16	fp32	bf16	fp32
DeepSpeed	bf16	bf16	fp32	bf16	fp32

表 2：兩種新 FSDP 模式總結(jié)及與 DeepSpeed 的對比

吞吐量測試結(jié)果

我們使用IBM Granite 7B模型 (其架構(gòu)為 Meta Llama2) 進行吞吐量比較。我們比較了模型的浮點算力利用率 (Model Flops Utilization，MFU) 和每 GPU 每秒詞元數(shù)這兩個指標(biāo)，并針對 FSDP (完全分片) 和 DeepSpeed (ZeRO3) 兩個場景進行了測量。

• IBM Granite 7Bhttps://hf.co/ibm-granite/granite-7b-base

如上文，我們使用 4 張 A100 GPU，超參如下：

• batch size 為 8

• 模型加載為 torch.bfloat16

• 使用 torch.bfloat16 混合精度

表 3 表明 FSDP 和 DeepSpeed 的表現(xiàn)類似，這與我們的預(yù)期相符。

隨著大規(guī)模對齊技術(shù) (如InstructLab及GLAN) 的流行，我們計劃對結(jié)合各種提高吞吐量的方法 (如，序列組裝 + 4D 掩碼、torch.compile、選擇性 checkpointing) 進行全面的吞吐量對比基準(zhǔn)測試。

• InstructLabhttps://github.com/instructlab

• GLANhttps://arxiv.org/abs/2402.13064

框架	每 GPU 每秒詞元數(shù)	**每步耗時 (s) **	**浮點算力利用率 (MFU) **
FSDP (混合精度模式)	3158.7	10.4	0.41
DeepSpeed	3094.5	10.6	0.40

表 3：四張 A100 GPU 上 FSDP 和 DeepSpeed 之間的大致吞吐量比較。

最后的話

我們提供了新的概念指南以幫助用戶在兩個框架之間遷移。該指南可以幫助用戶厘清以下問題：

• 如何實現(xiàn)等效的分片策略？

• 如何進行高效的模型加載？

• FSDP 和 DeepSpeed 中如何管理權(quán)重預(yù)取？

• 與 DeepSpeed 對等的 FSDP 封裝是什么？

我們在 🤗 Accelerate 中考慮了配置這些框架的各種方式:

• 使用 accelerate launch 從命令行配置

• 從🤗 Accelerate 提供給DeepSpeedhttps://hf.co/docs/accelerate/main/en/package_reference/deepspeed和FSDPhttps://hf.co/docs/accelerate/main/en/package_reference/fsdp的各種 Plugin 類中配置

🤗 Accelerate 使得在 FSDP 和 DeepSpeed 之間切換非常絲滑，大部分工作都只涉及更改 Accelerate 配置文件 (有關(guān)這方面的說明，請參閱新的概念指南) 。

除了配置變更之外，還有一些如檢查點處理方式的差異等，我們一并在指南中進行了說明。

本文中的所有實驗都可以使用原始 🤗 Accelerate 問題中的代碼重現(xiàn)。

• 概念指南https://hf.co/docs/accelerate/v0.31.0/en/concept_guides/fsdp_and_deepspeed

• 原始 🤗 Accelerate 問題https://github.com/huggingface/accelerate/issues/2624

我們計劃后續(xù)在更大規(guī)模 GPU 上進行吞吐量比較，并對各種不同技術(shù)進行比較，以在保持模型質(zhì)量的前提下更好地利用更多的 GPU 進行微調(diào)和對齊。

致謝

本工作凝聚了來自多個組織的多個團隊的共同努力。始于 IBM 研究中心，特別是發(fā)現(xiàn)該問題的 Aldo Pareja 和發(fā)現(xiàn)精度差距并解決該問題的 Fabian Lim。Zach Mueller 和Stas Bekman在提供反饋和修復(fù) accelerate 的問題上表現(xiàn)出色。Meta PyTorch 團隊的 Less Wright 對有關(guān) FSDP 參數(shù)的問題非常有幫助。最后，我們還要感謝?DeepSpeed?團隊對本文提供的反饋。

• Stas Bekmanhttps://github.com/stas00

• DeepSpeedhttps://www.deepspeed.ai/

---

英文原文: https://hf.co/blog/deepspeed-to-fsdp-and-back

原文作者: Yu Chin Fabian, aldo pareja, Zachary Mueller, Stas Bekman

譯者: Matrix Yao (姚偉峰)，英特爾深度學(xué)習(xí)工程師，工作方向為 transformer-family 模型在各模態(tài)數(shù)據(jù)上的應(yīng)用及大規(guī)模模型的訓(xùn)練推理。