[NPU][Megatron GRPO] 开启 PP 后 P2P 通信在 HcclGroupEnd 报错并导致进程崩溃

Checklist / 检查清单

• I have searched existing issues, and this is a new bug report. / 我已经搜索过现有的 issues，确认这是一个新的 bug report。

Bug Description / Bug 描述

问题现象

在 Ascend NPU 上使用 ms-swift Megatron GRPO 训练，设置 pipeline_model_parallel_size > 1 后，PP stage 间 P2P 通信报错并导致训练进程直接崩溃。

错误最终在 torch_npu._C._npu_synchronize 暴露，核心报错为：

HcclGroupEnd: not in a group call. Didn't call HcclGroupStart before

GRPO 样本默认是变长序列，本次没有开启 padding_free。错误发生在 PP 变长 tensor shape 通信/后续 P2P 通信附近，可见 _communicate_shapes()、isend/irecv 或 batch_isend_irecv 相关调用栈。

已尝试的配置

以下配置均仍会复现：

• --batch_p2p_comm true
• --batch_p2p_comm false
• 关闭 P2P overlap

期望行为

NPU 上开启 PP 后，GRPO 变长序列能够正常完成 stage 间 shape 和 tensor 的 P2P 通信，不应出现 HCCL GroupStart/GroupEnd 状态未配对或进程崩溃。

How to Reproduce / 如何复现

环境

• 硬件：Ascend 910C
• ms-swift：4.5.0
• MindSpeed：0.16.0
• megatron-core：0.16.0
• PyTorch：2.9.0+cpu
• torch_npu：2.9.0.post4
• CANN：9.0.0

复现步骤

1. 使用 ms-swift Megatron GRPO 启动训练。
2. 不开启 padding_free，输入样本/response 保持 GRPO 默认的变长序列。
3. 设置 --pipeline_model_parallel_size 2（或更大）。
4. 开始训练，PP stage 间执行 shape/tensor P2P 通信时进程直接崩溃。

关键参数示例：

swift rlhf \\
  --rlhf_type grpo \\
  --pipeline_model_parallel_size 2 \\
  --batch_p2p_comm false \\
  ...

将 pipeline_model_parallel_size 设为 1 时不进入该 PP 报错路径。

Additional Information / 补充信息

模型：Qwen3.5-35B-A3B

相关代码线索

ms-swift 在 swift/megatron/init.py 中对 Megatron _batched_p2p_ops 做了 patch，并强制设置：

kwargs['group'] = None

另外，MindSpeed 已有将 NPU PP 的 batch P2P 替换为独立 P2P ops 的修改：

• https://gitee.com/ascend/MindSpeed/pulls/2707

希望协助确认：

1. ms-swift 的 group=None patch 在 NPU PP 场景是否符合预期；
2. ms-swift 是否需要接入/启用 MindSpeed 对 NPU PP P2P 的替换；
3. _communicate_shapes() 的变长 shape 通信是否需要使用相同的 PP process group 和一致的异步通信顺序。

Start with the Megatron communication patch in swift/megatron/init.py, especially the _batched_p2p_ops group=None setting, and trace _communicate_shapes() during variable-length GRPO batches. Reproduce with pipeline_model_parallel_size 2, padding_free disabled, and batch_p2p_comm both true and false. Done means PP shape and tensor communication completes on Ascend NPU without the HcclGroupStart/End error or process crash.