feat(video): enable canonical Nemotron Omni training for V2 MoE checkpoints

src/megatron/bridge/models/nemotron_omni/modeling_nemotron_omni.py

@@ -287,10 +287,20 @@ def _merge_projected_media(
         media_token_id: int,
         attention_mask: Optional[torch.Tensor],
     ) -> torch.Tensor:
-        """Replace each valid media placeholder with exactly one feature row."""
+        """Replace each valid media placeholder with exactly one feature row.
+
+        ``attention_mask`` is a token-validity mask here, not MCore's 4-D
+        causal attention mask.  Requiring an exact shape match prevents a
+        causal mask from broadcasting the placeholder mask to ``[B, 1, S, S]``.
+        """
 
         media_mask = input_ids == media_token_id
         if attention_mask is not None:
+            if attention_mask.shape != input_ids.shape:
+                raise ValueError(
+                    "The media token-validity mask must have the same shape as input_ids: "
+                    f"got mask={tuple(attention_mask.shape)}, input_ids={tuple(input_ids.shape)}."
+                )
             media_mask = media_mask & attention_mask.bool()
 
         expected_features = int(media_mask.sum().item())
@@ -679,12 +689,23 @@ def forward(
             if image_embeddings is None:
                 image_embeddings = combined_embeddings.new_empty((0, combined_embeddings.shape[-1]))
 
+            # MBridge collators use a 2-D attention mask as a token-validity
+            # mask, while NeMo RL's dense Megatron path supplies MCore's 4-D
+            # causal mask (where True means blocked).  Only the former can
+            # filter media placeholders.  Padding masks are unambiguous and
+            # take precedence for collator-owned packed inputs.
+            media_token_validity_mask = None
+            if padding_mask is not None:
+                media_token_validity_mask = ~padding_mask
+            elif attention_mask is not None and attention_mask.dim() == input_ids.dim():
+                media_token_validity_mask = attention_mask
+
             combined_embeddings = self._merge_projected_media(
                 combined_embeddings,
                 input_ids,
                 image_embeddings,
                 self.image_token_index,
-                ~padding_mask if packed_seq_params is not None and padding_mask is not None else attention_mask,
+                media_token_validity_mask,
             )
 
             if self.sound_token_index > 0:
@@ -695,7 +716,7 @@ def forward(
                     input_ids,
                     sound_embeddings,
                     self.sound_token_index,
-                    ~padding_mask if packed_seq_params is not None and padding_mask is not None else attention_mask,
+                    media_token_validity_mask,
                 )
 
         if packed_seq_params is not None:

src/megatron/bridge/models/nemotron_omni/nemotron_omni_bridge.py

@@ -202,7 +202,9 @@ def megatron_to_hf_config(cls, provider) -> dict:
 
     def _llava_mapping_registry(self) -> MegatronMappingRegistry:
         """Build mappings for the historical LLaVA wrapper namespace."""
-        vl_registry = super().mapping_registry()
+        # Call the explicit legacy implementation. NemotronVLBridge.mapping_registry
+        # can route V2-labeled MoE checkpoints back to this canonical bridge.
+        vl_registry = self._legacy_mapping_registry()
         mapping_list = list(vl_registry.mappings)
 
         # MoE language decoder (not present in the dense VL variant).

src/megatron/bridge/models/nemotron_vl/nemotron_vl_bridge.py

@@ -29,6 +29,18 @@
 from megatron.bridge.models.nemotron_vl.nemotron_vl_provider import NemotronVLModelProvider
 
 
+def _is_legacy_v2_omni_config(hf_config) -> bool:
+    """Identify Nano Omni checkpoints serialized with the historical V2 label."""
+
+    architectures = getattr(hf_config, "architectures", None) or []
+    llm_config = getattr(hf_config, "llm_config", None)
+    return (
+        "NemotronH_Nano_VL_V2" in architectures
+        and llm_config is not None
+        and getattr(llm_config, "n_routed_experts", None) is not None
+    )
+
+
 @MegatronModelBridge.register_bridge(
     source="NemotronH_Nano_VL_V2",
     target=NemotronVLModel,
@@ -48,10 +60,37 @@ class NemotronVLBridge(MegatronModelBridge):
     # Provider translation
     # ------------------------------------------------------------------
 
-    def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> NemotronVLModelProvider:  # type: ignore[override]
+    def _canonical_omni_bridge(self):
+        """Create the canonical bridge lazily to avoid a module import cycle."""
+
+        from megatron.bridge.models.nemotron_omni.nemotron_omni_bridge import NemotronOmniBridge
+
+        bridge = NemotronOmniBridge()
+        bridge.hf_pretrained = getattr(self, "hf_pretrained", None)
+        bridge.hf_config = getattr(self, "hf_config", None)
+        return bridge
+
+    def provider_bridge(self, hf_pretrained: PreTrainedCausalLM):  # type: ignore[override]
         hf_config = hf_pretrained.config
         llm_config = hf_config.llm_config
 
+        if _is_legacy_v2_omni_config(hf_config):
+            # Some Nano Omni checkpoints were exported before the dedicated
+            # architecture name existed. Route only the MoE-shaped V2 configs
+            # to the canonical expanded-sequence model; dense V2 checkpoints
+            # continue to use the historical NemotronVLModel/LLaVAModel path.
+            bridge = self._canonical_omni_bridge()
+            bridge.hf_pretrained = hf_pretrained
+            bridge.hf_config = hf_config
+            provider = bridge.provider_bridge(hf_pretrained)
+
+            # The historical tokenizer uses <img>/<\/img> IDs 19/20, whereas
+            # the public V3 checkpoint uses 21/22. Prefer serialized numeric
+            # values when present and otherwise preserve the V2 contract.
+            provider.img_start_token_id = getattr(hf_config, "img_start_token_id", None) or 19
+            provider.img_end_token_id = getattr(hf_config, "img_end_token_id", None) or 20
+            return provider
+
         # Use base class helper for common config mapping
         provider_kwargs = self.hf_config_to_provider_kwargs(llm_config)
 
@@ -80,6 +119,13 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> NemotronVLModelP
     # ------------------------------------------------------------------
 
     def mapping_registry(self) -> MegatronMappingRegistry:  # noqa: D401
+        if _is_legacy_v2_omni_config(getattr(self, "hf_config", None)):
+            return self._canonical_omni_bridge().mapping_registry()
+        return self._legacy_mapping_registry()
+
+    def _legacy_mapping_registry(self) -> MegatronMappingRegistry:
+        """Return the historical wrapper-prefixed Nemotron-VL mappings."""
+
         param_mappings = {
             # vision model
             "llava_model.vision_model.class_token": "vision_model.radio_model.model.patch_generator.cls_token.token",

tests/unit_tests/models/nemotron_omni/test_nemotron_omni_conversion.py

@@ -41,6 +41,7 @@
 )
 from megatron.bridge.models.nemotron_vl.modeling_nemotron_vl import NemotronVLModel
 from megatron.bridge.models.nemotron_vl.nemotron_vl_bridge import NemotronVLBridge
+from megatron.bridge.models.nemotron_vl.nemotron_vl_provider import NemotronVLModelProvider
 from megatron.bridge.training.config import ConfigContainer
 
 
@@ -120,6 +121,18 @@ def _mock_omni_hf_config():
     )
 
 
+def _mock_legacy_v2_omni_hf_config():
+    """Represent Nano Omni weights exported before the V3 architecture name."""
+
+    hf_config = _mock_omni_hf_config()
+    hf_config.architectures = ["NemotronH_Nano_VL_V2"]
+    hf_config.model_type = "NemotronH_Nano_VL_V2"
+    del hf_config.sound_config
+    del hf_config.sound_context_token_id
+    hf_config.vision_config.args = {"register_multiple": 10}
+    return hf_config
+
+
 def test_public_nemotron_omni_architecture_is_registered():
     hf_config = _mock_omni_hf_config()
 
@@ -131,6 +144,46 @@ def test_public_nemotron_omni_architecture_is_registered():
     assert isinstance(get_model_bridge("NemotronH_Super_Omni_Reasoning_V3", hf_config=hf_config), NemotronOmniBridge)
 
 
+def test_legacy_v2_moe_checkpoint_routes_to_canonical_nemotron_omni():
+    hf_config = _mock_legacy_v2_omni_hf_config()
+    hf_pretrained = Mock(spec=PreTrainedCausalLM)
+    hf_pretrained.config = hf_config
+
+    bridge = get_model_bridge("NemotronH_Nano_VL_V2", hf_config=hf_config)
+    provider = bridge.provider_bridge(hf_pretrained)
+    registry = bridge.mapping_registry()
+
+    assert isinstance(bridge, NemotronVLBridge)
+    assert isinstance(provider, NemotronOmniModelProvider)
+    assert provider.nemotron_omni_contract == NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT
+    assert provider.image_token_index == 18
+    assert provider.img_start_token_id == 19
+    assert provider.img_end_token_id == 20
+    assert provider.has_sound is False
+    assert provider.separate_video_embedder is True
+    assert provider.temporal_patch_dim == 2
+    assert provider.temporal_ckpt_compat is True
+    video_mapping = registry.hf_to_megatron_lookup(
+        "vision_model.radio_model.model.patch_generator.video_embedder.weight"
+    )
+    assert video_mapping.megatron_param == "vision_model.video_embedder.weight"
+    assert all(not mapping.megatron_param.startswith("llava_model.") for mapping in registry.mappings)
+
+
+def test_dense_legacy_v2_checkpoint_keeps_nemotron_vl_path():
+    hf_config = _mock_legacy_v2_omni_hf_config()
+    del hf_config.llm_config.n_routed_experts
+    hf_pretrained = Mock(spec=PreTrainedCausalLM)
+    hf_pretrained.config = hf_config
+
+    bridge = get_model_bridge("NemotronH_Nano_VL_V2", hf_config=hf_config)
+    provider = bridge.provider_bridge(hf_pretrained)
+    registry = bridge.mapping_registry()
+
+    assert isinstance(provider, NemotronVLModelProvider)
+    assert any(mapping.megatron_param.startswith("llava_model.") for mapping in registry.mappings)
+
+
 def test_nemotron_omni_provider_bridge_maps_public_config_fields():
     hf_config = _mock_omni_hf_config()
     hf_pretrained = Mock(spec=PreTrainedCausalLM)

tests/unit_tests/models/nemotron_omni/test_nemotron_omni_model.py

@@ -316,6 +316,24 @@ def test_image_forward_replaces_expanded_placeholders_without_changing_length():
     assert torch.equal(output[3, 0], torch.tensor([9.0, 9.0, 9.0]))
 
 
+def test_image_forward_does_not_use_mcore_causal_mask_as_token_validity():
+    image_features = torch.tensor([[101.0, 102.0, 103.0], [201.0, 202.0, 203.0]])
+    model = _BoundaryModel(image_features)
+    input_ids = torch.tensor([[7, 18, 18, 9]])
+    causal_attention_mask = torch.triu(torch.ones(1, 1, 4, 4, dtype=torch.bool), diagonal=1)
+
+    output = model(
+        input_ids=input_ids,
+        attention_mask=causal_attention_mask,
+        images=torch.ones(1),
+    )
+
+    assert output.shape == (4, 1, 3)
+    assert torch.equal(output[1, 0], image_features[0])
+    assert torch.equal(output[2, 0], image_features[1])
+    assert torch.equal(model.language_model.last_kwargs["attention_mask"], causal_attention_mask)
+
+
 def test_audio_forward_replaces_expanded_placeholders_without_changing_length():
     sound_features = torch.tensor([[101.0, 102.0, 103.0], [201.0, 202.0, 203.0]])
     model = _BoundaryModel(torch.empty(0, 3), sound_features)