From a4c748bf2321b4f85419241d0b4577adbdbdc85c Mon Sep 17 00:00:00 2001 From: xgui Date: Thu, 1 Oct 2026 01:38:19 +0000 Subject: [PATCH 01/10] [megatron][vlm] Allow remove_microbatch_padding for VLMs (CP=1) Megatron-Bridge's Qwen3-VL model rebuilds 3D mRoPE positions per packed sub-sequence when it receives a THD stream with position_ids=None, so the existing packed path works for VLMs. Context parallelism stays blocked. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01KuNZVD7sZm2e3tySMqTsCG Signed-off-by: xgui --- .../workers/megatron/megatron_model_wrapper.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py index bc798bee64..40a73354e3 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py @@ -299,12 +299,15 @@ def __call__(self, *args, **kwargs): return self.forward(*args, **kwargs) def _assert_vlm_supported(self): - """Guard the VLM parallelism constraints carried over from the FSDP path. - - 3D RoPE and multimodal token positions make sample/microbatch packing, - context parallelism, and sequence parallelism unsafe for VLMs today. + """Guard the VLM parallelism constraints. + + With ``remove_microbatch_padding`` the model receives a [1, T] THD stream and + Megatron-Bridge's Qwen3-VL model rebuilds 3D mRoPE positions per packed + sub-sequence from ``packed_seq_params`` (``rope.get_rope_index``). + TODO(xgui): context parallelism. preprocess_packed_seqs pre-shards the stream + per CP rank, which the bridge model only accepts with explicit rank-local + 3D position ids. """ - assert not self.remove_microbatch_padding, "VLM + microbatch padding removal unsupported" assert mpu.get_context_parallel_world_size() == 1, "VLM + context parallelism unsupported" assert ( mpu.get_tensor_model_parallel_world_size() == 1 or self.cfg.policy.sequence_parallel_size == 1 From a08b6b8fad368862ca0425dc97812b211abc9b0f Mon Sep 17 00:00:00 2001 From: xgui Date: Thu, 1 Oct 2026 05:45:09 +0000 Subject: [PATCH 02/10] [megatron][vlm] Refuse sample packing only for Qwen3VLModel with GatedDeltaNet layers The double-packing guard from #1769 matched every Qwen3VLModel, so it also blocked dense Qwen3-VL, which takes SkyRL's [1, T] THD stream without re-packing. Found by the 4xH100 A/B run (worker init ValueError). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01KuNZVD7sZm2e3tySMqTsCG Signed-off-by: xgui --- .../distributed/megatron/megatron_utils.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py b/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py index 391134ef67..cc26bbc9fa 100644 --- a/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py +++ b/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py @@ -754,15 +754,19 @@ def preprocess_packed_seqs( def model_packs_sequences_internally(model: Union[nn.Module, List[nn.Module]]) -> bool: """Whether the model packs sequences inside its own ``forward``. - 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. + True for a ``Qwen3VLModel`` with GatedDeltaNet layers (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. + Attention-only ``Qwen3VLModel`` (Qwen3-VL) takes SkyRL's [1, T] THD stream + as-is and rebuilds mRoPE per packed sub-sequence, so it is allowed. Returns + ``False`` when mbridge / Qwen3VL is not importable, so other models are + unaffected. """ try: from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model import ( Qwen3VLModel, ) + from megatron.core.ssm.gated_delta_net.gdn import GatedDeltaNet except ImportError: return False @@ -770,7 +774,10 @@ def model_packs_sequences_internally(model: Union[nn.Module, List[nn.Module]]) - 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): + if any( + isinstance(m, Qwen3VLModel) and any(isinstance(sub, GatedDeltaNet) for sub in m.modules()) + for m in unwrapped_list + ): return True return False From 1321590ab470d2866bb73568e4fdfe52656b97fd Mon Sep 17 00:00:00 2001 From: xgui Date: Sat, 3 Oct 2026 23:12:27 +0000 Subject: [PATCH 03/10] [megatron][vlm] Remove the Qwen3VLModel double-packing guard Megatron-Bridge#4532 (in the pinned 8e7077c6) makes Qwen3VLModel treat a caller-packed [1, T] THD stream as authoritative: no second compaction, and the caller's packed_seq_params reach every layer unchanged. The guard from #1769 (model_packs_sequences_internally) guarded against the pre-#4532 behaviour, so drop it, its stale GPU-test case, and comments that said the VL path cannot pack. language_model_only stays: it skips the vision tower for text-only Qwen3.5 training. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01KuNZVD7sZm2e3tySMqTsCG Signed-off-by: xgui --- examples/train/fp8/README.md | 3 +- .../distributed/megatron/megatron_utils.py | 33 +------------------ .../megatron/megatron_model_wrapper.py | 17 ++-------- .../workers/megatron/megatron_worker.py | 4 +-- .../workers/megatron/model_bridges.py | 7 ++-- .../gpu_ci/megatron/test_megatron_models.py | 7 ++-- .../gpu_ci/megatron/test_megatron_vlm_init.py | 3 -- .../megatron/test_mtp_packed_vs_unpacked.py | 7 ++-- .../megatron/test_mtp_weight_roundtrip.py | 3 +- 9 files changed, 16 insertions(+), 68 deletions(-) diff --git a/examples/train/fp8/README.md b/examples/train/fp8/README.md index 6ee2320439..4a2db41d48 100644 --- a/examples/train/fp8/README.md +++ b/examples/train/fp8/README.md @@ -30,8 +30,7 @@ Notes: `trainer.placement.policy_num_gpus_per_node` and the inference engines for a disaggregated placement. - **Qwen3.5 runs text-only.** All scripts set `language_model_only=true` on the policy, ref and - inference engine: Qwen3.5 otherwise loads through the VL bridge, which packs sequences inside - its own forward and is rejected together with SkyRL sample packing. + inference engine, so the trainer builds only the language model (no vision tower). - **GDN kernels on Blackwell.** The Blackwell scripts `export FLA_TILELANG=0` so fla uses its Triton GatedDeltaNet kernels; the TileLang packed backward aborts on B200 (it shows up as a CUDA "misaligned address" in the first backward). Leave it unset on Hopper, where the Triton diff --git a/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py b/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py index cc26bbc9fa..871a05c53e 100644 --- a/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py +++ b/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py @@ -38,7 +38,7 @@ get_moe_layer_wise_logging_tracker, reduce_aux_losses_tracker_across_ranks, ) -from megatron.core.utils import get_attr_wrapped_model, unwrap_model +from megatron.core.utils import get_attr_wrapped_model from skyrl.backends.skyrl_train.distributed.megatron.packing_utils import ( get_packed_seq_align_size, @@ -751,37 +751,6 @@ def preprocess_packed_seqs( return input_ids, packed_seq_params -def model_packs_sequences_internally(model: Union[nn.Module, List[nn.Module]]) -> bool: - """Whether the model packs sequences inside its own ``forward``. - - True for a ``Qwen3VLModel`` with GatedDeltaNet layers (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. - Attention-only ``Qwen3VLModel`` (Qwen3-VL) takes SkyRL's [1, T] THD stream - as-is and rebuilds mRoPE per packed sub-sequence, so it is allowed. Returns - ``False`` when mbridge / Qwen3VL is not importable, so other models are - unaffected. - """ - try: - from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model import ( - Qwen3VLModel, - ) - from megatron.core.ssm.gated_delta_net.gdn import GatedDeltaNet - 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) and any(isinstance(sub, GatedDeltaNet) for sub in m.modules()) - for m in unwrapped_list - ): - return True - return False - - def remove_left_padding( input_ids: torch.Tensor, attention_mask: torch.Tensor, diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py index 40a73354e3..d7b2782d46 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py @@ -19,7 +19,6 @@ from skyrl.backends.skyrl_train.distributed.megatron.megatron_utils import ( get_model_config, make_batch_generator, - model_packs_sequences_internally, preprocess_packed_seqs, recover_left_padding, remove_left_padding, @@ -220,19 +219,6 @@ def __init__( # [B, S, vocab//TP] logits + its fp32 grad). See model_utils. self._fused_lm_head = bool(getattr(self.cfg, "fused_lm_head_logprob", False)) self._fused_lm_head_backend = getattr(self.cfg, "fused_lm_head_logprob_backend", "torch") - # Some models (e.g. Qwen3.5 via the VL bridge -> Qwen3VLModel) pack - # sequences inside their own forward; SkyRL sample packing would then - # double-pack and corrupt the GDN cu_seqlens, so refuse it. For Qwen3.5, - # use language_model_only=True (native GPTModel GDN path) to pack. - if self.remove_microbatch_padding and model_packs_sequences_internally(self.actor_module): - raise ValueError( - "remove_microbatch_padding=True (sample packing) is not supported for models that " - "pack sequences inside their own forward (e.g. the Qwen3.5 VL Qwen3VLModel): it " - "double-packs and corrupts the GatedDeltaNet cu_seqlens. Set " - "trainer.policy.language_model_only=True to route Qwen3.5 to the native GPTModel GDN " - "packing path, or set trainer.remove_microbatch_padding=False." - ) - # Pending grad-sync request recorded by `_defer_finalize_model_grads`, replayed # by `run_pending_grad_sync`. See those methods for why the sync is deferred. self._pending_grad_sync: Optional[dict] = None @@ -303,7 +289,8 @@ def _assert_vlm_supported(self): With ``remove_microbatch_padding`` the model receives a [1, T] THD stream and Megatron-Bridge's Qwen3-VL model rebuilds 3D mRoPE positions per packed - sub-sequence from ``packed_seq_params`` (``rope.get_rope_index``). + sub-sequence from ``packed_seq_params`` (``rope.get_rope_index``), without + re-packing the stream (NVIDIA-NeMo/Megatron-Bridge#4532). TODO(xgui): context parallelism. preprocess_packed_seqs pre-shards the stream per CP rank, which the bridge model only accepts with explicit rank-local 3D position ids. diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index a66d6b9c3f..0a5f762b26 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -253,8 +253,8 @@ def init_configs( bridge = AutoBridge.from_hf_pretrained(bridge_source, trust_remote_code=True) # For Qwen3.5, language_model_only routes to the native GPTModel + GDN - # path (which supports sample packing) instead of the VL Qwen3VLModel - # (which doesn't). Must run before to_megatron_provider; no-op otherwise. + # path instead of the VL Qwen3VLModel, so the vision tower is not built. + # Must run before to_megatron_provider; no-op otherwise. if language_model_only and maybe_force_qwen35_text_bridge(bridge, hf_config): logger.info( "language_model_only=True: forcing Qwen3.5 text->GPTModel bridge " diff --git a/skyrl/backends/skyrl_train/workers/megatron/model_bridges.py b/skyrl/backends/skyrl_train/workers/megatron/model_bridges.py index 37f9229148..89614dfa1a 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/model_bridges.py +++ b/skyrl/backends/skyrl_train/workers/megatron/model_bridges.py @@ -95,10 +95,9 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM): # Qwen3.5 (language-model-only) -> GPTModel. # - # Qwen3.5 checkpoints dispatch to the VL bridge -> Qwen3VLModel, which packs - # sequences inside its own forward and breaks under SkyRL sample packing. When - # only the LM is wanted (language_model_only=True), route to the native - # GPTModel + GDN thd path instead. The stock text bridges assume a flat text + # Qwen3.5 checkpoints dispatch to the VL bridge -> Qwen3VLModel, which builds + # the vision tower. When only the LM is wanted (language_model_only=True), + # route to the native GPTModel + GDN path instead. The stock text bridges assume a flat text # checkpoint (top-level config, hf_prefix="model."); these subclasses adapt # them to the unified VL checkpoint (text_config, model.language_model.*), # like Qwen35VLBridge but targeting GPTModel. diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py index 119b788219..cc13ad9f3c 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py @@ -87,10 +87,9 @@ def get_test_actor_config(model_name) -> SkyRLTrainConfig: cfg.trainer.remove_microbatch_padding = False if "qwen3.5" in model_name.lower(): # Qwen3.5 hybrid GDN checkpoints report a ...ForConditionalGeneration arch - # and auto-dispatch to the VL bridge -> Qwen3VLModel, which self-packs and - # double-packs against SkyRL's sample packing (corrupting the GDN - # cu_seqlens). language_model_only routes them to the native GPTModel + GDN - # thd path instead, which supports packed sequences directly. + # and auto-dispatch to the VL bridge -> Qwen3VLModel. language_model_only + # routes them to the native GPTModel + GDN thd path, which is what this + # text-only test exercises. cfg.trainer.remove_microbatch_padding = True cfg.trainer.policy.language_model_only = True cfg.trainer.ref.language_model_only = True diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py index 80c0cabb8a..63131b7bb4 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py @@ -277,12 +277,9 @@ def run_hf_forward(batch, model_name): @pytest.mark.parametrize( ("tp", "cp", "sequence_parallel_size", "remove_microbatch_padding", "gpus_per_node", "expected_substring"), [ - # Qwen3VL packs sequences internally, `remove_microbatch_padding` is not supported - (1, 1, 1, True, 2, "pack sequences inside their own forward"), (2, 1, 2, False, 2, "sequence parallelism"), ], ids=[ - "microbatch_padding", "sequence_parallel", ], ) diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_packed_vs_unpacked.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_packed_vs_unpacked.py index f78d23919a..9cb8b60113 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_packed_vs_unpacked.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_packed_vs_unpacked.py @@ -33,9 +33,8 @@ from skyrl.train.utils.utils import validate_cfg from tests.backends.skyrl_train.gpu.utils import init_worker_with_type -# MiMo-7B-RL, not Qwen3.5: Qwen3.5's GDN layers cannot sample-pack at all (megatron-LM PR #2644), -# so its recipes run remove_microbatch_padding=false and the wrapper rejects packing outright -- -# there is no packed path to probe. MiMo is dense (Qwen2-style attention) and ships a native MTP +# MiMo-7B-RL, not Qwen3.5: the Qwen3.5 MTP recipes run remove_microbatch_padding=false, so they +# have no packed path to probe. MiMo is dense (Qwen2-style attention) and ships a native MTP # head, and its spec-decode recipe packs for real, so this exercises the path production uses. MODEL_NAME = "XiaomiMiMo/MiMo-7B-RL" @@ -123,7 +122,7 @@ def soft_ce_over_depths(main_logits, students, mask, mask_fn, vp=False): # ========================================================================= # (2) BATCHED LEFT-PADDED, non-packed (remove_left_padding) -- the ACTUAL - # Qwen3.5 training path (GDN can't pack -> REMOVE_MICROBATCH_PADDING=false). + # Qwen3.5 MTP training path (its recipes set REMOVE_MICROBATCH_PADDING=false). # ========================================================================= B = len(token_ids_list) L = max(len(x) for x in token_ids_list) diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_weight_roundtrip.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_weight_roundtrip.py index 6c2b1d9335..1ef22819b5 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_weight_roundtrip.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_weight_roundtrip.py @@ -109,8 +109,7 @@ def _make_policy_cfg(model_name: str) -> SkyRLTrainConfig: cfg.trainer.policy.megatron_config.tensor_model_parallel_size = 1 cfg.trainer.policy.megatron_config.pipeline_model_parallel_size = 1 cfg.trainer.policy.megatron_config.context_parallel_size = 1 - # This test only round-trips MTP head weights; packing is irrelevant and Qwen3.5's GDN layers - # cannot sample-pack anyway (the wrapper rejects it -- see the 9B recipe's remove_microbatch_padding=false). + # This test only round-trips MTP head weights; packing is irrelevant. cfg.trainer.remove_microbatch_padding = False # Skip optimizer init for weight sync test cfg.trainer.policy.inference_only_init = True From dc535ae437940ee200c3477cedf198cd03040b41 Mon Sep 17 00:00:00 2001 From: xgui Date: Sun, 4 Oct 2026 17:25:45 +0000 Subject: [PATCH 04/10] [megatron][vlm] Add packed-vs-unpacked per-token logprob test for Qwen3-VL and Qwen3.5-VL Variable-length image samples plus one text-only row, microbatches of 4 and a trailing single-sample microbatch, TP=1 and TP=2+SP. Compares every scored token and splits the stats by slot in the packed microbatch, so a sample-boundary leak (mRoPE restart, or GatedDeltaNet state/conv carried over) shows up as later slots being worse than slot 0. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MZKuFwA4VYK2tfFDVgsWt2 Signed-off-by: xgui --- .../gpu_ci/megatron/test_megatron_vlm_init.py | 150 ++++++++++++++++++ 1 file changed, 150 insertions(+) diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py index 63131b7bb4..f2e7e414de 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py @@ -359,3 +359,153 @@ async def test_vlm_train(ray_init_fixture, worker_type, tp, pp, gpus_per_node): training_losses.append(step_i_outputs["loss"]) assert training_losses[0] > training_losses[-1] + + +# Packed-vs-unpacked parity. Microbatches of 4 samples (plus a trailing single-sample +# microbatch), so every packed row has samples after the first: a sample-boundary leak +# (wrong mRoPE restart, or GDN state / conv carried across samples) can only show up there. +PACKING_MICRO_BATCH = 4 +PACKING_PROMPTS = [ + ("Describe this picture in one short sentence.", (56, 56), "A small square of noise."), + ("What shapes can you see here? Answer briefly.", (112, 84), "Mostly scattered dots with no clear shape."), + ("Is this image bright or dark?", (224, 168), "It is a mix of bright and dark pixels."), + ("Count the colors.", (84, 224), "Too many colors to count; it looks like random noise."), + ("Give a title for this image.", (140, 140), "Static."), + ("Without any image: what is two plus three?", None, "Two plus three is five."), + ("What would you call this texture?", (196, 112), "A grainy, television-static texture."), + ("Summarize the image.", (70, 154), "Random colored noise in a tall rectangle."), + ("One word for this image?", (168, 56), "Noise."), +] + + +def get_packing_parity_batch(model_name: str) -> TrainingInputBatch: + """Variable-length image (and one text-only) samples, left-padded, scoring every token.""" + processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True) + gen = torch.Generator().manual_seed(0) + rows = [] + for prompt, size, answer in PACKING_PROMPTS: + user_content = [{"type": "text", "text": prompt}] + images = None + if size is not None: + from PIL import Image + + w, h = size + pixels = torch.randint(0, 256, (h, w, 3), generator=gen, dtype=torch.uint8).numpy() + images = [Image.fromarray(pixels)] + user_content = [{"type": "image"}] + user_content + messages = [ + {"role": "user", "content": user_content}, + {"role": "assistant", "content": answer}, + ] + text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=False) + out = processor(text=[text], images=images, return_tensors="pt") + rows.append((out["input_ids"][0].tolist(), out.get("pixel_values"), out.get("image_grid_thw"))) + + pixel_dim = next(pv.shape[-1] for _, pv, _ in rows if pv is not None) + max_len = max(len(ids) for ids, _, _ in rows) + pad_token_id = processor.tokenizer.pad_token_id or processor.tokenizer.eos_token_id + sequences, attention_mask, pixel_values, image_grid_thw = [], [], [], [] + for ids, pv, grid in rows: + pad = max_len - len(ids) + sequences.append([pad_token_id] * pad + ids) + attention_mask.append([0] * pad + [1] * len(ids)) + # Text-only rows carry empty vision tensors. + pixel_values.append(pv if pv is not None else torch.zeros(0, pixel_dim)) + image_grid_thw.append(grid if grid is not None else torch.zeros(0, 3, dtype=torch.long)) + + batch_size = len(rows) + num_actions = max_len - 1 # score every position that has a predecessor + loss_mask = torch.tensor(attention_mask)[:, -num_actions:].float() + zeros = torch.zeros(batch_size, num_actions) + data = TrainingInputBatch( + { + "sequences": torch.tensor(sequences), + "attention_mask": torch.tensor(attention_mask), + "action_log_probs": zeros, + "base_action_log_probs": zeros, + "rollout_logprobs": zeros, + "values": zeros, + "returns": zeros, + "advantages": zeros, + "loss_mask": loss_mask, + "response_mask": loss_mask, + "pixel_values": TensorList(pixel_values), + "image_grid_thw": TensorList(image_grid_thw), + } + ) + data.metadata = {"response_length": num_actions} + return data + + +def _megatron_vlm_forward_logprobs(model_name, batch, tp, remove_microbatch_padding) -> torch.Tensor: + cfg = get_test_actor_config(model_name=model_name) + cfg.trainer.strategy = "megatron" + cfg.trainer.placement.policy_num_gpus_per_node = tp + cfg.trainer.policy.megatron_config.tensor_model_parallel_size = tp + cfg.trainer.policy.megatron_config.pipeline_model_parallel_size = 1 + cfg.trainer.micro_forward_batch_size_per_gpu = PACKING_MICRO_BATCH + cfg.trainer.micro_train_batch_size_per_gpu = PACKING_MICRO_BATCH + cfg.trainer.remove_microbatch_padding = remove_microbatch_padding + try: + actor_group = init_worker_with_type("policy", shared_pg=None, colocate_all=False, num_gpus_per_node=tp, cfg=cfg) + refs = actor_group.async_run_ray_method("mesh", "forward", data=batch, loss_fn="cross_entropy") + output = WorkerOutput.cat(actor_group.actor_infos, ray.get(refs)) + return loss_fn_outputs_to_tensor(output.loss_fn_outputs, key="logprobs").float() + finally: + ray.shutdown() + ray_init_for_tests() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model_name", "tp"), + [ + ("Qwen/Qwen3-VL-2B-Instruct", 1), + ("Qwen/Qwen3-VL-2B-Instruct", 2), + # Qwen3.5 goes through the same Qwen3VLModel but with GatedDeltaNet layers, whose + # state and conv must reset at every packed sample boundary. + ("Qwen/Qwen3.5-0.8B", 1), + ("Qwen/Qwen3.5-0.8B", 2), + ], + ids=["qwen3_vl_tp1", "qwen3_vl_tp2_sp", "qwen3_5_vl_tp1", "qwen3_5_vl_tp2_sp"], +) +@pytest.mark.megatron +async def test_megatron_vlm_packed_vs_unpacked(ray_init_fixture, model_name, tp): + """Sample packing must not change per-token logprobs of image+text samples. + + Compares every scored token of each sample, split by the sample's slot in its + packed microbatch: slot 0 differs only by kernel numerics (nothing precedes it); + later slots would additionally show any leak across the sample boundary. + """ + batch = get_packing_parity_batch(model_name) + unpacked = _megatron_vlm_forward_logprobs(model_name, batch, tp, remove_microbatch_padding=False) + packed = _megatron_vlm_forward_logprobs(model_name, batch, tp, remove_microbatch_padding=True) + assert packed.shape == unpacked.shape + + num_actions = batch.metadata["response_length"] + am = batch["attention_mask"].bool() + # A token is scored only if it and its predecessor are both real (the first token of + # each sample is predicted from padding unpacked, but from the previous sample packed). + scored = (am[:, 1:] & am[:, :-1])[:, -num_actions:] + diff = (packed - unpacked).abs() + assert torch.isfinite(packed[scored]).all() and torch.isfinite(unpacked[scored]).all() + + slots = torch.arange(len(PACKING_PROMPTS)) % PACKING_MICRO_BATCH + stats = {} + for name, rows in (("first", slots == 0), ("later", slots > 0)): + d = diff[rows][scored[rows]] + stats[name] = (d.max().item(), d.mean().item(), d.numel()) + print(f"\n[packing parity] {model_name} tp={tp}") + for name, (mx, mean, n) in stats.items(): + print(f" {name:5s} samples: max={mx:.4f} mean={mean:.5f} tokens={n}") + for i in range(len(PACKING_PROMPTS)): + d = diff[i][scored[i]] + print(f" sample {i} slot {slots[i].item()}: max={d.max().item():.4f} mean={d.mean().item():.5f}") + + first_max, first_mean, _ = stats["first"] + later_max, later_mean, _ = stats["later"] + # Overall agreement, same bar as test_megatron_forward's average-diff check. + assert diff[scored].mean().item() < 2e-2, stats + # No boundary leak: later samples must not be materially worse than first samples. + assert later_mean <= 3 * first_mean + 5e-3, stats + assert later_max <= 3 * first_max + 0.25, stats From ba4ad23ceaa5eb2b5e804353814b586df8888193 Mon Sep 17 00:00:00 2001 From: xgui Date: Sun, 4 Oct 2026 21:06:38 +0000 Subject: [PATCH 05/10] [megatron][vlm] Fix packed-vs-unpacked test alignment; score answer tokens only The test read logprobs through the forward(loss_fn=cross_entropy) path, which returns each sample's values compacted to its loss_mask and left-aligned, but indexed them as right-aligned, so it compared the wrong positions (and zeros). Use the inference forward path (the RL old/ref-logprob path, right-aligned [B, response_length]) and score only each row's assistant answer, as RL does: image-placeholder targets of random-noise images have huge, bf16-sensitive logprobs that are never trained on. Found by the 4xH100 run. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MZKuFwA4VYK2tfFDVgsWt2 Signed-off-by: xgui --- .../gpu_ci/megatron/test_megatron_vlm_init.py | 59 +++++++++++-------- 1 file changed, 36 insertions(+), 23 deletions(-) diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py index f2e7e414de..b2bc47e1f9 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py @@ -379,7 +379,13 @@ async def test_vlm_train(ray_init_fixture, worker_type, tp, pp, gpus_per_node): def get_packing_parity_batch(model_name: str) -> TrainingInputBatch: - """Variable-length image (and one text-only) samples, left-padded, scoring every token.""" + """Variable-length image (and one text-only) samples laid out like an RL batch. + + Each row is prompt + assistant answer, left-padded. ``response_length`` is the + longest answer and ``loss_mask`` marks each row's answer tokens (right-aligned), + so only text targets are scored -- image-placeholder targets of random-noise + images have huge, bf16-sensitive logprobs and are never trained on. + """ processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True) gen = torch.Generator().manual_seed(0) rows = [] @@ -398,24 +404,32 @@ def get_packing_parity_batch(model_name: str) -> TrainingInputBatch: {"role": "assistant", "content": answer}, ] text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=False) + prompt_text = processor.apply_chat_template(messages[:1], tokenize=False, add_generation_prompt=True) out = processor(text=[text], images=images, return_tensors="pt") - rows.append((out["input_ids"][0].tolist(), out.get("pixel_values"), out.get("image_grid_thw"))) - - pixel_dim = next(pv.shape[-1] for _, pv, _ in rows if pv is not None) - max_len = max(len(ids) for ids, _, _ in rows) + prompt_ids = processor(text=[prompt_text], images=images, return_tensors="pt")["input_ids"][0].tolist() + ids = out["input_ids"][0].tolist() + # Score what follows the shared prompt prefix (some templates render the + # generation prompt slightly differently from a completed turn, e.g. think tags). + prefix = next((k for k, (a, b) in enumerate(zip(ids, prompt_ids)) if a != b), len(prompt_ids)) + assert 0 < len(ids) - prefix <= len(ids) // 2, (prefix, len(ids)) + rows.append((ids, len(ids) - prefix, out.get("pixel_values"), out.get("image_grid_thw"))) + + pixel_dim = next(pv.shape[-1] for _, _, pv, _ in rows if pv is not None) + max_len = max(len(ids) for ids, _, _, _ in rows) + num_actions = max(resp_len for _, resp_len, _, _ in rows) pad_token_id = processor.tokenizer.pad_token_id or processor.tokenizer.eos_token_id - sequences, attention_mask, pixel_values, image_grid_thw = [], [], [], [] - for ids, pv, grid in rows: + sequences, attention_mask, loss_mask, pixel_values, image_grid_thw = [], [], [], [], [] + for ids, resp_len, pv, grid in rows: pad = max_len - len(ids) sequences.append([pad_token_id] * pad + ids) attention_mask.append([0] * pad + [1] * len(ids)) + loss_mask.append([0] * (num_actions - resp_len) + [1] * resp_len) # Text-only rows carry empty vision tensors. pixel_values.append(pv if pv is not None else torch.zeros(0, pixel_dim)) image_grid_thw.append(grid if grid is not None else torch.zeros(0, 3, dtype=torch.long)) batch_size = len(rows) - num_actions = max_len - 1 # score every position that has a predecessor - loss_mask = torch.tensor(attention_mask)[:, -num_actions:].float() + loss_mask = torch.tensor(loss_mask, dtype=torch.float) zeros = torch.zeros(batch_size, num_actions) data = TrainingInputBatch( { @@ -438,6 +452,7 @@ def get_packing_parity_batch(model_name: str) -> TrainingInputBatch: def _megatron_vlm_forward_logprobs(model_name, batch, tp, remove_microbatch_padding) -> torch.Tensor: + """Inference forward (the RL old/ref-logprob path): [B, response_length], right-aligned.""" cfg = get_test_actor_config(model_name=model_name) cfg.trainer.strategy = "megatron" cfg.trainer.placement.policy_num_gpus_per_node = tp @@ -448,7 +463,7 @@ def _megatron_vlm_forward_logprobs(model_name, batch, tp, remove_microbatch_padd cfg.trainer.remove_microbatch_padding = remove_microbatch_padding try: actor_group = init_worker_with_type("policy", shared_pg=None, colocate_all=False, num_gpus_per_node=tp, cfg=cfg) - refs = actor_group.async_run_ray_method("mesh", "forward", data=batch, loss_fn="cross_entropy") + refs = actor_group.async_run_ray_method("mesh", "forward", data=batch) output = WorkerOutput.cat(actor_group.actor_infos, ray.get(refs)) return loss_fn_outputs_to_tensor(output.loss_fn_outputs, key="logprobs").float() finally: @@ -471,22 +486,20 @@ def _megatron_vlm_forward_logprobs(model_name, batch, tp, remove_microbatch_padd ) @pytest.mark.megatron async def test_megatron_vlm_packed_vs_unpacked(ray_init_fixture, model_name, tp): - """Sample packing must not change per-token logprobs of image+text samples. + """Sample packing must not change per-token logprobs of the answer tokens. - Compares every scored token of each sample, split by the sample's slot in its - packed microbatch: slot 0 differs only by kernel numerics (nothing precedes it); - later slots would additionally show any leak across the sample boundary. + Stats are split by each sample's slot in its packed microbatch: a leak across a + sample boundary can only make later slots worse than slot 0. Packed and unpacked + run different kernels, so they agree only to bf16 precision (and a whole + microbatch can shift together when the kernel path changes). """ batch = get_packing_parity_batch(model_name) unpacked = _megatron_vlm_forward_logprobs(model_name, batch, tp, remove_microbatch_padding=False) packed = _megatron_vlm_forward_logprobs(model_name, batch, tp, remove_microbatch_padding=True) - assert packed.shape == unpacked.shape - num_actions = batch.metadata["response_length"] - am = batch["attention_mask"].bool() - # A token is scored only if it and its predecessor are both real (the first token of - # each sample is predicted from padding unpacked, but from the previous sample packed). - scored = (am[:, 1:] & am[:, :-1])[:, -num_actions:] + assert packed.shape == unpacked.shape == (len(PACKING_PROMPTS), num_actions) + + scored = batch["loss_mask"].bool() diff = (packed - unpacked).abs() assert torch.isfinite(packed[scored]).all() and torch.isfinite(unpacked[scored]).all() @@ -504,8 +517,8 @@ async def test_megatron_vlm_packed_vs_unpacked(ray_init_fixture, model_name, tp) first_max, first_mean, _ = stats["first"] later_max, later_mean, _ = stats["later"] - # Overall agreement, same bar as test_megatron_forward's average-diff check. - assert diff[scored].mean().item() < 2e-2, stats + # bf16 agreement on answer tokens (HF bf16 vs fp32 differs by ~5e-2 mean on these). + assert diff[scored].mean().item() < 3e-2, stats # No boundary leak: later samples must not be materially worse than first samples. - assert later_mean <= 3 * first_mean + 5e-3, stats + assert later_mean <= 3 * first_mean + 1e-2, stats assert later_max <= 3 * first_max + 0.25, stats From 094a8ef60728c6e4c0c2fffb4036ba4256578e51 Mon Sep 17 00:00:00 2001 From: xgui Date: Mon, 5 Oct 2026 04:36:15 +0000 Subject: [PATCH 06/10] [megatron][vlm] Packing parity test: compare against each sample run alone Packed and unpacked microbatches agree only to bf16 rounding, and unpacked microbatches move by the same amount when only their composition changes (measured on 4xH100: 0.034 mean on answer tokens at TP2+SP for both). A fixed packed-vs-unpacked bar therefore tests batching noise. Use each sample run alone (unpacked, micro=1) as the reference and require packed microbatches to be no further from it than unpacked ones; keep the slot-based leak check and a loose absolute guard. At TP>1 the text-only row is left out of the reference: alone it hits a Megatron-Bridge split_deepstack_embs crash on image-free microbatches. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MZKuFwA4VYK2tfFDVgsWt2 Signed-off-by: xgui --- .../gpu_ci/megatron/test_megatron_vlm_init.py | 91 ++++++++++++------- 1 file changed, 59 insertions(+), 32 deletions(-) diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py index b2bc47e1f9..647888ce72 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py @@ -451,15 +451,17 @@ def get_packing_parity_batch(model_name: str) -> TrainingInputBatch: return data -def _megatron_vlm_forward_logprobs(model_name, batch, tp, remove_microbatch_padding) -> torch.Tensor: +def _megatron_vlm_forward_logprobs( + model_name, batch, tp, remove_microbatch_padding, micro_batch=PACKING_MICRO_BATCH +) -> torch.Tensor: """Inference forward (the RL old/ref-logprob path): [B, response_length], right-aligned.""" cfg = get_test_actor_config(model_name=model_name) cfg.trainer.strategy = "megatron" cfg.trainer.placement.policy_num_gpus_per_node = tp cfg.trainer.policy.megatron_config.tensor_model_parallel_size = tp cfg.trainer.policy.megatron_config.pipeline_model_parallel_size = 1 - cfg.trainer.micro_forward_batch_size_per_gpu = PACKING_MICRO_BATCH - cfg.trainer.micro_train_batch_size_per_gpu = PACKING_MICRO_BATCH + cfg.trainer.micro_forward_batch_size_per_gpu = micro_batch + cfg.trainer.micro_train_batch_size_per_gpu = micro_batch cfg.trainer.remove_microbatch_padding = remove_microbatch_padding try: actor_group = init_worker_with_type("policy", shared_pg=None, colocate_all=False, num_gpus_per_node=tp, cfg=cfg) @@ -486,39 +488,64 @@ def _megatron_vlm_forward_logprobs(model_name, batch, tp, remove_microbatch_padd ) @pytest.mark.megatron async def test_megatron_vlm_packed_vs_unpacked(ray_init_fixture, model_name, tp): - """Sample packing must not change per-token logprobs of the answer tokens. - - Stats are split by each sample's slot in its packed microbatch: a leak across a - sample boundary can only make later slots worse than slot 0. Packed and unpacked - run different kernels, so they agree only to bf16 precision (and a whole - microbatch can shift together when the kernel path changes). + """Sample packing must add no error beyond ordinary batching. + + Packed and padded (unpacked) microbatches run different kernels on different + shapes, so in bf16 they agree only to rounding -- and so does unpacked with + itself when the microbatch composition changes. The reference is therefore + each sample run alone (unpacked, one sample per microbatch), and the check is + that packed microbatches of 4 are no further from it than unpacked ones. + A sample-boundary leak (mRoPE restart, GatedDeltaNet state/conv carried + across samples) would also make later slots of a packed microbatch worse + than slot 0. """ batch = get_packing_parity_batch(model_name) + num_actions = batch.metadata["response_length"] unpacked = _megatron_vlm_forward_logprobs(model_name, batch, tp, remove_microbatch_padding=False) packed = _megatron_vlm_forward_logprobs(model_name, batch, tp, remove_microbatch_padding=True) - num_actions = batch.metadata["response_length"] assert packed.shape == unpacked.shape == (len(PACKING_PROMPTS), num_actions) + # Reference: every sample alone. At TP>1 (sequence parallel) a microbatch with no image + # crashes Megatron-Bridge's split_deepstack_embs (fixed upstream in Megatron-Bridge#3869, + # not yet merged), so the text-only row is left out of the reference there. + has_image = torch.tensor([size is not None for _, size, _ in PACKING_PROMPTS]) + ref_rows = torch.arange(len(PACKING_PROMPTS)) if tp == 1 else has_image.nonzero().flatten() + ref_batch = TrainingInputBatch({k: (None if v is None else v[ref_rows]) for k, v in batch.items()}) + ref_batch.metadata = batch.metadata + alone = _megatron_vlm_forward_logprobs(model_name, ref_batch, tp, remove_microbatch_padding=False, micro_batch=1) + scored = batch["loss_mask"].bool() - diff = (packed - unpacked).abs() - assert torch.isfinite(packed[scored]).all() and torch.isfinite(unpacked[scored]).all() - - slots = torch.arange(len(PACKING_PROMPTS)) % PACKING_MICRO_BATCH - stats = {} - for name, rows in (("first", slots == 0), ("later", slots > 0)): - d = diff[rows][scored[rows]] - stats[name] = (d.max().item(), d.mean().item(), d.numel()) - print(f"\n[packing parity] {model_name} tp={tp}") - for name, (mx, mean, n) in stats.items(): - print(f" {name:5s} samples: max={mx:.4f} mean={mean:.5f} tokens={n}") - for i in range(len(PACKING_PROMPTS)): - d = diff[i][scored[i]] - print(f" sample {i} slot {slots[i].item()}: max={d.max().item():.4f} mean={d.mean().item():.5f}") - - first_max, first_mean, _ = stats["first"] - later_max, later_mean, _ = stats["later"] - # bf16 agreement on answer tokens (HF bf16 vs fp32 differs by ~5e-2 mean on these). - assert diff[scored].mean().item() < 3e-2, stats - # No boundary leak: later samples must not be materially worse than first samples. - assert later_mean <= 3 * first_mean + 1e-2, stats - assert later_max <= 3 * first_max + 0.25, stats + for name, t in (("unpacked", unpacked), ("packed", packed), ("alone", alone)): + rows_scored = scored[ref_rows] if name == "alone" else scored + assert torch.isfinite(t[rows_scored]).all(), name + ref_scored = scored[ref_rows] + packed_vs_alone = (packed[ref_rows] - alone).abs() + unpacked_vs_alone = (unpacked[ref_rows] - alone).abs() + packed_vs_unpacked = (packed - unpacked).abs() + + slots = (torch.arange(len(PACKING_PROMPTS)) % PACKING_MICRO_BATCH)[ref_rows] + print(f"\n[packing parity] {model_name} tp={tp} (mean/max abs logprob diff on answer tokens)") + for name, d, m in ( + ("packed vs alone", packed_vs_alone, ref_scored), + ("unpacked vs alone", unpacked_vs_alone, ref_scored), + ("packed vs unpacked", packed_vs_unpacked, scored), + ): + print(f" {name:18s}: mean={d[m].mean().item():.5f} max={d[m].max().item():.4f}") + for j, i in enumerate(ref_rows.tolist()): + m = ref_scored[j] + print( + f" sample {i} slot {slots[j].item()}: packed-alone={packed_vs_alone[j][m].mean().item():.5f} " + f"unpacked-alone={unpacked_vs_alone[j][m].mean().item():.5f}" + ) + + packed_err = packed_vs_alone[ref_scored].mean().item() + unpacked_err = unpacked_vs_alone[ref_scored].mean().item() + # Packing must not add error beyond what ordinary batching already shows. + assert packed_err <= 1.5 * unpacked_err + 5e-3, (packed_err, unpacked_err) + # Guard against gross breakage (a wrong position or boundary is >> bf16 noise). + assert packed_vs_unpacked[scored].mean().item() < 5e-2 + # No boundary leak: later slots of a packed microbatch are not worse than slot 0. + first, later = slots == 0, slots > 0 + first_d, later_d = packed_vs_alone[first][ref_scored[first]], packed_vs_alone[later][ref_scored[later]] + assert later_d.mean().item() <= 3 * first_d.mean().item() + 1e-2, (later_d.mean(), first_d.mean()) + assert later_d.max().item() <= 3 * first_d.max().item() + 0.25, (later_d.max(), first_d.max()) From a20c0a9f4793b51b9a3b169b3fa76a5a8236c9f1 Mon Sep 17 00:00:00 2001 From: xgui Date: Mon, 5 Oct 2026 22:39:02 +0000 Subject: [PATCH 07/10] [megatron][vlm] Packing parity test: add PP=2 cases Pixel values reach only the first pipeline stage while every stage rebuilds the packed mRoPE positions; cover Qwen3-VL-2B and Qwen3.5-0.8B at PP=2. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MZKuFwA4VYK2tfFDVgsWt2 Signed-off-by: xgui --- .../gpu_ci/megatron/test_megatron_vlm_init.py | 37 +++++++++++-------- 1 file changed, 22 insertions(+), 15 deletions(-) diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py index 647888ce72..d1d9fa2b2b 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py @@ -452,19 +452,21 @@ def get_packing_parity_batch(model_name: str) -> TrainingInputBatch: def _megatron_vlm_forward_logprobs( - model_name, batch, tp, remove_microbatch_padding, micro_batch=PACKING_MICRO_BATCH + model_name, batch, tp, remove_microbatch_padding, micro_batch=PACKING_MICRO_BATCH, pp=1 ) -> torch.Tensor: """Inference forward (the RL old/ref-logprob path): [B, response_length], right-aligned.""" cfg = get_test_actor_config(model_name=model_name) cfg.trainer.strategy = "megatron" - cfg.trainer.placement.policy_num_gpus_per_node = tp + cfg.trainer.placement.policy_num_gpus_per_node = tp * pp cfg.trainer.policy.megatron_config.tensor_model_parallel_size = tp - cfg.trainer.policy.megatron_config.pipeline_model_parallel_size = 1 + cfg.trainer.policy.megatron_config.pipeline_model_parallel_size = pp cfg.trainer.micro_forward_batch_size_per_gpu = micro_batch cfg.trainer.micro_train_batch_size_per_gpu = micro_batch cfg.trainer.remove_microbatch_padding = remove_microbatch_padding try: - actor_group = init_worker_with_type("policy", shared_pg=None, colocate_all=False, num_gpus_per_node=tp, cfg=cfg) + actor_group = init_worker_with_type( + "policy", shared_pg=None, colocate_all=False, num_gpus_per_node=tp * pp, cfg=cfg + ) refs = actor_group.async_run_ray_method("mesh", "forward", data=batch) output = WorkerOutput.cat(actor_group.actor_infos, ray.get(refs)) return loss_fn_outputs_to_tensor(output.loss_fn_outputs, key="logprobs").float() @@ -475,19 +477,22 @@ def _megatron_vlm_forward_logprobs( @pytest.mark.asyncio @pytest.mark.parametrize( - ("model_name", "tp"), + ("model_name", "tp", "pp"), [ - ("Qwen/Qwen3-VL-2B-Instruct", 1), - ("Qwen/Qwen3-VL-2B-Instruct", 2), + ("Qwen/Qwen3-VL-2B-Instruct", 1, 1), + ("Qwen/Qwen3-VL-2B-Instruct", 2, 1), + # PP=2: pixel values reach only the first stage; every stage rebuilds mRoPE positions. + ("Qwen/Qwen3-VL-2B-Instruct", 1, 2), # Qwen3.5 goes through the same Qwen3VLModel but with GatedDeltaNet layers, whose # state and conv must reset at every packed sample boundary. - ("Qwen/Qwen3.5-0.8B", 1), - ("Qwen/Qwen3.5-0.8B", 2), + ("Qwen/Qwen3.5-0.8B", 1, 1), + ("Qwen/Qwen3.5-0.8B", 2, 1), + ("Qwen/Qwen3.5-0.8B", 1, 2), ], - ids=["qwen3_vl_tp1", "qwen3_vl_tp2_sp", "qwen3_5_vl_tp1", "qwen3_5_vl_tp2_sp"], + ids=["qwen3_vl_tp1", "qwen3_vl_tp2_sp", "qwen3_vl_pp2", "qwen3_5_vl_tp1", "qwen3_5_vl_tp2_sp", "qwen3_5_vl_pp2"], ) @pytest.mark.megatron -async def test_megatron_vlm_packed_vs_unpacked(ray_init_fixture, model_name, tp): +async def test_megatron_vlm_packed_vs_unpacked(ray_init_fixture, model_name, tp, pp): """Sample packing must add no error beyond ordinary batching. Packed and padded (unpacked) microbatches run different kernels on different @@ -501,8 +506,8 @@ async def test_megatron_vlm_packed_vs_unpacked(ray_init_fixture, model_name, tp) """ batch = get_packing_parity_batch(model_name) num_actions = batch.metadata["response_length"] - unpacked = _megatron_vlm_forward_logprobs(model_name, batch, tp, remove_microbatch_padding=False) - packed = _megatron_vlm_forward_logprobs(model_name, batch, tp, remove_microbatch_padding=True) + unpacked = _megatron_vlm_forward_logprobs(model_name, batch, tp, remove_microbatch_padding=False, pp=pp) + packed = _megatron_vlm_forward_logprobs(model_name, batch, tp, remove_microbatch_padding=True, pp=pp) assert packed.shape == unpacked.shape == (len(PACKING_PROMPTS), num_actions) # Reference: every sample alone. At TP>1 (sequence parallel) a microbatch with no image @@ -512,7 +517,9 @@ async def test_megatron_vlm_packed_vs_unpacked(ray_init_fixture, model_name, tp) ref_rows = torch.arange(len(PACKING_PROMPTS)) if tp == 1 else has_image.nonzero().flatten() ref_batch = TrainingInputBatch({k: (None if v is None else v[ref_rows]) for k, v in batch.items()}) ref_batch.metadata = batch.metadata - alone = _megatron_vlm_forward_logprobs(model_name, ref_batch, tp, remove_microbatch_padding=False, micro_batch=1) + alone = _megatron_vlm_forward_logprobs( + model_name, ref_batch, tp, remove_microbatch_padding=False, micro_batch=1, pp=pp + ) scored = batch["loss_mask"].bool() for name, t in (("unpacked", unpacked), ("packed", packed), ("alone", alone)): @@ -524,7 +531,7 @@ async def test_megatron_vlm_packed_vs_unpacked(ray_init_fixture, model_name, tp) packed_vs_unpacked = (packed - unpacked).abs() slots = (torch.arange(len(PACKING_PROMPTS)) % PACKING_MICRO_BATCH)[ref_rows] - print(f"\n[packing parity] {model_name} tp={tp} (mean/max abs logprob diff on answer tokens)") + print(f"\n[packing parity] {model_name} tp={tp} pp={pp} (mean/max abs logprob diff on answer tokens)") for name, d, m in ( ("packed vs alone", packed_vs_alone, ref_scored), ("unpacked vs alone", unpacked_vs_alone, ref_scored), From 10a1d7787e6e12a4cf3774f6eeb6681d7932a4de Mon Sep 17 00:00:00 2001 From: xgui Date: Mon, 5 Oct 2026 23:07:12 +0000 Subject: [PATCH 08/10] [megatron][vlm] Gate VLM sample packing on models that pack themselves Sample packing for VLMs relies on the model rebuilding 3D mRoPE positions from the packed stream. Allow it only when every model chunk is Megatron-Bridge's Qwen3VLModel (Qwen3-VL, Qwen3.5-VL) or sets model_owns_packing=True (NeMo-RL's opt-in attribute), and raise a clear error for other VLMs instead of running an untested path. Unpacked VLM training is unchanged. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MZKuFwA4VYK2tfFDVgsWt2 Signed-off-by: xgui --- .../distributed/megatron/megatron_utils.py | 28 ++++++++++++++- .../megatron/megatron_model_wrapper.py | 15 +++++++- .../test_model_owns_vlm_packing.py | 35 +++++++++++++++++++ 3 files changed, 76 insertions(+), 2 deletions(-) create mode 100644 tests/backends/skyrl_train/distributed/test_model_owns_vlm_packing.py diff --git a/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py b/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py index 871a05c53e..39e08d812b 100644 --- a/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py +++ b/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py @@ -38,7 +38,7 @@ get_moe_layer_wise_logging_tracker, reduce_aux_losses_tracker_across_ranks, ) -from megatron.core.utils import get_attr_wrapped_model +from megatron.core.utils import get_attr_wrapped_model, unwrap_model from skyrl.backends.skyrl_train.distributed.megatron.packing_utils import ( get_packed_seq_align_size, @@ -751,6 +751,32 @@ def preprocess_packed_seqs( return input_ids, packed_seq_params +def model_owns_vlm_packing(model: Union[nn.Module, List[nn.Module]]) -> bool: + """Whether the VLM handles a packed [1, T] stream itself. + + True when every model chunk is Megatron-Bridge's ``Qwen3VLModel`` (Qwen3-VL, + Qwen3.5-VL), which rebuilds 3D mRoPE positions per packed sub-sequence from + ``packed_seq_params``, or sets ``model_owns_packing = True`` (NeMo-RL's opt-in + attribute for models that pack and split for context parallelism themselves). + """ + try: + from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model import ( + Qwen3VLModel, + ) + except ImportError: + Qwen3VLModel = None + + chunks = model if isinstance(model, (list, tuple)) else [model] + for chunk in chunks: + unwrapped = unwrap_model(chunk) + if getattr(unwrapped, "model_owns_packing", False): + continue + if Qwen3VLModel is not None and isinstance(unwrapped, Qwen3VLModel): + continue + return False + return bool(chunks) + + def remove_left_padding( input_ids: torch.Tensor, attention_mask: torch.Tensor, diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py index d7b2782d46..02779eb32b 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py @@ -10,6 +10,7 @@ finalize_model_grads_with_expert_adapter_sync, ) from megatron.core.pipeline_parallel import get_forward_backward_func +from megatron.core.utils import unwrap_model from omegaconf import OmegaConf from skyrl.backends.skyrl_train.distributed.megatron.fused_lm_head import ( @@ -19,6 +20,7 @@ from skyrl.backends.skyrl_train.distributed.megatron.megatron_utils import ( get_model_config, make_batch_generator, + model_owns_vlm_packing, preprocess_packed_seqs, recover_left_padding, remove_left_padding, @@ -214,6 +216,9 @@ def __init__( self.policy_loss_fn = policy_loss_fn self.remove_microbatch_padding = self.cfg.remove_microbatch_padding self.is_vlm = is_vlm + # Sample packing for VLMs needs a model that computes mRoPE positions from the packed + # stream itself; other VLMs only run unpacked (see _assert_vlm_supported). + self.model_owns_vlm_packing = is_vlm and model_owns_vlm_packing(self.actor_module) # Fuse the LM-head projection into the chunked log-prob/entropy via the # GPTModel output_processor hook (avoids materializing the full # [B, S, vocab//TP] logits + its fp32 grad). See model_utils. @@ -290,11 +295,19 @@ def _assert_vlm_supported(self): With ``remove_microbatch_padding`` the model receives a [1, T] THD stream and Megatron-Bridge's Qwen3-VL model rebuilds 3D mRoPE positions per packed sub-sequence from ``packed_seq_params`` (``rope.get_rope_index``), without - re-packing the stream (NVIDIA-NeMo/Megatron-Bridge#4532). + re-packing the stream (NVIDIA-NeMo/Megatron-Bridge#4532). Other VLMs only run + unpacked (``model_owns_vlm_packing``). TODO(xgui): context parallelism. preprocess_packed_seqs pre-shards the stream per CP rank, which the bridge model only accepts with explicit rank-local 3D position ids. """ + if self.remove_microbatch_padding and not self.model_owns_vlm_packing: + model_cls = type(unwrap_model(self.actor_module[0])).__name__ + raise ValueError( + "trainer.remove_microbatch_padding=true (sample packing) is supported for VLMs only on " + "Megatron-Bridge's Qwen3VLModel (Qwen3-VL, Qwen3.5-VL), which rebuilds mRoPE positions per " + f"packed sample; got {model_cls}. Set trainer.remove_microbatch_padding=false." + ) assert mpu.get_context_parallel_world_size() == 1, "VLM + context parallelism unsupported" assert ( mpu.get_tensor_model_parallel_world_size() == 1 or self.cfg.policy.sequence_parallel_size == 1 diff --git a/tests/backends/skyrl_train/distributed/test_model_owns_vlm_packing.py b/tests/backends/skyrl_train/distributed/test_model_owns_vlm_packing.py new file mode 100644 index 0000000000..0641290560 --- /dev/null +++ b/tests/backends/skyrl_train/distributed/test_model_owns_vlm_packing.py @@ -0,0 +1,35 @@ +"""model_owns_vlm_packing: which VLMs may use sample packing on the Megatron backend.""" + +import pytest +import torch.nn as nn + +megatron_utils = pytest.importorskip("skyrl.backends.skyrl_train.distributed.megatron.megatron_utils") + + +class _OptIn(nn.Module): + model_owns_packing = True + + +def test_plain_model_does_not_own_packing(): + assert not megatron_utils.model_owns_vlm_packing([nn.Linear(2, 2)]) + + +def test_opt_in_attribute(): + assert megatron_utils.model_owns_vlm_packing([_OptIn()]) + assert megatron_utils.model_owns_vlm_packing(_OptIn()) + + +def test_every_chunk_must_own_packing(): + assert not megatron_utils.model_owns_vlm_packing([_OptIn(), nn.Linear(2, 2)]) + + +def test_no_chunks(): + assert not megatron_utils.model_owns_vlm_packing([]) + + +def test_qwen3_vl_model_class(): + model_mod = pytest.importorskip("megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model") + # Skip construction: an uninitialized instance is enough for the isinstance check. + fake = model_mod.Qwen3VLModel.__new__(model_mod.Qwen3VLModel) + nn.Module.__init__(fake) + assert megatron_utils.model_owns_vlm_packing([fake]) From 8c835f39cd3d3632d13202e7db8d8a8baaafb187 Mon Sep 17 00:00:00 2001 From: xgui Date: Tue, 6 Oct 2026 04:08:00 +0000 Subject: [PATCH 09/10] [test] VLM packing parity: init Ray with the CI env per forward; shape-agnostic empty pixel_values ray_init_for_tests omits NCCL_CUMEM_ENABLE=0 and the other gpu_ci env vars, so use the gpu_ci ray_init context per forward. Build empty pixel_values for text-only rows from a real row's shape. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MZKuFwA4VYK2tfFDVgsWt2 Signed-off-by: xgui --- .../gpu/gpu_ci/megatron/test_megatron_vlm_init.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py index d1d9fa2b2b..386d3f57b8 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py @@ -26,6 +26,7 @@ from skyrl.train.config.sft_config import SFTConfig, SFTPlacementConfig from skyrl.train.sft_trainer import SFTTrainer from skyrl.train.utils.utils import validate_cfg +from tests.backends.skyrl_train.gpu.gpu_ci.conftest import ray_init from tests.backends.skyrl_train.gpu.utils import ( init_worker_with_type, ray_init_for_tests, @@ -414,7 +415,7 @@ def get_packing_parity_batch(model_name: str) -> TrainingInputBatch: assert 0 < len(ids) - prefix <= len(ids) // 2, (prefix, len(ids)) rows.append((ids, len(ids) - prefix, out.get("pixel_values"), out.get("image_grid_thw"))) - pixel_dim = next(pv.shape[-1] for _, _, pv, _ in rows if pv is not None) + ref_pv = next(pv for _, _, pv, _ in rows if pv is not None) max_len = max(len(ids) for ids, _, _, _ in rows) num_actions = max(resp_len for _, resp_len, _, _ in rows) pad_token_id = processor.tokenizer.pad_token_id or processor.tokenizer.eos_token_id @@ -425,7 +426,7 @@ def get_packing_parity_batch(model_name: str) -> TrainingInputBatch: attention_mask.append([0] * pad + [1] * len(ids)) loss_mask.append([0] * (num_actions - resp_len) + [1] * resp_len) # Text-only rows carry empty vision tensors. - pixel_values.append(pv if pv is not None else torch.zeros(0, pixel_dim)) + pixel_values.append(pv if pv is not None else ref_pv.new_zeros((0, *ref_pv.shape[1:]))) image_grid_thw.append(grid if grid is not None else torch.zeros(0, 3, dtype=torch.long)) batch_size = len(rows) @@ -463,16 +464,14 @@ def _megatron_vlm_forward_logprobs( cfg.trainer.micro_forward_batch_size_per_gpu = micro_batch cfg.trainer.micro_train_batch_size_per_gpu = micro_batch cfg.trainer.remove_microbatch_padding = remove_microbatch_padding - try: + # A fresh Ray runtime per forward frees the GPUs for the next one. + with ray_init(): actor_group = init_worker_with_type( "policy", shared_pg=None, colocate_all=False, num_gpus_per_node=tp * pp, cfg=cfg ) refs = actor_group.async_run_ray_method("mesh", "forward", data=batch) output = WorkerOutput.cat(actor_group.actor_infos, ray.get(refs)) return loss_fn_outputs_to_tensor(output.loss_fn_outputs, key="logprobs").float() - finally: - ray.shutdown() - ray_init_for_tests() @pytest.mark.asyncio @@ -492,7 +491,7 @@ def _megatron_vlm_forward_logprobs( ids=["qwen3_vl_tp1", "qwen3_vl_tp2_sp", "qwen3_vl_pp2", "qwen3_5_vl_tp1", "qwen3_5_vl_tp2_sp", "qwen3_5_vl_pp2"], ) @pytest.mark.megatron -async def test_megatron_vlm_packed_vs_unpacked(ray_init_fixture, model_name, tp, pp): +async def test_megatron_vlm_packed_vs_unpacked(model_name, tp, pp): """Sample packing must add no error beyond ordinary batching. Packed and padded (unpacked) microbatches run different kernels on different From b26cf01e52ed4513cde468ec0dceebe4ac6b639a Mon Sep 17 00:00:00 2001 From: xgui Date: Tue, 6 Oct 2026 04:10:09 +0000 Subject: [PATCH 10/10] [train] Note that VLM SFT stays unpacked; packing is enabled for RL Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MZKuFwA4VYK2tfFDVgsWt2 Signed-off-by: xgui --- skyrl/train/sft_trainer.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/skyrl/train/sft_trainer.py b/skyrl/train/sft_trainer.py index 1df1d28ab9..6f208d0c03 100644 --- a/skyrl/train/sft_trainer.py +++ b/skyrl/train/sft_trainer.py @@ -930,10 +930,12 @@ def setup(self): self.is_vlm = check_is_vlm(self.cfg.trainer.policy.model.path) if self.is_vlm: self.processor = get_processor(self.cfg.trainer.policy.model.path, **tokenizer_kwargs) - # Sequence packing / microbatch padding removal are unsupported for - # VLMs (3D RoPE + image token positions). ``remove_microbatch_padding`` - # defaults to True, so disable both unconditionally and mirror the - # change onto the already-built trainer config the workers receive. + # VLM microbatch padding removal is enabled for RL only (Megatron, on + # Megatron-Bridge's Qwen3VLModel); VLM SFT runs unpacked for now. + # ``remove_microbatch_padding`` defaults to True, so disable both + # unconditionally and mirror the change onto the already-built trainer + # config the workers receive. + # TODO(xgui): enable remove_microbatch_padding for VLM SFT on Megatron. if self.sft_cfg.use_sequence_packing or self.sft_cfg.remove_microbatch_padding: logger.warning("VLM detected: disabling sequence packing / microbatch padding removal.") self.sft_cfg.use_sequence_packing = False