Hugging Face Accelerate 两个后端的故事：FSDP 与 DeepSpeed

社区中有两个流行的零冗余优化器 (Zero Redundancy Optimizer，ZeRO)算法实现，一个来自DeepSpeed，另一个来自PyTorch。Hugging FaceAccelerate对这两者都进行了集成并通过接口暴露出来，以供最终用户在训练/微调模型时自主选择其中之一。

本文重点介绍了 Accelerate 对外暴露的这两个后端之间的差异。为了让用户能够在这两个后端之间无缝切换，我们在 Accelerate 中合并了一个精度相关的 PR及一个新的概念指南。

• 零冗余优化器 (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

• 一个精度相关的 PRhttps://github.com/huggingface/accelerate/issues/2624

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

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。

表 1：FSDP 与 DeepSpeed 混合精度处理异同

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

几个要点：

• 正如

  声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/一键难忘520/article/detail/883937