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 团队对本博客提供的反馈。
