diff --git a/src/paddlefleet/transformer/multi_token_prediction.py b/src/paddlefleet/transformer/multi_token_prediction.py index 108cd0232..2db2bb582 100755 --- a/src/paddlefleet/transformer/multi_token_prediction.py +++ b/src/paddlefleet/transformer/multi_token_prediction.py @@ -900,25 +900,6 @@ def forward(self, dict_args: dict): :, (depth + 1) : (depth + 1 + seq_len) ].contiguous() - # 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", diff --git a/tests/single_card_tests/transformer/test_mtp_magic_send.py b/tests/single_card_tests/transformer/test_mtp_magic_send.py index 471f75982..5c5535eb6 100644 --- a/tests/single_card_tests/transformer/test_mtp_magic_send.py +++ b/tests/single_card_tests/transformer/test_mtp_magic_send.py @@ -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):