fix(qwen3.5-vl): cp_size 1 needs no packed-sequence divisibility check - #2376
Open
shifangx wants to merge 1 commit into
Open
fix(qwen3.5-vl): cp_size 1 needs no packed-sequence divisibility check#2376shifangx wants to merge 1 commit into
shifangx wants to merge 1 commit into
Conversation
get_packed_cp_local_indices maps a CP rank's local tokens back to the packed stream using Megatron's two-chunk zigzag layout, which splits every sequence into 2 * cp_size chunks -- hence the divisibility requirement. But the check runs unconditionally, and at cp_size 1 there is nothing to split: one rank owns the whole stream and the correct local view is the identity. qwen3_5_vl.py:163 (_inject_vision_embeddings) calls this on every training step with cp_size = the configured context-parallel size, so with --context-parallel-size 1 -- what examples/geo3k_vlm/run_geo3k_qwen35.sh itself sets -- any pack with an odd token count is rejected: ValueError: Packed sequence length 4551 must be divisible by 2 * CP size 1 About half of all packed sequences have an odd length, so the Megatron VL path cannot complete its first optimizer step as configured. Seen on GEO3K at rollout 0 on 1 node and on 2 nodes alike. The fix short-circuits cp_size == 1 to the identity range and leaves the zigzag path untouched: wherever the loop previously succeeded at cp_size 1 (even lengths) it computed exactly that range already -- chunk = len/2, rank 0 takes [start, start+chunk) then [start+chunk, end) -- so this changes no output that was previously produced, it only stops rejecting inputs that were always valid. Empty cu_seqlens now returns an empty tensor instead of an IndexError. The existing test covers cp_size=2 with length-8 sequences, i.e. only the case that works; this adds the cp-1 odd-length and multi-sequence cases. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.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
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
get_packed_cp_local_indices(slime_plugins/models/qwen3_5_vl_utils.py) rejects anypacked sequence whose length is not divisible by
2 * cp_size�~@~T including atcp_size == 1, where there is nothing to split.Reproduction
examples/geo3k_vlm/run_geo3k_qwen35.shsets--context-parallel-size 1. Run aswritten on GEO3K (Qwen3.5-35B-A3B, Megatron backend), it dies at the first
optimizer step of rollout 0:
Reproduced on 1 node �
W 8 B200 and on 2 nodes �W 8 B200, failing on the same packedlength. Roughly half of all packed sequences have an odd token count, so this is
not a rare edge case: the Megatron VL path cannot complete a single gradient step
at CP 1.
Why
The function maps a CP rank's local tokens back to the packed stream using
Megatron's two-chunk zigzag layout, which splits every sequence into
2 * cp_sizechunks �~@~T hence the divisibility requirement. The check runs unconditionally, but at
cp_size == 1one rank owns the whole stream and the correct local view is theidentity.
The fix
Short-circuit
cp_size == 1to the identity range and leave the zigzag pathuntouched.
It changes no output that was previously produced: at
cp_size == 1with an evenlength the existing loop already computes exactly that range (
chunk = len / 2,rank 0 takes
[start, start + chunk)then[start + chunk, end)). It only stopsrejecting inputs that were always valid. Empty
cu_seqlensnow returns an emptytensor instead of raising
IndexError.Tests
tests/test_qwen3_5_vl_native.pycoveredcp_size=2with length-8 sequences �~@~T thecase that works. This adds
test_thd_cp_indices_are_the_identity_without_context_parallelismfor the CP-1 odd-length and multi-sequence cases. The existing
cp_size=2expectations are unchanged.