Skip to content

fix(qwen3.5-vl): cp_size 1 needs no packed-sequence divisibility check - #2376

Open
shifangx wants to merge 1 commit into
THUDM:mainfrom
shifangx:shifang/fix-cp1-packed-indices
Open

fix(qwen3.5-vl): cp_size 1 needs no packed-sequence divisibility check#2376
shifangx wants to merge 1 commit into
THUDM:mainfrom
shifangx:shifang/fix-cp1-packed-indices

Conversation

@shifangx

@shifangx shifangx commented Sep 10, 2026

Copy link
Copy Markdown

What

get_packed_cp_local_indices (slime_plugins/models/qwen3_5_vl_utils.py) rejects any
packed sequence whose length is not divisible by 2 * cp_size �~@~T including at
cp_size == 1, where there is nothing to split.

Reproduction

examples/geo3k_vlm/run_geo3k_qwen35.sh sets --context-parallel-size 1. Run as
written on GEO3K (Qwen3.5-35B-A3B, Megatron backend), it dies at the first
optimizer step of rollout 0:

slime_plugins/models/qwen3_5_vl.py:227       forward
slime_plugins/models/qwen3_5_vl.py:163       _inject_vision_embeddings
slime_plugins/models/qwen3_5_vl_utils.py:22  get_packed_cp_local_indices
ValueError: Packed sequence length 4551 must be divisible by 2 * CP size 1

Reproduced on 1 node �W 8 B200 and on 2 nodes �W 8 B200, failing on the same packed
length. 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_size
chunks �~@~T hence the divisibility requirement. The check runs unconditionally, but at
cp_size == 1 one rank owns the whole stream and the correct local view is the
identity.

The fix

Short-circuit cp_size == 1 to the identity range and leave the zigzag path
untouched.

It changes no output that was previously produced: at cp_size == 1 with an even
length the existing loop already computes exactly that range (chunk = len / 2,
rank 0 takes [start, start + chunk) then [start + chunk, end)). It only stops
rejecting inputs that were always valid. Empty cu_seqlens now returns an empty
tensor instead of raising IndexError.

Tests

tests/test_qwen3_5_vl_native.py covered cp_size=2 with length-8 sequences �~@~T the
case that works. This adds test_thd_cp_indices_are_the_identity_without_context_parallelism
for the CP-1 odd-length and multi-sequence cases. The existing cp_size=2
expectations are unchanged.

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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant