背景

在 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 量级。

现象

典型 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 的差异

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. 中间修复尝试与结论

关键结论：要对齐 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，应用层不再各自实现。

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

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

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

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

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

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

不建议的改法

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