Hugging Face Accelerate:统一 DeepSpeed 与 FSDP 的精度

Hugging Face Accelerate 已推出更新,以统一 PyTorch Fully Sharded Data Parallel(FSDP)和 Microsoft DeepSpeed 的精度处理。此更改已集成至 Accelerate 0.30.0 版本,使用户能够在这两种 ZeRO 冗余优化器实现之间切换,而不会因内部精度默认值不同而出现收敛差异。

精度差异与收敛性

DeepSpeed 与 FSDP 在参数精度处理上的差异可能导致训练结果出现分歧。在使用 bfloat16 加载的 Mistral-7B 基础模型进行测试时,DeepSpeed 展现出稳定的收敛,而 FSDP 除非手动将学习率按 GPU 数量放大或降低至 1e-5,否则无法降低损失。

这种行为归因于内部上转。DeepSpeed 的 DeepSpeedZeroOptimizer_Stage3 会通过 _create_fp32_partitions 自动将可训练参数组上转为 float32,并在设计上保持主权重的全精度。这使得优化器能够在低精度下不稳定的学习率下实现收敛。

精度工作流比较

DeepSpeed 与 FSDP 在管理“展平”参数和优化器初始化方式上根本不同:

过程 FSDP DeepSpeed
准备阶段 使用 torch_dtype 忽略 torch_dtype,在 float32 中创建
优化器初始化 在 torch_dtype 中创建参数 在 float32 中创建参数
优化器(预步) 如有上转则上转至 torch_dtype 将所有内容上转至 float32
优化器(实际步) 在 torch_dtype 中执行 在 float32 中执行

虽然 DeepSpeed 的原生上转确保了收敛,但会使内存消耗增加约 2 倍,这在使用少量 GPU 进行训练时影响显著。相反,torch 原生的 FSDP 实现并不强制上转,在内存受限的场景下提供了更大的灵活性,使优化器能够在低精度下运行。

Accelerate 0.30.0 中的新 FSDP 模式

为统一这两大框架,Hugging Face Accelerate 现在支持两种不同的 FSDP 模式,以匹配用户在收敛稳定性和内存效率之间的优先级:

  1. Mixed-Precision Mode:在准备阶段和优化器步骤中将参数上转至 fp32(与 DeepSpeed 对齐),而训练保持 bf16。
  2. Memory-Constrained Mode:在准备、训练和优化器步骤中全程使用低精度 (bf16)。

框架比较表

框架 模型加载 混合精度 准备阶段 训练 优化器
FSDP(内存受限) bf16 无 bf16 bf16 bf16
FSDP(混合精度) bf16 bf16 fp32 bf16 fp32
DeepSpeed bf16 bf16 fp32 bf16 fp32

吞吐性能

在四块 A100 GPU 上使用 IBM Granite 7B 模型进行基准测试表明,FSDP(对齐模式)和 DeepSpeed(Zero3)在吞吐性能上几乎相同:

框架 每设备每秒标记数 步长时间 (秒) 模型 FLOPs 利用率 (MFU)
FSDP(对齐) 3158.7 10.4 0.41
DeepSpeed 3094.5 10.6 0.40

迁移与配置

Hugging Face 已发布一篇概念指南帮助用户在 FSDP 与 DeepSpeed 之间迁移。切换主要通过 Accelerate 配置文件或 DeepSpeed 与 FSDP 插件类完成。迁移时的关键考虑因素包括:

  • 分片策略:在不同框架之间实现等效的分片。
  • 模型加载:采用高效的加载模式。
  • 权重预取:管理权重在 GPU 之间的移动方式。
  • 检查点:处理各框架在保存和加载模型状态时的差异。

Sources

相关

  • Dispatch
  • Dispatch
  • Dispatch
  • Dispatch
  • Dispatch