fix magic send sp/cp logic

[FeixLiu]

PR Category

Distributed Strategy

PR Types

Bug fixes

Description

fix magic send sp/cp logic

是否引起精度变化

否

[risemeup1111]

发现一个需要先修复的 SP/EP MoE 路径回归风险，详细说明已放在行内评论。当前 CI 还有多项测试在运行中；我本地尝试跑 tests/single_card_tests/transformer/test_mtp_magic_send.py，但评审环境缺少 paddle 包，未能完成本地复验。

^{Powered by Nyanpasu with gpt-5.5 xhigh, please check the suggestions carefully.}

[risemeup1111]

优先级：P1

问题： 这里现在把 mtp_input_ids_local 保持为全局 [B, S] 后直接传给 MTP transformer，但 decoder_input/hidden_states 在 sequence parallel 下已经被切到 [S/tp, B, H]。当 MTP 层后面是 EP MoE（expert_model_parallel_size > 1）时，MoELayer.forward 不会先 gather SP hidden states，TopKRouter.forward 的 3D SP 分支也没有对 input_ids 做 TP scatter，会在形状断言处看到 [B, S] vs [B, S/tp] 并失败；当前测试用 mock 的 _proj_and_transformer_layer 只数了 MTP 层内 scatter 次数，没覆盖到这个 router 断言。

影响： enable_mtp_magic_send + sequence_parallel + tensor_model_parallel_size > 1 + expert_model_parallel_size > 1 的 MoE 训练会在路由阶段因输入 id 形状不匹配中断，或者无法正确屏蔽 padding。

期望处理： 请保留“避免重复 CP/SP scatter”的目标，但把 SP 本地化补到实际消费 input_ids 的路径；例如在 router 的 3D SP 分支中，仅当 input_ids.shape[1] 仍是全局长度时再做一次 scatter，并为上述组合增加回归测试。

处理要求：请针对该评论修复并提交新的 commit。

参考实现形态：

if input_ids is not None and self.sequence_parallel:
    batch_size_, seq_len_ = input_ids.shape
    if (
        seq_len_ != seq_len
        and self.tensor_model_parallel_size > 1
        and seq_len_ == seq_len * self.tensor_model_parallel_size
    ):
        input_ids = ScatterOp.apply(input_ids.reshape([-1])).reshape(
            [batch_size_, -1]
        )

[PaddlePaddle-bot]

🤖 Paddle-CI-Agent | pr_review | 2026-07-15 07:50:05

📋 Review 摘要

PR 概述：调整 MTP magic send 在 CP/SP 下对 mtp_input_ids_local 的散射职责。
变更范围：src/paddlefleet/transformer/multi_token_prediction.py 与对应单卡测试。
影响面 Tag：Transformer MTP MoE Distributed Strategy

问题

级别	文件	概述
🔴 Bug	src/paddlefleet/transformer/multi_token_prediction.py:901	SP + EP 的 MTP MoE 路径会把全局 input_ids 传给本地 router，触发 shape mismatch 或 padding mask 错位

📝 PR 规范检查

标题缺少规范 Tag；描述结构符合当前模板。

标题建议（可直接复制）：

[Bug fixes] fix magic send sp/cp logic

总体评价

CP 的 ids scatter 可以下沉到 router，但 SP 的 3D MoE router 分支当前没有同步下沉，删除本层 scatter 会破坏 sequence_parallel=True 且 expert_model_parallel_size>1 的 magic send 训练路径。建议补齐 router 侧 3D SP scatter，或在 MTP 层仅保留 EP>1 场景下的 SP scatter，并增加覆盖真实 router 的回归测试。

[PaddlePaddle-bot]

🔴 Bug 删除这里对 mtp_input_ids_local 的 SP 本地化后，sequence_parallel=True 且 expert_model_parallel_size>1 的 MTP MoE 路径会把全局 [B, S] ids 传到本地 router。

decoder_input 在上面已经被 ScatterOp 切成 [S/tp, B, H]；MoELayer.forward 只有 EP<=1 时才 gather，EP>1 会保持本地 hidden states；而 TopKRouter.forward 的 3D input 分支只做 CP scatter，随后断言 input_ids.shape == [B, S/tp]。因此当前 mtp_input_ids_local 保持 [B, S] 会在 router 处触发 shape mismatch，或导致 padding mask 与本地 token 错位。新增测试用 proj_override 提前截住了 _proj_and_transformer_layer，没有覆盖真实 router。

建议修复方式：
在保留本次 CP 下沉意图的同时，让 SP ids 与 MoE 实际输入形状对齐。例如只在 EP>1 的 3D router/MTP 路径执行 SP scatter：

if input_ids is not None and self.sequence_parallel and self.config.expert_model_parallel_size > 1:
    b, _ = input_ids.shape
    input_ids = ScatterOp.apply(input_ids.reshape([-1])).reshape([b, -1])

并补一个能走到真实 TopKRouter 的 SP+EP magic send 回归用例。

[From00]

LGTM