mindspore-ai / mindspore-ai/hyper-parallel
[Bug]: fully_shard replicate_params 的 BF16 临时参数覆盖 FP32 主参数
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 53
- Forks
- 63
- Avg merge
- 23h 45m
- Merged PRs (30d)
- 63
Description
Checklist
- 1. I have searched the existing issues (https://gitcode.com/mindspore/hyper-parallel/issues).
- 2. I have read the relevant documentation.
- 3. I have created a minimal reproduction case that clearly demonstrates the issue, including a complete code example and the error message with full traceback and error logs.
🐛 Describe the bug
fully_shard 管理的 replicate_params 以 FP32 初始化,并通过 MixedPrecisionPolicy(param_dtype=torch.bfloat16) 使用 BF16 参与前向计算时,Torch 与 MindSpore 两端的 HSDPParam.to_sharded() 都会把临时 unsharded 参数复制回 sharded 主参数。
由于临时参数已被转换成 BF16,这次复制会将 BF16 数值写回 FP32 sharded_param。虽然目标张量的 dtype 仍显示为 FP32,但参数有效精度已永久舍入为 BF16;后续优化器更新不再基于原始 FP32 主参数。
最小复现流程:
model = ExistingModel().to("npu") # 参数初始化为 FP32
replicate_params = set(model.target.parameters())
fully_shard(
model,
mesh=mesh,
reshard_after_forward=True,
mp_policy=MixedPrecisionPolicy(param_dtype=torch.bfloat16),
replicate_params=replicate_params,
)
master_before = [param.detach().clone() for param in model.target.parameters()]
model(inputs)
# 当前行为:forward/reshard 后 FP32 主参数被 BF16 临时参数覆盖
该问题是数值精度劣化,不会产生 Python traceback。
Expected behavior
to_sharded() 只负责恢复 module 上的 sharded 参数对象并释放 unsharded 临时存储,不应把低精度临时参数反向覆盖到 FP32 主参数。前向/reshard 后,replicate_params 的 FP32 主参数 dtype 与数值都应保持不变,优化器继续基于 FP32 主参数更新。
Additional context
Torch 与 MindSpore 实现存在相同复制逻辑,需要保持两端语义一致。回归验证复用现有 Torch replicate_params 精度场景:standalone 网络整体转换为 BF16,fully_shard 网络仍以 FP32 初始化并设置 param_dtype=BF16;不新增测试用例或 module 类。
Environment info
- Repository:
mindspore/hyper-parallel - Base:
upstream/masteratbb62140e - Backends: PyTorch and MindSpore
- Target device: Ascend NPU
schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 303
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/303
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Inspect the Torch and MindSpore implementations of HSDPParam.to_sharded(), then review the existing Torch replicate_params precision scenario described in the issue. Verify that a BF16 temporary parameter does not change the FP32 replicate parameter's dtype or values after forward and reshard, while both implementations retain consistent semantics.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100