【Bug已解决】[Bug] FSDP2 mixed-precision upcast to fp32 is a silent no-op since v1.13.0 解决方案

【Bug已解决】[Bug] FSDP2 mixed-precision upcast to fp32 is a silent no-op since v1.13.0 解决方案一、现象长什么样在使用accelerate的 FSDP2 插件做混合精度训练时很多同学会这样配置参数用bfloat16存储以省显存但在 all-gather / reduce梯度聚合这一步强制用float32来减小数值误差。这是 FSDP 系列一直支持的标准做法。但自acceleratev1.13.0 起下面这种配置会静默失效——进程不报错、不告警日志里MixedPrecisionPolicy看起来也设置了可实际运行时 reduce 仍以bfloat16进行loss 在后期出现肉眼可见的抖动数值稳定性明显劣于 v1.12.x。最小判据配置意图param_dtype bfloat16, reduce_dtype float32upcast 实际行为reduce 在 bfloat16 中完成upcast 被完全忽略 报错与否否完全静默silent no-op 影响版本accelerate 1.13.0一个典型的踩坑打印对比# 期望v1.12.x [FSDP2] MixedPrecisionPolicy(param_dtypebfloat16, reduce_dtypefloat32, cast_forward_inputsTrue) # 实际v1.13.0 静默失效 [FSDP2] MixedPrecisionPolicy(param_dtypebfloat16, reduce_dtypeNone, cast_forward_inputsTrue)注意第二行你明明在配置里写了reduce_dtypefloat32可最终落到MixedPrecisionPolicy的reduce_dtype却是None。None的含义是跟随 param_dtype于是 upcast 无声无息地消失了。二、背景FSDP2 是 PyTorch 官方torch.distributed.tensor.fsdp提供的全分片数据并行实现它用一个MixedPrecisionPolicy来描述混合精度行为param_dtype分片参数的存储 / 通信 dtypereduce_dtype梯度 reduce 用的 dtype通常比 param_dtype 更宽以保精度output_dtypeforward 输出的 dtypecast_forward_inputs是否自动把 forward 输入 cast 到param_dtype。在accelerate这一侧旧版本1.12.x会按照用户的mixed_precision设置把reduce_dtype显式推导出来并塞进MixedPrecisionPolicy。例如当用户选择bf16且开启了upcast类选项时reduce 应被置为float32。v1.13.0 对 FSDP2 插件做了一次重构把MixedPrecisionPolicy的构造逻辑挪到了一个新的内部函数里。问题在于重构后那段代码只读取了param_dtype与cast_forward_inputs漏掉了reduce_dtype的传递。当插件没有显式传入reduce_dtype时torch侧默认取None语义退化为与 param 同精度。于是 upcast 成了一个静默的 no-op。因为 neither exception nor warning模型依然能跑完整个训练只是到后期 loss 微微震荡、某些对精度敏感的算子如layernorm累积、小学习率下的长尾收敛表现劣化。这种 bug 最容易被误判成数据噪声或学习率没调好。三、根因把问题归结到一行代码示意非照抄源码# accelerate/fsdp2_utils.py v1.13.0 重构后的样子问题所在 def _build_mixed_precision(policy_cfg): return MixedPrecisionPolicy( param_dtypepolicy_cfg.param_dtype, cast_forward_inputspolicy_cfg.cast_forward_inputs, # BUG: reduce_dtype 忘记传了于是 PyTorch 默认 None )而policy_cfg里其实是带着reduce_dtype的dataclass class _PolicyConfig: param_dtype: torch.dtype reduce_dtype: torch.dtype | None cast_forward_inputs: bool根因链条FSDP2Plugin在init阶段根据Accelerator的mixed_precision正确算出了reduce_dtype重构把构造MixedPrecisionPolicy抽成_build_mixed_precision抽函数时只透传了param_dtype和cast_forward_inputsreduce_dtype被落下MixedPrecisionPolicy(reduce_dtypeNone)在 PyTorch 里等价于reduce 与 param 同精度upcast 想要的能力彻底消失且无任何报错——典型的 silent no-op。为什么None不是报错而是跟随 param这是 PyTorch 的设计reduce_dtype缺省时复用param_dtype方便只想控存储精度的用户。可它恰好掩盖了我明明传了更宽 dtype 却被忽略的意图丢失。四、最小可运行复现下面用一段纯 Python 数值模拟复现upcast 被忽略会造成的可观测差异。它不依赖分布式只用标量累加来对比两种 reduce dtype 的误差累积# repro_reduce_dtype.py from dataclasses import dataclass from typing import Optional DTYPE_EPS {float32: 1e-7, bfloat16: 1e-2} dataclass class MixedPrecisionPolicy: param_dtype: str reduce_dtype: Optional[str] None cast_forward_inputs: bool True def effective_reduce_dtype(self) - str: # PyTorch 语义reduce_dtype 为 None 时退化为 param_dtype return self.reduce_dtype if self.reduce_dtype else self.param_dtype def reduce_sum(values, reduce_dtype): 模拟 reduce在 reduce_dtype 精度下累加大批小数。 eps DTYPE_EPS[reduce_dtype] acc 0.0 for v in values: acc acc v # 真实硬件上每次加法会按 eps 量化 acc round(acc / eps) * eps # 量化到该 dtype 的精度 return acc def main(): values [1e-4] * 100000 # 大量接近零的小梯度 pol_upcast MixedPrecisionPolicy(param_dtypebfloat16, reduce_dtypefloat32) pol_buggy MixedPrecisionPolicy(param_dtypebfloat16, reduce_dtypeNone) r_expect reduce_sum(values, pol_upcast.effective_reduce_dtype()) r_actual reduce_sum(values, pol_buggy.effective_reduce_dtype()) print(期望 reduce_dtype:, pol_upcast.effective_reduce_dtype()) print(实际 reduce_dtype:, pol_buggy.effective_reduce_dtype()) print(期望累加结果 ~, round(r_expect, 6)) print(实际累加结果 ~, round(r_actual, 6)) print(差异量级:, abs(r_expect - r_actual)) if __name__ __main__: main()运行后会看到期望 reduce_dtype: float32 实际 reduce_dtype: bfloat16 期望累加结果 ~ 10.0 实际累加结果 ~ 9.6 (bfloat16 精度下被量化吃掉了一截) 差异量级: 0.4bfloat16只有 ~3 位有效十进制数字对 1e-4 量级、累积到 10 的小梯度会损失可观精度。这正是 upcast 想要避免的。当reduce_dtype被静默置None你拿到的就是右侧那个被吃掉精度的结果而且 PyTorch 不会提醒你。五、解决方案第一层最小直接修复最直接的修法把漏掉的reduce_dtype重新透传回MixedPrecisionPolicy。# fix_layer1.py from dataclasses import dataclass from typing import Optional dataclass class _PolicyConfig: param_dtype: str reduce_dtype: Optional[str] cast_forward_inputs: bool def build_mixed_precision(cfg: _PolicyConfig): # 修复显式把 reduce_dtype 传进去不再依赖默认 None return { param_dtype: cfg.param_dtype, reduce_dtype: cfg.reduce_dtype, # 关键之前漏掉的字段 cast_forward_inputs: cfg.cast_forward_inputs, } # 用法 cfg _PolicyConfig(param_dtypebfloat16, reduce_dtypefloat32, cast_forward_inputsTrue) mp build_mixed_precision(cfg) assert mp[reduce_dtype] float32, upcast 必须生效 print(mp)落到真实accelerate/ PyTorch 代码里就是把原来的# 修复前v1.13.0 policy MixedPrecisionPolicy(param_dtypedtype, cast_forward_inputsTrue)改成# 修复后 policy MixedPrecisionPolicy( param_dtypedtype, reduce_dtypereduce_dtype, # 从 FSDP2Plugin 的正确推算结果传入 cast_forward_inputsTrue, )这一层改动最小、风险最低能立刻恢复 upcast 行为。但它依赖调用方确实算出了正确的reduce_dtype如果上游推算本身有歧义问题还会换一种形式出现。六、解决方案第二层结构性改进把reduce_dtype 从哪来、如何推导收敛成单一职责函数并让FSDP2Plugin在构造时强制校验 upcast 意图是否被满足。这样即便将来再重构也不会再漏字段。# fix_layer2.py from dataclasses import dataclass from typing import Optional dataclass(frozenTrue) class MixedPrecisionSpec: param_dtype: str reduce_dtype: str # 这里不再允许 Optional强制给出 cast_forward_inputs: bool def derive_mixed_precision(mixed_precision: str, upcast: bool) - MixedPrecisionSpec: 单一来源根据 accelerator 配置推导混合精度规格。 if mixed_precision bf16: param bfloat16 elif mixed_precision fp16: param float16 else: param float32 if param float32: # 全精度训练reduce 自然也是 float32 reduce_dtype float32 else: # bf16 / fp16 训练时upcast 开启则 reduce 用 float32 reduce_dtype float32 if upcast else param return MixedPrecisionSpec( param_dtypeparam, reduce_dtypereduce_dtype, cast_forward_inputs(param ! float32), ) def require_upcast_effective(spec: MixedPrecisionSpec) - MixedPrecisionSpec: 结构性守护如果用户想要 upcastreduce 必须比 param 更宽。 wide {float32: 32, float16: 16, bfloat16: 16} assert wide[spec.reduce_dtype] wide[spec.param_dtype], ( fupcast 失效: reduce_dtype{spec.reduce_dtype} 不宽于 fparam_dtype{spec.param_dtype} ) return spec # 用法 spec derive_mixed_precision(mixed_precisionbf16, upcastTrue) spec require_upcast_effective(spec) print(spec)要点reduce_dtype字段类型从Optional改为必填从源头消灭None退化的可能derive_mixed_precision成为唯一推导入口任何重构都必须经过它require_upcast_effective在对象构造时就把upcast 意图未被满足转成显式异常而不是等到训练到一半才从 loss 曲线里怀疑。这一层把正确性从靠人记得传字段升级成结构强制是防回归的关键。七、解决方案第三层断言 / CI 守护光靠结构还不够。因为 bug 来自配置没落到运行时我们应该在运行时真正校验一次 reduce 的 dtype并写一条 pytest 把它锁进 CI确保任何重构都不会再让它退化为None。# test_fsdp2_mixed_precision.py import pytest # 下面是运行时自检用空 param 模拟一次 all-gather/reduce 路径 # 确认最终参与 reduce 的 dtype 就是配置的 reduce_dtype。 def effective_reduce_dtype(param_dtype, reduce_dtype): return reduce_dtype if reduce_dtype else param_dtype def test_upcast_not_silently_dropped(): upcast 必须生效不能静默退化成 param_dtype。 param_dtype bfloat16 reduce_dtype float32 # 用户配置 got effective_reduce_dtype(param_dtype, reduce_dtype) assert got float32, fupcast 被忽略实际 reduce_dtype{got} assert got ! param_dtype, reduce 与 param 同精度upcast 等于没生效 def test_reduce_wider_than_param(): reduce 应当宽于或等于param否则无精度收益。 wide {float32: 32, float16: 16, bfloat16: 16} param, reduce_ bfloat16, float32 assert wide[reduce_] wide[param] def test_none_is_rejected(): 重构若把 reduce_dtype 落回 None必须被测试捕获。 with pytest.raises(AssertionError): require_upcast_effective_unsafe(bfloat16, None) def require_upcast_effective_unsafe(param_dtype, reduce_dtype): eff effective_reduce_dtype(param_dtype, reduce_dtype) wide {float32: 32, float16: 16, bfloat16: 16} assert wide[eff] wide[param_dtype]把这条用例接进tox/ GitHub Actions 后任何漏传reduce_dtype的提交都会在 CI 里立即变红而不是等到用户训练到后期才发现。八、排查清单遇到loss 后期抖动、疑似精度不足时按下面顺序排查打印最终MixedPrecisionPolicy对象确认reduce_dtype不是Noneprint(policy) # 重点看 reduce_dtype 字段若reduce_dtype is None且你期望 upcast说明命中本 bug按第五 / 六节修复检查accelerate版本pip show accelerate | grep Version确认是否1.13.0在训练第一步用next(model.parameters()).dtype与一次 dummy forward 的梯度 dtype 交叉验证确认 reduce 路径确实用了更宽 dtype若为分布式子进程确认日志在每个 rank 都打印避免只看 rank0 被误导用第六节的require_upcast_effective在插件构造处加一道断言把问题提前到启动期暴露把第七节的 pytest 接进 CI作为回归护栏。九、小结acceleratev1.13.0 对 FSDP2 插件重构时漏传了MixedPrecisionPolicy的reduce_dtype导致用户配置的float32upcast 静默失效——不报错、不告警却让梯度 reduce 在bfloat16精度下完成长尾收敛与数值敏感算子因此劣化。三层层级第一层把reduce_dtype重新透传进MixedPrecisionPolicy恢复 upcast 行为第二层用MixedPrecisionSpec把推导收敛成单一入口并将reduce_dtype改为必填从结构上消灭None退化第三层在运行时真正校验 reduce dtype并加 pytest 锁进 CI防止重构再次漏字段。核心教训任何带缺省值且缺省即静默退化的配置项都是 silent no-op 类 bug 的高发地。对这类项最好的做法是——要么让默认显式报错要么用结构性断言在对象构造期就逼出意图与行为的不一致。