[Performance][MindSpore][PP+HSDP] VPP Stage AllReduce 阻塞 BWD_SEND，造成 474 ms PP 气泡

背景

在 300B MoE 的 MindSpore PyNative 训练中观察到，VPP Stage 的参数梯度归约位于上游 PP 反向梯度发送的关键路径上。HSDP 组内负载不均产生的 AllReduce 长尾会沿下游到上游逐级传播，最终表现为计算流上数百毫秒的 EVENT_WAIT 和 PP bubble。

本问题仅涉及 MindSpore 后端；Torch 后端不在本 issue 范围内。

相关但不重复的已有工作：

• #120：HSDP 内部异步融合 AllReduce 流水线优化；
• #184：FSDP/HSDP 通信融合、RS/AR overlap 与零拷贝方案。

本 issue 聚焦 PP/VPP scheduler 与 HSDP reduction 的跨子系统调度顺序：即使 HSDP 内部 collective 已异步下发，只要其 wait() 落在 PP BWD_SEND 之前，仍会把参数梯度归约长尾暴露到 PP 关键路径。

复现场景

• 模型：300B MoE
• 后端：MindSpore PyNative
• 序列长度：4K
• Pipeline Parallel：PP4
• Virtual Pipeline：VPP2
• microbatch 数：16
• pipeline_parallel_overlap_p2p: True
• pipeline_parallel_overlap_b_f: True
• pipeline_parallel_p2p_transport: multi_stream
• profiling：profiling/pynative/4k/pp4_vpp2_4k_0912.tar.gz

Stage 分配：

pipeline_parallel_interleave_num: 2
pipeline_parallel_layers_per_stage:
  - "0-5,24-29"
  - "6-11,30-35"
  - "12-17,36-41"
  - "18-23,42-44"

当前调度逻辑

旧版 add_fsdp_reduce_grad() 在检测到某个 virtual Stage 的最后一个 microbatch backward 后，立即插入 FSDP_REDUCE_GRAD：

BWD(micro=last, stage=S)
FSDP_REDUCE_GRAD(stage=S)
BWD_SEND / BATCH_SEND_RECV
后续本地 VPP Stage 计算

例如 rank 12288 的 Stage 7 尾部近似为：

BWD(micro15, Stage7)
FSDP_REDUCE_GRAD(Stage7)
BATCH_SEND_RECV:
  BWD_SEND(micro15, Stage7) -> rank8192
  BWD_RECV(micro12, Stage3) <- rank0
BWD(micro12, Stage3)

rank 8192 的 Stage 6 同理：

BWD(micro15, Stage6)
FSDP_REDUCE_GRAD(Stage6)
BATCH_SEND_RECV:
  BWD_SEND(micro15, Stage6) -> rank4096
  BWD_RECV(micro12, Stage2) <- rank12288
BWD(micro12, Stage2)

stage.launch_reduce_grad() 不只是异步发起 collective，还会推进 HSDP RS/AR 状态机并等待/应用先前 pending 的 reduction。MindSpore 的 CommHandle.wait() 会在当前 device stream 上插入 aclrtStreamWaitEvent。旧逻辑没有切换 stream，因此 wait 落到主计算流（本例为 Stream 47），阻塞后续 batchSendRecv 的提交。

Profiling 证据

以用户看到的 rank 4096 事件为起点：

EVENT_WAIT on Stream 47
start:    5390.003955 ms
duration: 474.035280 ms

它对应 rank 4096 与 rank 8192 之间的 PP hcom_batchSendRecv__982_37_1，约 8 MiB BFP16。实际 P2P 匹配完成只需约 0.5 ms，绝大部分时间在等下游 rank 8192 提交匹配通信。

完整依赖链：

Rank	相对时间	状态
12288	5153.335–5626.377 ms	Stream 47 暴露一串 HSDP AllReduce 尾巴，约 473.043 ms
12288	5626.379 ms	AllReduce 结束后约 2.25 us 立即启动到 rank8192/rank0 的 PP batchSendRecv
8192	5219.773–5627.465 ms	等待 rank12288 的反向梯度，约 407.692 ms
8192	5627.467–5730.712 ms	反向计算，约 103.245 ms
8192	5730.713–5863.529 ms	等待 32-rank、50,331,648 FP32（约 192 MiB）的 HSDP AllReduce，约 132.816 ms
8192	5863.532 ms	AllReduce 结束后约 2.55 us 立即启动到 rank4096 的 PP batchSendRecv
4096	5390.004–5864.039 ms	等待 rank8192 返回反向梯度，共 474.035 ms

因果关系：

rank12288: BWD -> HSDP AR wait -> BWD_SEND to rank8192
                                      |
rank8192:       wait recv -> BWD -> HSDP AR wait -> BWD_SEND to rank4096
                                                       |
rank4096:                  wait recv (474 ms) <---------+

