[megatron][vlm] Context parallelism for Qwen3-VL / Qwen3.5-VL

[xinyuangui2]

Context parallelism (CP > 1) for Qwen3-VL / Qwen3.5-VL on Megatron, on top of #2357 (VLM sequence packing).

1. Every CP rank gets the full packed stream

Qwen3-VL computes 3D mRoPE positions and places image features over the whole packed sequence, so SkyRL can't cut the sequence per CP rank first, as it does for text models. For these models preprocess_packed_seqs(shard_for_cp=False) sends the full stream to every CP rank, and Megatron-Bridge's Qwen3VLModel splits it itself:

• Given a full stream, it computes this rank's partition: model.py#L632-L647
• It places image/deepstack features, then slices the embeddings: model.py#L852-L865
• It computes mRoPE on the full stream, then slices the positions: model.py#L931-L955
• A pre-split stream is accepted only with explicit rank-local 3D position ids, which SkyRL doesn't compute: model.py#L630-L631

Text models and other VLMs keep SkyRL's per-rank split (model_owns_vlm_packing is false for them).

2. VLM + CP needs per-token loss mode

Bridge's Qwen-VL providers turn on calculate_per_token_loss whenever CP > 1 (qwen3_vl_provider.py#L137-L139, qwen35_vl_provider.py#L220-L222), and the model asserts it (model.py#L413-L416). In this mode Megatron normalizes differently:

