[Bug] hp_fsdp2 + SkipDTensorDispatch 下 AdamW bf16 更新与 torch fsdp2 数值不一致

背景

在 veomni 上使用 4× Ascend 910B3、bf16、AdamW、seed=42、开启确定性（HCCL_DETERMINISTIC 等）的配置，对比：

• torch fsdp2：PyTorch 官方 FSDP2 + PyTorch DTensor
• hp_fsdp2：hyper-parallel FSDP2 + HP DTensor

发现从 step2 起 loss 逐步偏离（step1 一致），50 step 后 max loss diff 约 1e-3 量级。

---

现象

阶段	fsdp2 vs hp_fsdp2
step1 loss	完全一致
step2+ loss	逐步偏离
backward / grad clip 后	grad bit-exact 一致
optimizer.step() 后 param	开始分叉
MoE routing 统计	逐步一致（排除离散路由问题）

典型 case（lm_head.weight，step1）：

• grad ≈ -1.56e-08（极小，低于 bf16 ULP）
• fsdp2：param 不变（正确舍入）
• hp（问题版本）：param 跳 +0.000122（误跳 1 个 bf16 刻度）

---

根因分析

1. 分叉点：optimizer.step()，不是 backward / 梯度归约

numeric debug 打点结论：

• after_bwd / after_clip / before_optim：grad 完全一致
• after_optim：param 开始不一致

comm_fusion、fused=False 等已排除为主要原因。

2. 核心机制：SkipDTensorDispatch + 原生 torch.optim.AdamW

训练框架在 hp_fsdp2 下通常这样调用 optimizer：

with SkipDTensorDispatch():
    optimizer.step()

SkipDTensorDispatch 会 关闭 HP DTensor 的 op dispatch，将参数 unwrap 为 local bf16 tensor，直接跑 PyTorch 原生 AdamW。

相关代码路径：

• hyper_parallel/core/dtensor/dtensor.py：SkipDTensorDispatch → disable_dtensor_dispatch()
• hyper_parallel/core/shard/_op_dispatch.py：bypass 分支对 in-place op（含 addcdiv_）在 local tensor 上直接执行

if self._should_bypass_dispatch(op_name):
    result = op_call(*self._unwrap_args(args), **self._unwrap_kwargs(kwargs))
    ...

3. 与 torch fsdp2 的差异

	torch fsdp2	hp_fsdp2（问题路径）
参数类型	PyTorch DTensor	HP DTensor
optimizer 路径	PyTorch DTensor dispatch（无 Skip）	Skip → local bf16 裸执行
AdamW 内核	foreach=True（fused=False 时）	若强制 foreach=False 则 further 偏离
param 更新	torch._foreach_addcdiv_ on bf16	同 op，但 operand/state 路径不同 或整段 fp32 cast 舍入点不同

AdamW 最后一步：

param.addcdiv_(exp_avg, denom, value=-step_size)

在 bf16 上，当数学更新量 小于 bf16 ULP 时：

• PyTorch DTensor 路径：小更新可被正确舍入为 0
• Skip + local bf16 裸写（尤其 NPU kernel）：可能 误跳 1 个 bf16 刻度

一步里这样的元素很少，但下一步 forward 使用不同权重，loss 从 step2 开始累积偏离。

4. 中间修复尝试与结论

尝试	结果
fused=False	有改善，未消除
整段 optimizer 升 fp32 再 cast 回 bf16	step2 改善，但 exp_avg 与 fsdp2 不一致；param 写回仍差 1~2 ULP
自定义 AdamW：state 保持 bf16 + foreach 对齐 + param 用 bf16 foreach_addcdiv	step1 param/exp_avg bit-exact；50 step loss 完全一致

关键结论：要对齐 fsdp2，不是「算得更准（fp32）」，而是「走同一条 bf16 kernel 路径」。在 state 已与 fsdp2 对齐的前提下，param 写回应使用与 PyTorch AdamW 相同的 torch._foreach_addcdiv_，而不是 fp32 算完再 .to(bf16)。

---

建议 hyper-parallel 侧改进

短期（应用层 workaround，已在 veomni 验证）

在 SkipDTensorDispatch 内 不要直接调用 optimizer.step()，改为：

1. 自定义 AdamW，mirrors torch/optim/adam.py 的 _multi_tensor_adam（bf16 state 更新）
2. foreach / fused 与 fsdp2 保持一致（fused=False, foreach=True）
3. param 写回使用 torch._foreach_addcdiv_ on bf16 local shards（与 fsdp2 相同）

长期（hyper-parallel 框架层）

1. 提供官方 optimizer_step helper（或 document best practice），避免应用层各自 hack
2. 评估是否可以让 AdamW 关键 in-place op（addcdiv_ / lerp_ / mul_）在 Skip 模式下仍走 与 PyTorch DTensor 等价的 numerics，而非简单 unwrap 到 local
3. 增加 UT：hp_fsdp2 vs torch fsdp2 在 bf16 AdamW step 上 param delta bit-exact（可参考 veomni numeric debug 工具）

---

复现环境

• 硬件：4× Ascend 910B3
• torch==2.8.0 + torch_npu 2.8.0.post2
• bf16 + AdamW + seed=42 + 确定性开关
• 对比脚本：veomni_dev/scripts/zwd/run_fsdp2_vs_hp_fp32_master.sh（及 numeric debug 工具）

