fix(vlm): forward get_rope_index to neat packing for mRoPE models

The VLM recipe never passed the model's get_rope_index callable to
neat_pack_dataset_vlm. With it absent, PackedDatasetWrapper sets
has_mrope=False and emits 1D position_ids per pack. The collater then
forwards 2D [B, L] position_ids to the model, which short-circuits
get_rope_index inside model.forward and the language model expands the
same 1D positions across all 3 mRoPE channels.

Net effect: packed Qwen2.5-VL / Qwen3-VL / Qwen3-VL-MoE / Qwen3-Omni
training silently degraded mRoPE to plain 1D rotary, losing image
spatial/temporal positional information. Non-packed and non-mRoPE VLMs
were unaffected.

Plumbing only: extract get_rope_index via getattr(model_parts[0], ...)
in the recipe and forward it through build_dataloader to
neat_pack_dataset_vlm. Models without the method (Mistral3, LLaVA-OV,
KimiVL, Gemma4-VLM) keep the prior behavior since getattr returns None.

Adds two unit tests guarding the wiring against regression.

Signed-off-by: khazic <khazzz1c@gmail.com>

Author: khazic

nemo_automodel/recipes/vlm/finetune.py

@@ -348,6 +348,7 @@ def build_dataloader(
     local_batch_size,
     cfg_model=None,
     cfg_ps=None,
+    get_rope_index=None,
 ) -> tuple[DataLoader, ProcessorMixin]:
     """Build a DataLoader for the VLM dataset.
 
@@ -362,6 +363,11 @@ def build_dataloader(
         cfg_model: Model configuration (used to detect attention backend).
         cfg_ps: Packed sequence configuration (top-level ``packed_sequence:`` section).
             When provided, takes precedence over ``dataset.packing``.
+        get_rope_index: Optional ``model.get_rope_index`` callable. When provided,
+            VLM neat packing computes mRoPE 3D position IDs per sample so packed
+            mRoPE-aware models (Qwen2.5-VL, Qwen3-VL, ...) preserve multimodal
+            position semantics across pack boundaries instead of falling back to
+            plain 1D positions.
 
     Returns:
         The instantiated DataLoader and processor.
@@ -479,6 +485,7 @@ def build_dataloader(
                     packing_ratio=packing_cfg.get("packing_ratio", 1.0),
                     processor=processor,
                     balance_media_tokens=packing_cfg.get("balance_media_tokens", True),
+                    get_rope_index=get_rope_index,
                 )
                 _pad_id = getattr(processor.tokenizer, "pad_token_id", 0) or 0
                 _collate_max_length = packing_cfg.get("collate_max_length", None)
@@ -832,6 +839,11 @@ def setup(self):
             self.model_parts = [model]
             self.pp = None
 
+        # Extract mRoPE position-id builder from the model so VLM neat packing can
+        # produce 3D position_ids per sample. Without this, packed Qwen2.5-VL /
+        # Qwen3-VL training silently degrades mRoPE to plain 1D positions.
+        get_rope_index = getattr(self.model_parts[0], "get_rope_index", None)
+
         self.dataloader, self.processor = build_dataloader(
             self.cfg.dataset,
             self.cfg.dataloader,
@@ -842,6 +854,7 @@ def setup(self):
             local_batch_size=self.cfg.get("step_scheduler.local_batch_size", 1),
             cfg_model=self.cfg.model,
             cfg_ps=self.cfg.get("packed_sequence", None),
+            get_rope_index=get_rope_index,
         )
 
         # Build validation dataloader if the config provides it
@@ -855,6 +868,7 @@ def setup(self):
                 device_mesh=self.device_mesh,
                 seed=self.cfg.get("seed", 42),
                 local_batch_size=self.cfg.get("step_scheduler.local_batch_size", 1),
+                get_rope_index=get_rope_index,
             )
 
         self.best_metric_key = self.cfg.get("checkpoint.best_metric_key", "default")

tests/unit_tests/recipes/test_finetune_vlm_helpers.py

@@ -2432,3 +2432,133 @@ def test_fallback_mismatched_images(self):
         assert len(pv_chunks) == 2
         assert pv_chunks[0].shape[0] == 12  # all in first
         assert pv_chunks[1].shape[0] == 0   # empty
+
+
+# -----------------------------------------------------------------------------
+# get_rope_index forwarding tests for build_dataloader
+#
+# Guard against a regression where the VLM recipe forgot to pass
+# get_rope_index to neat_pack_dataset_vlm, silently degrading mRoPE to
+# plain 1D positions for packed Qwen2.5-VL / Qwen3-VL training.
+# -----------------------------------------------------------------------------
+
+
+def _make_packing_cfg(pack_size=128):
+    cfg = MagicMock()
+    cfg.pack_size = pack_size
+    cfg.pretokenize = True
+    cfg.max_length = pack_size
+    cfg.get.side_effect = lambda key, default=None: {
+        "pack_size": pack_size,
+        "drop_long_samples": True,
+        "max_packs": None,
+        "packing_ratio": 1.0,
+        "balance_media_tokens": True,
+        "collate_max_length": None,
+        "post_tokenize_hook_fn": None,
+    }.get(key, default)
+    return cfg
+
+
+def _make_dataset_cfg():
+    cfg = MagicMock(spec=["get", "instantiate", "path_or_dataset"])
+    cfg.get.side_effect = lambda key, default=None: {
+        "path_or_dataset": None,
+        "truncate": True,
+    }.get(key, default)
+    cfg.path_or_dataset = None
+    cfg.instantiate.return_value = []
+    return cfg
+
+
+def _patches_for_packing(neat_pack_side_effect):
+    processor = MagicMock()
+    processor.tokenizer.pad_token_id = 0
+    processor.chat_template = "{{ x }}"
+    return processor, [
+        patch("transformers.AutoProcessor.from_pretrained", return_value=processor),
+        patch("torch.utils.data.distributed.DistributedSampler"),
+        patch(
+            "nemo_automodel.components.datasets.vlm.datasets.PreTokenizedDatasetWrapper",
+            return_value=MagicMock(),
+        ),
+        patch(
+            "nemo_automodel.components.datasets.vlm.neat_packing_vlm.neat_pack_dataset_vlm",
+            side_effect=neat_pack_side_effect,
+        ),
+        patch("nemo_automodel.components.models.common.packing.configure_packing"),
+        patch(
+            "nemo_automodel.components.models.common.packing.get_attn_implementation",
+            return_value="sdpa",
+        ),
+    ]
+
+
+def test_build_dataloader_forwards_get_rope_index_to_packing():
+    """get_rope_index passed to build_dataloader must reach neat_pack_dataset_vlm."""
+    from contextlib import ExitStack
+
+    from nemo_automodel.recipes.vlm.finetune import build_dataloader
+
+    sentinel = MagicMock(name="get_rope_index")
+    captured = {}
+
+    def fake_neat_pack(*args, **kwargs):
+        captured.update(kwargs)
+        return MagicMock()
+
+    _, ctx_managers = _patches_for_packing(fake_neat_pack)
+
+    with ExitStack() as stack:
+        for cm in ctx_managers:
+            stack.enter_context(cm)
+        build_dataloader(
+            _make_dataset_cfg(),
+            MagicMock(get=MagicMock(return_value=None), instantiate=MagicMock(return_value=MagicMock())),
+            "test/model",
+            None,
+            None,
+            42,
+            1,
+            cfg_ps=_make_packing_cfg(pack_size=64),
+            get_rope_index=sentinel,
+        )
+
+    assert captured.get("get_rope_index") is sentinel, (
+        "build_dataloader must forward get_rope_index to neat_pack_dataset_vlm; "
+        f"got kwargs={list(captured.keys())}"
+    )
+
+
+def test_build_dataloader_default_get_rope_index_is_none():
+    """When the model does not expose get_rope_index, packing must receive None."""
+    from contextlib import ExitStack
+
+    from nemo_automodel.recipes.vlm.finetune import build_dataloader
+
+    captured = {}
+
+    def fake_neat_pack(*args, **kwargs):
+        captured.update(kwargs)
+        return MagicMock()
+
+    _, ctx_managers = _patches_for_packing(fake_neat_pack)
+
+    with ExitStack() as stack:
+        for cm in ctx_managers:
+            stack.enter_context(cm)
+        build_dataloader(
+            _make_dataset_cfg(),
+            MagicMock(get=MagicMock(return_value=None), instantiate=MagicMock(return_value=MagicMock())),
+            "test/model",
+            None,
+            None,
+            42,
+            1,
+            cfg_ps=_make_packing_cfg(pack_size=64),
+        )
+
+    assert "get_rope_index" in captured, (
+        "neat_pack_dataset_vlm must receive get_rope_index kwarg even when None"
+    )
+    assert captured["get_rope_index"] is None