Skip to content

fix magic send sp/cp logic #1466

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
From00 merged 1 commit into from
Jul 15, 2026
Merged

fix magic send sp/cp logic #1466

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 0 additions & 19 deletions src/paddlefleet/transformer/multi_token_prediction.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -900,25 +900,6 @@ def forward(self, dict_args: dict):
:, (depth + 1) : (depth + 1 + seq_len)

@risemeup1111 risemeup1111 Jul 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 优先级：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]
        )

].contiguous()

@PaddlePaddle-bot PaddlePaddle-bot Jul 14, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔴 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 回归用例。


# Scatter input_ids to match decoder_input's local shape (CP/SP)
if cp_world_size > 1 and self.config.experimental_dataflow:
# ContextParallelScatterOp expects 3D; unsqueeze then squeeze
mtp_input_ids_local = ContextParallelScatterOp.apply(
mtp_input_ids_local.unsqueeze(-1),
axis=1,
mode=self.config.cp_balance_mode,
).squeeze(-1)
if self.config.sequence_parallel:
# [B, S] -> [B, S/tp]: flatten, scatter, reshape back
B_local = mtp_input_ids_local.shape[0]
mtp_input_ids_local = mtp_input_ids_local.reshape(
[-1]
).unsqueeze(-1)
mtp_input_ids_local = ScatterOp.apply(mtp_input_ids_local)
mtp_input_ids_local = mtp_input_ids_local.squeeze(-1).reshape(
[B_local, -1]
)

# Trim rotary embeddings to seq_len (once; seq_len is constant across depths)
_rotary_keys = (
"rotary_pos_emb",
Expand Down
4 changes: 3 additions & 1 deletion tests/single_card_tests/transformer/test_mtp_magic_send.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -513,7 +513,9 @@ def scatter_fn(x):
"labels": paddle.randint(0, 100, [B, S_global]),
}
)
self.assertEqual(cnt[0], 2) # decoder_input + mtp_input_ids_local
self.assertEqual(
cnt[0], 1
) # decoder_input only (mtp_input_ids_local is NOT scattered here; router does it)
self.assertEqual(list(captured["di"].shape), [S_local, B, H])

def test_cp_scatter(self):
Expand Down
Loading