LLMs之Framework：Hugging Face Accelerate后端框架之FSDP和DeepSpeed的对比与分析

地址

论文地址：https://huggingface.co/blog/deepspeed-to-fsdp-and-back

时间

2024年6月13日

作者

Yu Chin Fabian、aldo pareja、Zachary Mueller、Stas Bekman

总结

总的来说，该文章阐述了FSDP和DeepSpeed在实现上的差异，Accelerate如何统一它们的行为，并提供指导帮助用户在两种后端之间切换。同时也讨论了低精度优化的内存与收敛权衡。

框架差异

DeepSpeed 和 FSDP 是两种流行的优化器分片(ZeRO算法)实现，各自有不同的精度处理方式。DeepSpeed 默认在 fp32 精度下操作，FSDP 更灵活，可以选择更低的精度。

DeepSpeed 和 PyTorch 中的 FSDP 是社区中两种流行的 ZeRO Redundancy Optimizer (Zero) 算法实现，用于分布式训练大型模型。Hugging Face Accelerate 允同时支持这两个框架之间无缝切换。

是否可互换

DeepSpeed 和 FSDP 在损失收敛上有明显差异，造成差异的原因在于精度处理不同。可以通过调整学习率可改善 FSDP 表现。

在实验中，使用 Mistral-7B 模型进行训练时，DeepSpeed 的损失收敛良好，而 FSDP 的损失没有减少，收敛效果较差。通过将学习率根据 GPU 数量进行调整，FSDP 的表现有所改善。

精度问题

DeepSpeed 使用 fp32，FSDP 允许低精度操作。

DeepSpeed 内部使用全精度（fp32）进行优化，这可能导致在低精度下不收敛的问题。FSDP 不强制使用全精度，允许在低精度(如BF16)下操作，从而提供了更大的灵活性。

DeepSpeed 内部总是将参数提升到 fp32 进行计算，这在小规模 GPU 环境中可能增加内存消耗。FSDP 则允许使用低精度操作，更适合内存受限的场景，但收敛可能较慢。DeepSpeed 的upcasting方式则收敛更快，但内存开销较大。

协调改进

FSDP 在 Accelerate 中增加了自动提升精度功能，支持混合精度和低精度模式。

为了更好协调 DeepSpeed 和 FSDP ，Hugging Face Accelerate 为 FSDP 增加了自动提升精度的功能。FSDP 现在可以在“混合精度”和“低精度”两种模式下操作。

为统一两种后端，Hugging Face Accelerate为FSDP 引入了"混合精度"模式，与DeepSpeed 行为一致，在FP32上进行优化。

因此Accelerate下FSDP 有两种模式：

>> 默认低精度模式，参数保持载入精度(如BF16)，内存占用少；

>> 混合精度模式，参数在FP32上优化，行为与DeepSpeed 一致；

性能对比

FSDP 和 DeepSpeed 在 IBM Granite 7B 上表现相似，处理速度和效率接近。

使用 IBM Granite 7B 模型进行吞吐量比较，FSDP 和 DeepSpeed 的表现相似。

在四个 A100 GPU 上进行测试，FSDP 和 DeepSpeed 的每秒处理的 token 数量和模型 FLOP 利用率相近。

结论与思考

提供迁移指南，通过配置文件可在两个框架间切换。

Hugging Face Accelerate提供了帮助用户在两个框架之间迁移的概念指南。通过修改 Accelerate 配置文件，可以轻松在 FSDP 和 DeepSpeed 之间切换。

Accelerate提供了新的概念指南，帮助用户在两种后端之间迁移，包括分片策略、模型加载、权重预取等。通过Accelerate的命令行和插件类，可以轻松切换FSDP 和DeepSpeed 后端。