验证结果（修复后）

bf16 foreach_addcdiv + state bf16 foreach 对齐后：

• step1 numeric debug：param / grad / exp_avg / exp_avg_sq max_abs=0
• 50 step loss 字符串逐步完全一致（0/50 mismatch）

---

关联

该问题在下游 veomni hp_fsdp2 集成中定位并 workaround；根因位于 hyper-parallel 的 SkipDTensorDispatch + optimizer numerics 路径，建议在框架层正式修复或提供官方方案。

---

补充：若要在 hyper-parallel 框架内对齐，建议修改的代码

方案 A（推荐）：新增官方 optimizer step API

目的：把已在 veomni 验证过的 custom AdamW 下沉到 hyper-parallel，应用层不再各自实现。

文件	改动
新建 hyper_parallel/platform/torch/fully_shard/optimizer_step.py	实现 optimizer_step_adamw(optimizer, *, zero_grad=True)：在 SkipDTensorDispatch + no_grad 内，对 HP DTensor 的 local bf16 shard 执行与 torch/optim/adam.py::_multi_tensor_adam 等价的 foreach 路径；state 更新（lerp_/addcmul_/mul_）保持 bf16；param 写回使用 torch._foreach_addcdiv_
hyper_parallel/platform/torch/fully_shard/__init__.py	export 新 API
hyper_parallel/__init__.py	加入 optimizer_step_adamw（或挂到 HSDPModule.optimizer_step()）
新建 tests/torch/fully_shard/_test_optimizer_numerics.py	对比 hp local foreach AdamW step 与 torch bf16 foreach AdamW 在相同 tensor 上 param/exp_avg bit-exact

不必改 _op_dispatch.py 的 bypass 逻辑；仍在 Skip 下对 local tensor 执行，但由框架保证 op 序列与 PyTorch AdamW 一致。

---

方案 B：增强 SkipDTensorDispatch / bypass 数值语义

目的：继续允许 optimizer.step()，但 bypass 路径与 PyTorch DTensor numerics 一致。

文件	改动
hyper_parallel/core/dtensor/dtensor.py	扩展 SkipDTensorDispatch：如 SkipDTensorDispatch(optimizer_safe=True) 或在文档中明确 optimizer 专用模式
hyper_parallel/core/shard/_op_dispatch.py	_should_bypass_dispatch / bypass 分支（约 L1028–L1251）：对 AdamW 涉及的 in-place op（addcdiv_/lerp_/mul_/addcmul_ 及 NPU 映射名如 InplaceAddExt 等）不要简单 unwrap → local 裸调；改为调用与 PyTorch ATen 等价的 numerics wrapper（bf16 上 foreach 同 kernel）
hyper_parallel/core/shard/_op_dispatch.py	_INPLACE_BYPASS_OPS（L244–245）：若新增 optimizer numerics wrapper，确保 in-place 后仍返回原 DTensor self
hyper_parallel/platform/torch/dtensor.py	__torch_function__（L52–78）：可选提供不 Skip 的 AdamW op dispatch（长期目标：应用层无需 Skip）

---

方案 C：让 optimizer.step() 无需 Skip（长期）

目的：HP DTensor 完整走 dispatch，与 PyTorch DTensor 行为一致。

文件	改动
hyper_parallel/core/shard/_op_dispatch.py + core/shard/ops/	为 addcdiv_/lerp_/_foreach_addcdiv_ 等注册 layout-infer + local compute 的 distributed op（参考 torch/distributed/tensor/_ops/_pointwise_ops.py 对 aten.addcdiv_.default 的注册）
hyper_parallel/platform/torch/fully_shard/param_group.py	flat buffer rebase（L481–485）：optimizer in-place 写回 flat buffer 后，确认 DTensor _local_tensor / .data 同步逻辑与 numerics 路径一致
hyper_parallel/platform/torch/fully_shard/param.py	unsharded_param refresh（L572–576）：optimizer 改 sharded local 后 all-gather 视图不失真（已有注释，需与 optim numerics 联调）
examples/torch/llama3/README.md	更新训练循环说明：若方案 C 完成，可不再要求 Skip + 原生 optimizer.step()

---

测试 / 文档（任意方案均需）

文件	改动
新建 tests/torch/fully_shard/_test_fsdp2_optimizer_vs_torch.py	小模型 + bf16 + foreach AdamW：HP FSDP2 param delta vs torch FSDP2 bit-exact
.agent/skills/code-review/distributed-guidelines.md 或用户文档	补充：禁止在 Skip 下直接 optimizer.step() 而不保证 AdamW numerics；推荐调用官方 helper

---

不建议的改法

• 仅在应用层把 param 升到 fp32 再 .to(bf16)：舍入点与 PyTorch foreach_addcdiv 不同，无法 bit-exact。
• 仅改 foreach=False：与 torch fsdp2（foreach=True）kernel 路径不一致。

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

Start with hyper_parallel/core/dtensor/dtensor.py and hyper_parallel/core/shard/_op_dispatch.py, then run veomni_dev/scripts/zwd/run_fsdp2_vs_hp_fp32_master.sh to reproduce the optimizer divergence. Review the proposed optimizer-step and comparison tests, and consider the documented framework-layer options. Done means bf16 AdamW parameters and state are bit-exact with torch FSDP2 across the stated test case.