From fd687867362394bbbf657c0f99a065a5f263216c Mon Sep 17 00:00:00 2001 From: conver334 Date: Sun, 5 Jul 2026 20:18:40 -0700 Subject: [PATCH] lite: CP-split BSHD inputs in qwen3_moe forward (fixes #5617) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _forward_step_bshd() forwarded batch.input_ids/labels unchanged, but GQAttention's non-THD branch (and its RoPE frequency slicing) assumes its input is already CP-zigzag-sharded, and the pipeline runtime sizes PP activations from the CP-local sequence length. With CP>1 this left the final PP stage's CP-local logits disagreeing in shape with the still full-sequence labels, crashing vocab-parallel cross entropy (and, with PP=1, silently misaligning loss). Zigzag-split input_ids/labels/loss_mask with zigzag_slice_for_cp() — the same primitive other models (mla.py, dsa.py, gated_delta_net.py) already use for CP — so both sides stay consistent. unpack_forward_output() already reconstructs zigzag-CP outputs, so only the input side was missing. Co-Authored-By: Claude Sonnet 5 Signed-off-by: conver334 --- .../lite/model/qwen3_moe/lite/protocol.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/experimental/lite/megatron/lite/model/qwen3_moe/lite/protocol.py b/experimental/lite/megatron/lite/model/qwen3_moe/lite/protocol.py index ab02ccab832..b96b199c81f 100644 --- a/experimental/lite/megatron/lite/model/qwen3_moe/lite/protocol.py +++ b/experimental/lite/megatron/lite/model/qwen3_moe/lite/protocol.py @@ -44,7 +44,8 @@ normalize_lora_config, trainable_param_stats, ) -from megatron.lite.primitive.parallel import ParallelState, init_parallel +from megatron.lite.primitive.parallel import ParallelState, init_parallel, zigzag_slice_for_cp +from megatron.lite.primitive.parallel.thd import parallel_state_from_model from megatron.lite.primitive.recompute import apply_recompute, parse_recompute_spec from megatron.lite.runtime.contracts import OptimizerConfig, ParallelConfig from megatron.lite.runtime.contracts.data import PackedBatch @@ -133,8 +134,21 @@ def _forward_step(model: nn.Module, batch: PackedBatch) -> dict: def _forward_step_bshd(model: nn.Module, batch: PackedBatch) -> dict: + ps = parallel_state_from_model(model) or ParallelState() + input_ids = batch.input_ids.reshape(1, -1) labels = batch.labels.reshape(1, -1) if batch.labels is not None else None - return model(input_ids=batch.input_ids.reshape(1, -1), labels=labels, packed_seq_params=None) + loss_mask = batch.loss_mask.reshape(1, -1) if batch.loss_mask is not None else None + if ps.cp_size > 1: + # GQAttention's non-THD branch expects CP-zigzag-sharded inputs, and + # unpack_forward_output() reconstructs zigzag-CP outputs; match both. + input_ids = zigzag_slice_for_cp(input_ids, ps.cp_rank, ps.cp_size, seq_dim=1) + if labels is not None: + labels = zigzag_slice_for_cp(labels, ps.cp_rank, ps.cp_size, seq_dim=1) + if loss_mask is not None: + loss_mask = zigzag_slice_for_cp(loss_mask, ps.cp_rank, ps.cp_size, seq_dim=1) + return model( + input_ids=input_ids, labels=labels, loss_mask=loss_mask, packed_seq_params=None + ) def unpack_forward_output(model: nn.Module, batch: PackedBatch, output) -> Any: