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

[xinyuangui2]

Stacked on #2357 (only the last commit is new here). Context parallelism (CP > 1) for VLMs on Megatron.

• For VLMs, preprocess_packed_seqs(shard_for_cp=False) sends the full packed stream to every CP rank (each sample still padded to 2*CP). Megatron-Bridge's Qwen3VLModel computes mRoPE on the full stream, places image/deepstack features, and applies one CP partition to all of it ([training] fix Qwen3-VL packed vlm_step MRoPE NVIDIA-NeMo/Megatron-Bridge#4532). Text models keep SkyRL's pre-split, which GPTModel requires.
• The bridge's Qwen VL providers force calculate_per_token_loss=True under CP and the model asserts it at construction. SkyRL's loss returns no token count and relies on the default scaling (same as text CP), so a pre-DDP-wrap hook (keep_skyrl_loss_scaling) clears it again.
• VLM + CP + MTP raises (TODO).
• New GPU test test_megatron_vlm_cp_vs_no_cp (Qwen3-VL-2B, Qwen3.5-0.8B; 2 GPUs): CP=2 vs CP=1 per-token logprobs and grad norm.

Not yet run on GPU.

🤖 Generated with Claude Code

https://claude.ai/code/session_01MZKuFwA4VYK2tfFDVgsWt2

[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]

The keep_skyrl_loss_scaling pre-wrap hook assumes that model is always an iterable (e.g., a list or tuple of modules). However, depending on the provider and configuration, the pre-wrap hook might be invoked with a single nn.Module instead of a list. If that happens, iterating over model will raise a TypeError: '...' object is not iterable at runtime.

To make this hook robust and prevent potential crashes, we should handle both a single nn.Module and a list/tuple of modules, similar to how it was done in the previous model_packs_sequences_internally helper.

Suggested change

	def keep_skyrl_loss_scaling(model: List[nn.Module]) -> List[nn.Module]:
	"""Pre-DDP-wrap hook: undo ``calculate_per_token_loss`` forced on by Megatron-Bridge.
	True for ``Qwen3VLModel`` (e.g. Qwen3.5 via the VL bridge), which would
	double-pack and corrupt the GDN ``cu_seqlens`` under SkyRL sample packing, so
	:class:`MegatronModelWrapper` refuses packing for it. Returns ``False`` when
	mbridge / Qwen3VL is not importable, so other models are unaffected.
	The bridge's Qwen VL providers turn on ``calculate_per_token_loss`` whenever
	CP > 1, and ``Qwen3VLModel.__init__`` asserts it. The flag only changes how
	Megatron normalizes the loss: DDP then skips its 1/(dp*cp) scaling and
	``finalize_model_grads`` divides by the token count the loss function returns.
	SkyRL's loss function returns ``(loss, metrics)`` without a token count and
	relies on the default scaling, the same as for text models under CP. So clear
	the flag after the model is built and before DDP reads it.
	"""
	try:
	from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model import (
	Qwen3VLModel,
	)
	except ImportError:
	return False
	chunks = model if isinstance(model, (list, tuple)) else [model]
	for chunk in chunks:
	unwrapped = unwrap_model(chunk)
	unwrapped_list = unwrapped if isinstance(unwrapped, (list, tuple)) else [unwrapped]
	if any(isinstance(m, Qwen3VLModel) for m in unwrapped_list):
	return True
	return False
	seen = set()
	for chunk in model:
	for module in chunk.modules():
	config = getattr(module, "config", None)
	if config is not None and id(config) not in seen and getattr(config, "calculate_per_token_loss", False):
	config.calculate_per_token_loss = False
	seen.add(id(config))
	return model
	def keep_skyrl_loss_scaling(model: Union[nn.Module, List[nn.Module]]) -> Union[nn.Module, List[nn.Module]]:
	"""Pre-DDP-wrap hook: undo ``calculate_per_token_loss`` forced on by Megatron-Bridge.
	The bridge's Qwen VL providers turn on ``calculate_per_token_loss`` whenever
	CP > 1, and ``Qwen3VLModel.__init__`` asserts it. The flag only changes how
	Megatron normalizes the loss: DDP then skips its 1/(dp*cp) scaling and
	``finalize_model_grads`` divides by the token count the loss function returns.
	SkyRL's loss function returns ``(loss, metrics)`` without a token count and
	relies on the default scaling, the same as for text models under CP. So clear
	the flag after the model is built and before DDP reads it.
	"""
	chunks = model if isinstance(model, (list, tuple)) else [model]
	seen = set()
	for chunk in chunks:
	for module in chunk.modules():
	config = getattr(module, "config", None)
	if config is not None and id(config) not in seen and getattr(config, "calculate_per_token_loss", False):
	config.calculate_per_token_loss = False
	seen.add(id(config))
	return model