Support a new model (NemotronH_Omni_Reasoning_V3) in Megatron-Bridge PTQ/QAD and export

modelopt/torch/export/plugins/mcore_common.py

@@ -34,6 +34,7 @@
     nemotron_h_causal_lm_export,
     nemotron_h_causal_lm_import,
 )
+from .mcore_nemotron_omni import NEMOTRON_H_OMNI_VISION_PREFIXES, nemotron_h_omni_causal_lm_export
 from .mcore_qwen import (
     qwen3_causal_lm_export,
     qwen3_causal_lm_import,
@@ -54,6 +55,7 @@
     "Llama4ForConditionalGeneration": llama4_causal_lm_export,
     "NemotronForCausalLM": nemotron_causal_lm_export,
     "NemotronHForCausalLM": nemotron_h_causal_lm_export,
+    "NemotronH_Omni_Reasoning_V3": nemotron_h_omni_causal_lm_export,
     "LlamaForCausalLMEagle": eagle_llama_causal_lm_export,
     "LlamaForCausalLMEagle3": eagle3_llama_causal_lm_export,
     "LlamaForCausalLMEagle3Deep": eagle3_deep_llama_causal_lm_export,
@@ -74,6 +76,7 @@
     "Qwen3VLForConditionalGeneration": QWEN3VL_VISION_PREFIXES,
     "Qwen3_5ForConditionalGeneration": QWEN3_5_VL_VISION_PREFIXES,
     "Qwen3_5MoeForConditionalGeneration": QWEN3_5_VL_VISION_PREFIXES,
+    "NemotronH_Omni_Reasoning_V3": NEMOTRON_H_OMNI_VISION_PREFIXES,
 }
 
 all_mcore_hf_import_mapping: dict[str, Any] = {

modelopt/torch/export/plugins/mcore_custom.py

@@ -255,16 +255,22 @@ def __init__(self, target_name_or_prefix: str = "", func_kwargs: dict[str, Any]
 
 def with_language_model_prefix(
     mapping: dict[str, CustomModuleMapping],
+    old_prefix: str = "model.",
+    new_prefix: str = "model.language_model.",
 ) -> dict[str, CustomModuleMapping]:
-    """Nest a text-model mapping under ``model.language_model.``; other prefixes are unchanged."""
+    """Nest a text-model mapping under the VLM's language-model prefix.
+
+    Targets starting with ``old_prefix`` have it replaced by ``new_prefix``; other targets are
+    unchanged. ``old_prefix=""`` matches every target, prefixing the whole mapping.
+    """
     result = {}
     for key, m in mapping.items():
         if not isinstance(m, CustomModuleMapping):
             result[key] = m  # plain flags such as ``use_packed_local_experts``
             continue
         prefix = m.target_name_or_prefix
-        if prefix.startswith("model."):
-            prefix = "model.language_model." + prefix[len("model.") :]
+        if prefix.startswith(old_prefix):
+            prefix = new_prefix + prefix[len(old_prefix) :]
         result[key] = type(m)(
             target_name_or_prefix=prefix, func_kwargs=copy.deepcopy(m.func_kwargs)
         )

modelopt/torch/export/plugins/mcore_nemotron_omni.py

@@ -0,0 +1,31 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Custom mapping from the new model (``NemotronH_Omni_Reasoning_V3``) to Megatron Core models.
+
+The new model nests the whole Nemotron-H language model, ``lm_head`` included, under
+``language_model.``, so the mapping is derived from Nemotron-H's (MTP included). The vision tower
+and projector are copied verbatim from HF via ``NEMOTRON_H_OMNI_VISION_PREFIXES``, not mapped.
+"""
+
+from .mcore_custom import with_language_model_prefix
+from .mcore_nemotron import nemotron_h_causal_lm_export
+
+# Weights copied straight from the HF checkpoint (never quantized).
+NEMOTRON_H_OMNI_VISION_PREFIXES = ("vision_model.", "vision_projector.", "mlp1.")
+
+nemotron_h_omni_causal_lm_export = with_language_model_prefix(
+    nemotron_h_causal_lm_export, old_prefix="", new_prefix="language_model."
+)

modelopt/torch/export/unified_export_megatron.py

@@ -883,6 +883,10 @@ def _copy_mtp_state_dict_from_pretrained(self) -> dict[str, torch.Tensor]:
         if self.rules.get("mtp_in_decoder_layers", False):
             return self._copy_decoder_mtp_layers_from_pretrained()
 
+        # The MTP root follows the arch's own mapping when it has one: "mtp." for Nemotron-H,
+        # "language_model.mtp." for a VLM that nests it.
+        mtp_rule = self.all_mcore_mappings.get("mtp.eh_proj")
+        mtp_prefix = mtp_rule.target_name_or_prefix.split("layers.")[0] if mtp_rule else "mtp."
         mtp_exists = False
 
         if os.path.isdir(self._hf_pretrained_model_name):
@@ -917,22 +921,22 @@ def _copy_mtp_state_dict_from_pretrained(self) -> dict[str, torch.Tensor]:
                 safetensors_index = json.load(f)
             model_dir = safetensors_index_file.parent
             for key in safetensors_index["weight_map"]:
-                if key.startswith("mtp.") and key not in self._state_dict:
+                if key.startswith(mtp_prefix) and key not in self._state_dict:
                     mtp_state_dict[key] = get_safetensor(model_dir, key)
                     mtp_exists = True
             if mtp_exists:
                 print(f"Exported MTP using {safetensors_index_file=}")
         elif single_safetensors_file is not None and single_safetensors_file.exists():
             with safe_open(str(single_safetensors_file), framework="pt", device="cpu") as f:
                 for key in f.keys():  # noqa: SIM118
-                    if key.startswith("mtp.") and key not in self._state_dict:
+                    if key.startswith(mtp_prefix) and key not in self._state_dict:
                         mtp_state_dict[key] = f.get_tensor(key)
                         mtp_exists = True
             if mtp_exists:
                 print(f"Exported MTP using {single_safetensors_file=}")
 
         if mtp_exists:
-            self.exclude_modules.append("mtp*")
+            self.exclude_modules.append(mtp_prefix.removesuffix(".") + "*")
         return mtp_state_dict
 
     def _copy_decoder_mtp_layers_from_pretrained(self) -> dict[str, torch.Tensor]:

tests/gpu_megatron/torch/export/plugins/test_mcore_export_mappings.py

@@ -17,7 +17,10 @@
 
 import pytest
 
-from modelopt.torch.export.plugins.mcore_common import all_mcore_hf_export_mapping
+from modelopt.torch.export.plugins.mcore_common import (
+    all_mcore_hf_export_mapping,
+    all_mcore_hf_vision_passthrough_mapping,
+)
 
 # ``_GPTModelExporter`` only emits ``k_scale`` / ``v_scale`` and sets ``kv_cache_quant_algo``
 # for layers whose architecture mapping defines ``core_attention``. A missing entry is silent:
@@ -34,6 +37,7 @@
     "Qwen3VLForConditionalGeneration": "model.language_model.layers.{}.self_attn.",
     "Qwen3_5ForConditionalGeneration": "model.language_model.layers.{}.self_attn.",
     "Qwen3_5MoeForConditionalGeneration": "model.language_model.layers.{}.self_attn.",
+    "NemotronH_Omni_Reasoning_V3": "language_model.backbone.layers.{}.mixer.",
     "GlmMoeDsaForCausalLM": "model.layers.{}.self_attn.",
 }
 
@@ -55,6 +59,7 @@ def test_export_mapping_emits_kv_cache_scales(arch, prefix):
     "Qwen3_5MoeForConditionalGeneration": "model.language_model.layers.{}.mlp.experts.{}",
     "Qwen3MoeForCausalLM": "model.layers.{}.mlp.experts.{}",
     "NemotronHForCausalLM": "backbone.layers.{}.mixer.experts.{}",
+    "NemotronH_Omni_Reasoning_V3": "language_model.backbone.layers.{}.mixer.experts.{}",
     "GlmMoeDsaForCausalLM": "model.layers.{}.mlp.experts.{}",
 }
 
@@ -84,3 +89,19 @@ def test_qwen3_5_moe_expert_names_match_released_checkpoint():
 
     fc2 = mapping["experts.linear_fc2"].target_name_or_prefix.format(7).format(3) + "."
     assert fc2 == "model.language_model.layers.7.mlp.experts.3.down_proj."
+
+
+def test_new_model_nests_every_target_under_language_model():
+    """The new model keeps the whole Nemotron-H language model, ``lm_head`` included, under one prefix."""
+    text = all_mcore_hf_export_mapping["NemotronHForCausalLM"]
+    vl = all_mcore_hf_export_mapping["NemotronH_Omni_Reasoning_V3"]
+    assert vl.keys() == text.keys()
+    for key, rule in vl.items():
+        assert type(rule) is type(text[key])
+        assert rule.target_name_or_prefix == "language_model." + text[key].target_name_or_prefix
+        assert rule.func_kwargs == text[key].func_kwargs
+    assert vl["output_layer"].target_name_or_prefix == "language_model.lm_head."
+    # MTP is exported from the live model; a passthrough would write it a second time.
+    assert vl["mtp.eh_proj"].target_name_or_prefix.startswith("language_model.mtp.")
+    passthrough = all_mcore_hf_vision_passthrough_mapping["NemotronH_Omni_Reasoning_V3"]
+    assert not any(p.startswith("language_model.") for p in passthrough)

tests/gpu_megatron/torch/export/plugins/test_moe_layout_choice.py

@@ -24,6 +24,7 @@
     [
         # These map fused grouped-GEMM experts, so they keep the faster layout.
         ("NemotronHForCausalLM", True),
+        ("NemotronH_Omni_Reasoning_V3", True),
         ("Qwen3_5MoeForConditionalGeneration", True),
         ("Qwen3MoeForCausalLM", False),
         ("DeepseekV3ForCausalLM", False),

tests/gpu_megatron/torch/export/test_unified_export_megatron.py

@@ -979,16 +979,41 @@ def test_qkv_slicing_records_hf_excludes_for_unquantized_fused_qkv():
     assert "backbone.layers.0.mixer.v_proj" in exporter.exclude_modules
 
 
-def _make_exporter_for_mtp(model_dir: Path) -> GPTModelExporter:
+def _make_exporter_for_mtp(model_dir: Path, arch: str | None = None) -> GPTModelExporter:
     """Create a minimal GPTModelExporter instance for testing _get_mtp_state_dict."""
     exporter = object.__new__(GPTModelExporter)
     exporter._hf_pretrained_model_name = str(model_dir)
     exporter._state_dict = {}  # MTP keys are absent — they should be picked up
     exporter.exclude_modules = []
     exporter.rules = {}  # a Qwen-style ``mtp.*`` checkpoint, not MTP stored as decoder layers
+    exporter.all_mcore_mappings = all_mcore_hf_export_mapping[arch] if arch else {}
     return exporter
 
 
+@pytest.mark.parametrize(
+    ("arch", "mtp_prefix"),
+    [
+        ("NemotronHForCausalLM", "mtp."),
+        ("NemotronH_Omni_Reasoning_V3", "language_model.mtp."),
+    ],
+)
+def test_mtp_state_dict_prefix_follows_mapping(tmp_path, arch, mtp_prefix):
+    """Without a live MTP, the source MTP is copied from the root the arch's mapping uses."""
+    model_dir = tmp_path / "fake_hf_model"
+    model_dir.mkdir()
+    tensors = {
+        "language_model.backbone.embeddings.weight": torch.zeros(64, 32),
+        f"{mtp_prefix}layers.0.enorm.weight": torch.ones(32),
+    }
+    save_file(tensors, str(model_dir / "model.safetensors"))
+
+    exporter = _make_exporter_for_mtp(model_dir, arch)
+    mtp_state_dict = exporter._get_mtp_state_dict()
+
+    assert list(mtp_state_dict) == [f"{mtp_prefix}layers.0.enorm.weight"]
+    assert exporter.exclude_modules == [mtp_prefix.removesuffix(".") + "*"]
+
+
 def test_mtp_state_dict_single_safetensors(tmp_path):
     """MTP weights are collected from a model with a single model.safetensors file."""
     model_dir = tmp_path / "fake_hf_model"