• The schedule doesn't divide the loss by token count or microbatch count: schedules.py#L334-L347
• DDP sums gradients over DP x CP instead of averaging, and asserts average_in_collective=False (distributed_data_parallel.py#L250-L255). SkyRL defaults it to True, so the worker turns it off.
• finalize_model_grads divides every gradient once by the global token count: finalize_model_grads.py#L694-L713

CP also requires packing (trainer.remove_microbatch_padding=true): SkyRL only splits for CP on the packed path.

3. Why loss_func multiplies by the token count

SkyRL's policy loss is already normalized, since advantages are pre-scaled per mini-batch (#1296). So the gradient we want is the plain sum of every microbatch's loss over all ranks, with no further division. To cancel Megatron's division by the token count N, the worker computes each call's N (padded packed tokens per CP rank, all-reduced over DP x CP), and loss_func returns loss x N together with this rank's own count (a 3-tuple). This is the same approach as verl's Megatron engine.

• KL / entropy are per-microbatch means, so they use N / (dp x num_real_microbatches): averaged over DP and real microbatches, as in the default mode.
• MoE aux / z losses keep working: the router pre-multiplies them by its token count and relies on this same division (router.py#L617-L640).
• The default mode (no per-token loss) is unchanged.

4. Several forward_backward calls per optimizer step

SkyRL can split one optimizer step across several forward_backward calls (e.g. Tinker). Megatron's schedule passes each call's token count to finalize_model_grads (schedules.py#L863-L870); Bridge's finalize wrapper passes it on (peft/utils.py#L1095-L1114). SkyRL already defers that finalize to optim_step (_defer_finalize_model_grads, replayed by run_pending_grad_sync). If the replay divided once by the summed count, the calls would be averaged instead of added. With three calls of N = 100, 200, 300:

buffer = 100·g1 + 200·g2 + 300·g3, divided by 600 = 1/6·g1 + 1/3·g2 + 1/2·g3   (wrong)

So the replay gets num_tokens=None, and _per_token_grad_scaling divides each call's gradients by its own N, scaling the buffer by N before the call and by 1/N after it:

call	enter (xN)	+ backward	exit (/N)	buffer
1	0	+ 100·g1	/100	g1
2	200·g1	+ 200·g2	/200	g1 + g2
3	300·(g1+g2)	+ 300·g3	/300	g1 + g2 + g3

This requires overlap_grad_reduce=false, since rescaling the buffer would race with reductions still in flight.

5. MTP stays unsupported with CP

SkyRL's decoupled MTP draft loss rolls teacher logits and labels with a plain torch.roll, which isn't CP-aware, so MTP already requires CP=1 for every model. VLM + CP + MTP raises early with a clear message.

Test: GPU test_megatron_vlm_cp_vs_no_cp (Qwen3-VL-2B, Qwen3.5-0.8B; TP1 and TP2+SP) compares CP=2 with CP=1 on logprobs and grad norm.

E2E: Harbor SWE-bench Multimodal, long rows (A vs B vs C)

Qwen3.5-35B-A3B (vision tower trained), Megatron TP2/EP8, 8xH100, 64k context, fused_lm_head_logprob=true; one step of 8 mini-swe-agent rollouts on a SWE-bench Multimodal task (Harbor + skycap, #2450), rows of 27k-65k tokens. Different runs sample different rollouts, so compare distributions.

	A unpacked, CP1	B packed, CP1	C packed, CP2
logprob diff mean / p99 / max	0.0240 / 0.237 / 2.54	0.0257 / 0.239 / 2.01	0.0249 / 0.236 / 2.12
fwd logprobs	81.5 s	62.7 s	61.0 s
policy_train	141 s	132 s	131 s
policy peak mem (allocated)	48.6 GB	48.1 GB	40.4 GB
avg rollout tokens	55.6k	39.8k	57.3k

At scale (C, packed + CP2, 88 rows of a pass@8 run, avg response 50.3k tokens, 1.76M loss-masked tokens): logprob diff mean 0.0233, p99 0.229, p99.9 0.447, 0.06% of tokens > 0.5, per-row means 0.013-0.031; fwd logprobs 212 s for the 88 rows.

CP=2 with packing matches CP=1 on logprobs and cuts policy peak memory by 8 GB; per-token loss is on (Qwen3VLModel asserts it under CP>1). MTP stays off.

GPU results (4xH100)

test_megatron_vlm_cp_vs_no_cp at 9d2fb1a, CP=2 vs CP=1 at the same TP and DP=1: 4/4 passed. The batch is on-policy (old logprobs from each layout's own forward) and lr=0, so both phases see the same weights.

	logprob mean / max	grad norm CP1 -> CP2	two calls / one call (CP1, CP2)
Qwen3-VL-2B, TP1	0.034 / 0.30	4729.0 -> 4679.5 (-1.0%)	1.0026, 1.0030
Qwen3.5-0.8B, TP1	0.036 / 0.28	6894.3 -> 6942.2 (+0.7%)	1.0047, 1.0042
Qwen3-VL-2B, TP2+SP	0.033 / 0.29	4497.2 -> 4446.2 (-1.1%)	1.0016, 1.0048
Qwen3.5-0.8B, TP2+SP	0.030 / 0.34	7007.9 -> 6982.6 (-0.4%)	1.0035, 1.0027

The logprob differences come from the CP split itself (attention partitioning and reduction order). They are the same size as other bf16 kernel-path changes and below HF's own bf16-vs-fp32 gap (0.05); a wrong split would give O(1) differences. Text-model regression test_megatron_worker.py -k cp: 11 passed (on an earlier revision).

RL, Qwen3.5-9B geometry3k, 40 steps, packed 12288, with #2344 and #2357 (W&B project):

	eval pass@1 (step 0 / 20 / 40)	train reward (steps 31-40)	logprob diff mean	step time
CP=2 (TP2xCP2)	0.298 / 0.747 / 0.790	0.720	0.0063	141 s
CP=1 (TP2xDP2)	0.308 / 0.754 / 0.820	0.741	0.0061	108 s

The curves track each other (eval SE ~0.016). CP=2 costs +30% step time at these 1-2k-token lengths.

Notes from debugging the test:

• An earlier revision showed Qwen3.5 TP2 at -7.3%. That was a test artifact: fixed old logprobs (-1) put tokens near the PPO clip boundary, and bf16-level logprob differences clipped different tokens in each layout. With clipping off it was -0.34%, with every layer within +/-1.5%.
• After one real optimizer step at CP=2, the weights and reduced gradients are bit-identical across CP ranks, and every trainable parameter is in a DDP grad buffer.
• One earlier run of the old post-step two-call check showed a one-off CP=2 outlier (+16-47% on Qwen3-VL). It didn't recur in three more runs; post-step two-call norms otherwise vary about +/-4% between runs.

Text only validation

without-commit
with-commit

https://claude.ai/code/session_01MZKuFwA4VYK2tfFDVgsWt2

---

Note

High Risk
Changes distributed training loss normalization, gradient finalize behavior, and packed-sequence CP layout for VLMs—incorrect scaling or sharding would silently skew RL gradients on long multimodal sequences.

Overview
Adds Megatron context parallelism (CP > 1) for Qwen3-VL / Qwen3.5-VL with packed sequences, plus an example Geometry-3K GRPO launch script for Qwen3.5-9B (TP2, CP2, packing, fused LM-head logprobs).

For Qwen3VLModel, packed inputs no longer get pre-sharded per CP rank: preprocess_packed_seqs gains shard_for_cp=False so every CP rank sees the full THD stream; the bridge model applies mRoPE, vision placement, and the CP split internally. Text and other VLMs still use SkyRL’s zigzag CP shard.

VLM + CP turns on Megatron-Bridge’s calculate_per_token_loss. The worker attaches global/local token counts, disables DDP average_in_collective, requires overlap_grad_reduce=false, and the wrapper adjusts loss scaling (policy loss × token count; KL/entropy/MTP use a separate microbatch scale). Deferred grad finalize drops summed num_tokens and uses _per_token_grad_scaling so multiple forward_backward calls per optimizer step add gradients instead of being averaged. Config validation rejects MTP with CP > 1.

GPU parity test test_megatron_vlm_cp_vs_no_cp and unit tests cover overlap-grad-reduce and MTP+CP validation.

^{Reviewed by Cursor Bugbot for commit 6b0dd06. Bugbot is set up for automated code reviews on this repo. Configure here.}

[gemini-code-assist]

Code Review

This pull request enables and tests context parallelism (CP) and sample packing for Vision-Language Models (VLMs) such as Qwen3-VL. It introduces a shard_for_cp option in sequence preprocessing to keep the full packed stream on each CP rank, and adds a keep_skyrl_loss_scaling pre-wrap hook to preserve default loss scaling under CP. New tests are added to verify packed-vs-unpacked parity and CP-vs-no-CP parity. Feedback on the changes suggests making the keep_skyrl_loss_scaling hook more robust by handling both single nn.Module and iterable inputs to prevent potential runtime TypeErrors.

[gemini-code-assist]

Code Review

This pull request enables context parallelism (CP) and sample packing support for Vision-Language Models (VLMs) like Qwen3-VL in the Megatron backend. Key changes include adding a shard_for_cp option to sequence preprocessing to allow models to handle CP splits internally, implementing per-token loss normalization when calculate_per_token_loss is active, and tracking global/local token counts across DP and CP ranks. Additionally, the PR adds comprehensive tests to verify packed-vs-unpacked parity and CP-vs-no-CP parity for VLMs. I have no feedback to provide as there are no review comments to evaluate.

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Cursor Bugbot has reviewed your changes and found 3 potential issues.

^{Reviewed by Cursor Bugbot for commit 0c82a1e. Configure here.}

[erictang000]

thanks for updating the description with the links, very clear what the problem is and the pr makes sense in that context! just a couple comments on where we should raise errors and some extra validation

one more comment is that it would be good to add an example vlm training example script that includes information on which of these flags need to be set, like adding the flags explicitly at the top on the config here:

SkyRL/examples/train/megatron/run_megatron_dapo_qwen3.5_35b_a3b.sh

Line 63 in c2a5ce0

# Qwen3.5 flags

[xinyuangui2]

thanks for updating the description with the links, very clear what the problem is and the pr makes sense in that context! just a couple comments on where we should raise errors and some extra validation

one more comment is that it would be good to add an example vlm training example script that includes information on which of these flags need to be set, like adding the flags explicitly at the top on the config here:

SkyRL/examples/train/megatron/run_megatron_dapo_qwen3.5_35b_a3b.sh

Line 63 in c2a5ce0

# Qwen3.5 flags

Added one example script, and the run