[Megatron Lite] Fix qwen3_moe BSHD forward under context parallelism (#5617)#5661
Open
conver334 wants to merge 1 commit into
Open
[Megatron Lite] Fix qwen3_moe BSHD forward under context parallelism (#5617)#5661conver334 wants to merge 1 commit into
conver334 wants to merge 1 commit into
Conversation
_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 <noreply@anthropic.com> Signed-off-by: conver334 <conver334@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Fixes #5617.
In the experimental Megatron-Lite runtime, a Qwen3-MoE eval-only BSHD (
use_thd=false) forward with context parallelism and pipeline parallelism enabled together (CP>1,PP>1) crashes in vocab-parallel cross entropy: the final pipeline stage produces CP-local logits but still receives full-sequence labels.Root cause
_forward_step_bshd()forwardedbatch.input_ids/batch.labelsunchanged, but:GQAttention's non-THD branch (and its RoPE frequency slicing) assumes its input is already CP-zigzag-sharded (see the "q is CP-zigzag pre-sliced" comment inprimitive/modules/gqa.py)._infer_pipeline_tensor_shape()sizes PP activations from the CP-local sequence length.So with
CP>1the last PP stage's CP-local logits disagreed in shape with the still-full-sequence labels, crashingvocab_parallel_cross_entropy(and, withPP=1, silently misaligning the loss). Only the THD path did CP splitting (viapack_thd_forward_kwargs); the BSHD input side was missing it.Fix
_forward_step_bshd()now zigzag-splitsinput_ids/labels/loss_maskwithzigzag_slice_for_cp()— the same primitive other models (mla.py,dsa.py,gated_delta_net.py) already use for CP — so the input layout matches what attention/RoPE expect.unpack_forward_output()already reconstructs zigzag-CP outputs back to full-sequence order viaunpack_thd_forward_output(), so only the input side needed the change.Testing
tests/unit/model/test_qwen_config_unit.pypasses (CPU-only, no GPU required).CP=2, PP=2eval-only forward now completes and matches the CP-off baseline for bothdist_optand FSDP2.Scope
BSHD-only (
use_thd=false) path, as reported. The THD path already handled CP and is unchanged.