rank 12288 的长尾涉及至少两类 replicate group：32-rank group 与 128-rank group。profiling 中 collective 的 Wait Time Ratio 接近 1，说明主要是 HSDP 组内到达不均衡，而非 8 MiB PP 传输带宽问题。当前 profiling 只采集了一条 PP replica，无法从该压缩包确定具体哪个 HSDP peer 是 straggler。

根因

参数梯度与输入梯度存在两条不同依赖：

参数梯度：BWD -> HSDP ReduceScatter/AllReduce -> optimizer
输入梯度：BWD -> BWD_SEND -> 上游 Stage backward

BWD_SEND 只依赖 BWD 已产生的输入梯度 dx，不依赖参数梯度 AllReduce 完成。当前 scheduler 将二者不必要地串行化为：

BWD -> HSDP reduction wait -> BWD_SEND

因此 HSDP collective 的 straggler tail 被放大成跨多个 PP rank 的关键路径气泡。现有 overlap_p2p=True 和 multi_stream 无法解决，因为发送端在 reduction wait 结束前尚未提交 P2P。

建议优化方案

建议为 MindSpore PP+VPP+HSDP 增加 scheduler-owned 的异步 Stage reduction：

BWD完成
  -> 优先提交关联的 BWD_SEND / batchSendRecv
  -> 在独立 MindSpore device stream 上推进该 Stage 的 HSDP reduction
  -> 主计算流继续后续本地 VPP Stage 计算和 P2P
  -> optimizer/最终梯度使用前，通过 device event 统一 drain

实现约束：

1. 不创建 Python 或 OS 线程，只使用 MindSpore NPU Stream 与 Event；
2. 仅修改 MindSpore 后端路径，Torch 行为保持不变；
3. 在主流记录 compute_ready event，reduction stream 等待后再调用 stage.launch_reduce_grad()；
4. reduction stream 完成后记录 reduce_done event，主流只在最终 optimizer/下一次必须消费梯度前等待；
5. 非最后一个本地 VPP Stage 的 reduction 可异步；最后一个 reduction/最终 drain 保持正确同步；
6. BWD_SEND 应优先于大块 AllReduce 提交，避免 collective 抢占通信资源后继续延迟 PP 关键路径；
7. 正确处理 reshard_after_backward、显式 FSDP_RESHARD 与 reduction stream 的所有权，避免同一 Stage 重复 reshard 或跨 stream 竞态；
8. 保持 P2P per-peer FIFO、现有 batch 配对形状及 multi-stream communicator 行为不变。

预期收益

在该 profiling 点，如果 rank 12288/8192 的参数梯度归约不再阻塞 BWD_SEND，rank 8192 的反向梯度有机会在 rank 4096 到达 5390 ms 等待点之前准备完成。该处 474 ms 的 PP wait 理论上可压缩到接近实际 P2P 匹配耗时（亚毫秒级）。

AllReduce 本身不会消失，实际端到端收益取决于它与后续 VPP chunk 计算、P2P 的覆盖比例和 HCCL 资源竞争，需要重新 profiling 验证，不能直接把 474 ms 全部计为端到端收益。

验收标准

• MindSpore 后端实现，不引入额外 Python/OS 线程；
• 非最后本地 VPP Stage 的最终 BWD 后，主计算流不再因 HSDP AllReduce EVENT_WAIT 阻塞关联的 BWD_SEND；
• profiling 中 PP batchSendRecv 能在 Stage reduction 完成前提交/推进；
• 第一组长 AllReduce 与后续 VPP chunk 计算形成可观测重叠；
• 8 卡 PP2/VPP2/FSDP（DP4）多步训练无死锁、loss/grad 与基线一致；
• 8 卡 PP2/VPP2/HSDP（2x2）多步训练无死锁、reduction 数量与梯度结果正确；
• PP1、VPP1、非 HSDP、overlap_p2p=False 等路径无行为回归；
• 在独占卡环境重新采集 profiling，报告 AllReduce overlap 比例和端到端 step time 改善。

风险点

• HSDP AllReduce 与 PP P2P 并发时可能竞争 HCCL/链路资源，应确保 PP 关键路径优先；
• MindSpore CommHandle.wait() 绑定调用时的 current stream，必须保证所有 reduction 状态机 wait 都在 reduction stream 上执行；
• HSDP pending queue/param buffer 生命周期必须持续到 reduce_done，不能被后续 Stage reshard 或 optimizer 提前复用；
• 异步 Stage reduction 之间需要显式 event 串行化，避免共享 HSDP 全局 pending queue 被并发推进。

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

Start at add_fsdp_reduce_grad() and stage.launch_reduce_grad(), then inspect how MindSpore CommHandle.wait() uses the current stream. Use profiling/pynative/4k/pp4_vpp2_4k_0912.tar.gz to confirm the BWD_SEND stall and trace the reduction state and reshard ownership. Done means asynchronous reduction preserves synchronization and passes the listed 8-card, PP1/VPP1, non-HSDP, and overlap-disabled checks without regressions.