[Bug] params_dtype=bfloat16 时 FSDP apply_reduced_grad 崩溃:reduced grad(fp32)与 param shard(bf16)dtype 不一致

问题描述

params_dtype = bfloat16(bf16 master 权重)时,动态图 FSDP 反向在归约后写回梯度处崩溃:

RuntimeError: The grad dtype and shape should be same as source tensor but got
  dtype: Float32 vs BFloat16  shape: {64640, 1024} vs {64640, 1024}

形状一致,仅 dtype 不一致:reduced grad 是 Float32,而 param shard 是 BFloat16。

调用栈(全程在 hyper-parallel FSDP 内):

hyper_parallel/platform/mindspore/autograd_compat.py:204  backward
hyper_parallel/platform/mindspore/fully_shard/hook_function.py:32  backward
hyper_parallel/platform/mindspore/fully_shard/scheduler.py:180  _backward_hook
hyper_parallel/core/fully_shard/hsdp_scheduler.py:259  _hsdp_backward_hook
hyper_parallel/platform/mindspore/fully_shard/state.py:559  post_backward
hyper_parallel/platform/mindspore/fully_shard/param.py:920  apply_reduced_grad
hyper_parallel/platform/mindspore/autograd_compat.py:75  grad
RuntimeError: The grad dtype and shape should be same as source tensor ...

复现环境

• MindSpore 2.10.0 / Ascend 910B2,CANN 9.1.0-beta.1
• hyper-parallel 15984d1(!821)
• 模型:DeepSeek-V3(MLA),动态图 PYNATIVE,mindformers master
• 配置:params_dtype: bfloat16 + compute_dtype: bfloat16,tensor_parallel: 2(2 卡即可复现);崩溃 weight 形状 {64640,1024} = vocab(129280/TP2)× hidden(1024)
• 同时叠加 PP(TP1×PP2)也是同一 bf16 根因先触发

根因分析

FSDP apply_reduced_grad(param.py:920)在把 reduce-scatter 后的梯度写回 param shard 时,经 autograd_compat.py:75 grad 校验 grad.dtype 必须等于 source(param)dtype。当 params_dtype=bf16 时:

• param shard 是 bf16;
• 但反向/归约得到的 reduced grad 是 fp32(loss 与梯度累加在 fp32 下进行,属常见做法);
• 二者 dtype 不一致 → 断言失败。

标准混合精度(params_dtype=float32 master + compute_dtype=bfloat16)下 grad 与 param 都是 fp32,故不触发;一旦 master 权重设为 bf16 即崩。

建议

FSDP 在 apply_reduced_grad 写回前,应把 reduced grad cast 到 param 的 dtype(或在 grad 不变量校验时允许 dtype 不同、由框架统一 cast),而不是直接断言 grad.dtype == param.dtype。这样 bf16 master 权重(grad 在 fp32 归约后降回 bf16)才能正常训练。

影响

任何 params_dtype=bfloat16(bf16 master 权重)的动态图 FSDP 训练首个 step 即崩,无法使用纯 bf16 master 权重配置。

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 215
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/215

Start with apply_reduced_grad in hyper_parallel/platform/mindspore/fully_shard/param.py around line 920, then inspect autograd_compat.py around line 75 to understand the dtype check. Reproduce with params_dtype and compute_dtype set to bfloat16 and tensor_parallel 2; done means the reduced fp32 gradient is accepted for the bf16 parameter shard without the reported crash.