diff --git a/modelopt/torch/export/plugins/mcore_common.py b/modelopt/torch/export/plugins/mcore_common.py index b74a4fc6aaa..fbfbea57133 100644 --- a/modelopt/torch/export/plugins/mcore_common.py +++ b/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] = { diff --git a/modelopt/torch/export/plugins/mcore_custom.py b/modelopt/torch/export/plugins/mcore_custom.py index ed3e00fd962..c34171cbd80 100644 --- a/modelopt/torch/export/plugins/mcore_custom.py +++ b/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) ) diff --git a/modelopt/torch/export/plugins/mcore_nemotron_omni.py b/modelopt/torch/export/plugins/mcore_nemotron_omni.py new file mode 100644 index 00000000000..a82badfeb7d --- /dev/null +++ b/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." +) diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 1d6a0fc99d7..ff0db546072 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/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,7 +921,7 @@ 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: @@ -925,14 +929,14 @@ def _copy_mtp_state_dict_from_pretrained(self) -> dict[str, torch.Tensor]: 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]: diff --git a/tests/gpu_megatron/torch/export/plugins/test_mcore_export_mappings.py b/tests/gpu_megatron/torch/export/plugins/test_mcore_export_mappings.py index 8b048e87bcd..965c44d7c29 100644 --- a/tests/gpu_megatron/torch/export/plugins/test_mcore_export_mappings.py +++ b/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) diff --git a/tests/gpu_megatron/torch/export/plugins/test_moe_layout_choice.py b/tests/gpu_megatron/torch/export/plugins/test_moe_layout_choice.py index 197f5923dac..7d1a301e612 100644 --- a/tests/gpu_megatron/torch/export/plugins/test_moe_layout_choice.py +++ b/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), diff --git a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py index 286700a9338..5e31b000674 100644 --- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py +++ b/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"