[Bug]: MindSpore FSDP shard size=1 时无效 backward prefetch 导致训练轨迹分叉

Checklist

• 已检索现有 Hyper-Parallel issues。
• 已阅读 fully_shard/prefetch 相关实现。
• 已构造可稳定比较的 A/A 复现场景，并完成逐层、逐 rank 梯度快照定位。

问题描述

在 MindSpore backend 下，模型模块虽然经过 fully_shard 包装，但实际 FSDP shard group size 为 1 时，跨 TransformerLayer 的 backward prefetch 仍调用 unshard(async_op=True)。

该场景不存在参数 AllGather，也没有可重叠的异步通信。MindSpore 参数本地 unshard 分支只进行本地 buffer 建立/复制并返回 allgather_handle=None；目标层真正进入 backward 时又执行一次 unshard，造成参数 buffer 生命周期及 replay/backward 调度交叠。相同 seed、相同配置重复运行会出现训练轨迹概率性分叉。

问题不是 DSA/Indexer 独有：关闭 DSA、只运行 MLA 时同样可以复现。

关联修复 PR：#1326

环境与版本

• Hyper-Parallel master 基线：eccce0bc227d007c948ae04076d91622002a13c8
• MindFormers master 基线：dbbf170529fd6464bda46e314caad686182972a7
• 8 卡单机
• deterministic: true
• AdamW
• 模型包含 MoE shared expert

最小复现拓扑

parallelism:
  data_parallel_shard: -1
  tensor_parallel: 2
  context_parallel: 4
  pipeline_parallel: 1

training:
  deterministic: true

model:
  experimental_attention_variant: mla
  n_shared_experts: 1

实际并行拓扑为：

data_parallel=1 (dp_shard=1 × dp_replicate=1)
tensor_parallel=2
context_parallel=4

因此 TransformerLayer 虽为 HSDPModule，但参数 shard_size=1，不存在真实 FSDP AllGather。

完整复现配置：

/home/l00913161/lzy/tp2_cp4_determinism_20260903/configs/mla_tp2_cp4_adamw_same_model_100.yaml

使用同一配置、同一 seed 连续运行两次，各训练 100 step，然后逐 step 比较 loss 和 grad_norm。

原始实现的 A/A 结果

MLA、TP2+CP4、AdamW、无重计算，100-step A/A：

指标	首个差异 step	MAE	P95	最大绝对误差
loss	4	3.6890e-5	1.0100e-4	2.6100e-4
grad_norm	3	3.0893e-4	8.1200e-4	1.4444e-2

拓扑隔离结果：

并行拓扑	loss A/A	grad_norm A/A
TP1+CP1	完全一致	完全一致
TP2+CP1	完全一致	完全一致
TP1+CP4	完全一致	完全一致
TP2+CP4	出现漂移	出现漂移

首次分叉证据

首步在 global grad norm 计算前采集 8 rank × 29 参数，共 232 个更新前梯度 SHA，两次运行全部相同，说明 backward 数学计算和 TP reduce 在首步没有产生不同值。

随后使用无 Host 同步快照：每 step 只在 NPU clone 4 个 shared FC 梯度，训练结束后统一计算 SHA。两次运行各得到 3200 条记录：

• step 1–45：8 rank、4 参数全部一致；
• step 46：仅 rank 2/3 的 decoder.layers.0.mlp.shared_experts.linear_fc1.weight 首次不同；
• step 46 的打印 loss/grad_norm 仍一致；
• step 47：差异扩散到所有 rank 的 shared FC 梯度和 loss/grad_norm；
• 100 step 共 1730 条 shared 梯度快照不同，54 个训练 step 的轨迹不同。

Prefetch 边界实验

设置	100-step A/A 结果
全部 prefetch 开启	漂移
全部关闭	完全一致
仅关闭 backward prefetch	完全一致
仅关闭 forward prefetch	仍漂移，grad step 26 起
仅保留 tail→last-layer backward prefetch	完全一致
只关闭 layer1→layer0 backward prefetch	完全一致

该结果将首个异常边界定位到后一层 backward pre-hook 对前一层发起的无效 prefetch。

调用链与根因

MindFormers 为相邻 TransformerLayer 建立 backward prefetch 链：

layer1 backward_pre_hook
  -> prefetch_state(layer0).prefetch()
  -> HSDPState.prefetch()
  -> unshard(async_op=True)

基础实现无条件调用：

def prefetch(self) -> None:
    self.unshard(async_op=True)

但 MindSpore 参数 unshard 在 shard_world_size <= 1 时不会发起 collective，只准备或复制本地 buffer，并将 async handle 保持为 None。因此这次 prefetch 没有通信收益，目标层实际进入 backward 时还会再次走同步 unshard。

这里必须以 _init_group_infos() 解析出的 shard_size 判断真实 FSDP communication group 大小，不能使用可能包含原 DTensor/TP 分片语义的 shard_world_size。

期望行为

• 当一个 HSDP state 的全部参数均为 shard_size=1 时，跳过 prefetch，不提前执行本地 unshard/buffer copy。
• 任一参数 shard_size>1 时，保留原有异步 prefetch，使真实 FSDP AllGather 能继续与计算重叠。
• 不改变前向、反向、梯度归约和 optimizer 数学语义。

修复方案

在 MindSpore backend 的 MindSporeHSDPStateV2.prefetch() 中限定：

def prefetch(self) -> None:
    if not any(getattr(param, "shard_size", 1) > 1 for param in self.hsdp_params):
        return
    super().prefetch()

修改放在 MindSpore backend，不改变共享 HSDP 核心逻辑，也不影响 PyTorch backend。

修复验证

• TP2+CP4、DP/FSDP1：300-step A/A，loss/grad_norm 逐 step 完全一致。
• MindSpore-backend 最终实现：TP2+CP4 100-step A/A，loss/grad_norm 完全一致。
• TP2+CP1、DP/FSDP4：100-step A/A 完全一致，真实参数 AllGather prefetch 保持启用。
• DSA2、TP2+CP4、full recompute，叠加 MindFormers 完整层 replay 修复：300-step loss/indexer_loss/grad_norm A/A 完全一致。
• 新增单测 2/2 通过：size-one 不调用 async unshard，shard-size=4 仍调用。
• git diff --check、py_compile 通过。

完整分析报告及原始日志索引：

/home/l00913161/lzy/tp2_cp4_determinism_20260903/TP2_CP4_DETERMINISM_ANALYSIS.md

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

Review linked PR #1326 and the MindSpore backend's MindSporeHSDPStateV2.prefetch() implementation. Check how _init_group_infos() supplies shard_size and inspect the two reported unit tests for size-one and shard-size=4 behavior. Done means size-one states skip async unshard while real FSDP AllGather prefetch remains enabled.