[megatron][vlm] Allow sequence packing for VLMs incl. Qwen3.5-VL (CP=1)

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

skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py

@@ -751,28 +751,30 @@ 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``.
+def model_owns_vlm_packing(model: Union[nn.Module, List[nn.Module]]) -> bool:
+    """Whether the VLM handles a packed [1, T] stream itself.
 
-    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 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:
-        return False
+        Qwen3VLModel = None
 
     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
+        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(

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,7 +20,7 @@
 from skyrl.backends.skyrl_train.distributed.megatron.megatron_utils import (
     get_model_config,
     make_batch_generator,
-    model_packs_sequences_internally,
+    model_owns_vlm_packing,
     preprocess_packed_seqs,
     recover_left_padding,
     remove_left_padding,
@@ -215,24 +216,14 @@ 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.
         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
@@ -299,12 +290,24 @@ 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``), without
+        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.
         """
-        assert not self.remove_microbatch_padding, "VLM + microbatch padding removal unsupported"
+        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

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 "

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.

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

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

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