From ae307c1719492572494f606ff92670d067457c53 Mon Sep 17 00:00:00 2001 From: Yue Date: Fri, 9 Oct 2026 03:34:05 +0000 Subject: [PATCH 1/3] Support a new model (NemotronH_Omni_Reasoning_V3) in Megatron-Bridge export The new model nests a Nemotron-H language model, lm_head included, under language_model., but its architecture had no entry in the export mappings. use_moe_grouped_gemm() therefore built its experts as SequentialMLP, and HF export had no rules for it. Register the architecture in its own mcore_nemotron_omni module: the Nemotron-H rules (MTP included) nested under language_model., plus a vision passthrough for vision_model., vision_projector. and mlp1. Generalize with_language_model_prefix() with old_prefix/new_prefix so it covers this layout; the defaults keep the Qwen-VL behavior. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yue --- CHANGELOG.rst | 1 + modelopt/torch/export/plugins/mcore_common.py | 3 ++ modelopt/torch/export/plugins/mcore_custom.py | 12 +++++-- .../export/plugins/mcore_nemotron_omni.py | 31 +++++++++++++++++++ .../plugins/test_mcore_export_mappings.py | 23 +++++++++++++- .../export/plugins/test_moe_layout_choice.py | 1 + 6 files changed, 67 insertions(+), 4 deletions(-) create mode 100644 modelopt/torch/export/plugins/mcore_nemotron_omni.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index f287c409c71..7aa7e14cdc3 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -48,6 +48,7 @@ Changelog *Megatron Framework (M-LM / M-Bridge)* +- Support a new model (``NemotronH_Omni_Reasoning_V3``) in the ``examples/megatron_bridge`` scripts: its MoE experts now build as grouped-GEMM (``TEGroupedMLP``), and Hugging Face export writes the quantized language model under ``language_model.`` while copying the vision tower and projector from the source checkpoint. - Add optional Top-P (nucleus) truncation to the Megatron top-k logits KD loss (now ``TopLogitsKLLoss``) via ``logit_kl_top_p`` and ``logit_kl_top_p_min_k`` in ``DistillationConfig``: after the global Top-K selection, only the smallest prefix whose cumulative teacher probability reaches ``top_p`` (with a floor of ``min_k`` entries) contributes to the KL, mirroring ``--logits-save-top-p`` / ``--logits-save-top-p-min-k`` in Megatron-LM's logits saver. - Add an end-to-end W4A4 NVFP4 PTQ and QAD tutorial for Qwen3.6-35B-A3B also covering evaluation and vLLM throughput benchmarking. See `examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/README.md `_ for details. - Add ``--mlflow `` to the ``examples/megatron_bridge`` scripts that write a checkpoint -- ``prune_minitron.py``, ``quantize.py``, ``distill.py``, ``export_quantized_megatron_to_hf.py`` and ``export_distilled_megatron_to_hf.py`` (MLflow's own ``MLFLOW_TRACKING_URI`` is honoured too). Each run records the invocation, its arguments as searchable params and its log -- a distillation records its training metrics instead of the rank-0 log -- and writes ``.experiment.json`` into the checkpoint it produced, so a pruning, a quantization, the distillation that refines its checkpoint and the export that deploys it can be traced to one another; uploading the checkpoints themselves stays off unless ``--mlflow_log_checkpoints`` is passed. 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..9b71891fb1b 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. + """ 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/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), From 9540d5b4ee2641a88e147edb54f2396514539c25 Mon Sep 17 00:00:00 2001 From: Yue Date: Fri, 9 Oct 2026 05:07:24 +0000 Subject: [PATCH 2/3] Copy the source MTP from the arch's own MTP root when there is no live MTP _copy_mtp_state_dict_from_pretrained matched a hard-coded "mtp." root, so for NemotronH_Omni_Reasoning_V3 (MTP under language_model.mtp.) a model built without MTP exported no MTP head. Derive the root from the arch's mtp.eh_proj rule instead; archs without one keep "mtp.", and Nemotron-H resolves to "mtp." unchanged. Also note in with_language_model_prefix that old_prefix="" prefixes every target. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yue --- modelopt/torch/export/plugins/mcore_custom.py | 2 +- .../torch/export/unified_export_megatron.py | 10 ++++--- .../export/test_unified_export_megatron.py | 27 ++++++++++++++++++- 3 files changed, 34 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/export/plugins/mcore_custom.py b/modelopt/torch/export/plugins/mcore_custom.py index 9b71891fb1b..c34171cbd80 100644 --- a/modelopt/torch/export/plugins/mcore_custom.py +++ b/modelopt/torch/export/plugins/mcore_custom.py @@ -261,7 +261,7 @@ def with_language_model_prefix( """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. + unchanged. ``old_prefix=""`` matches every target, prefixing the whole mapping. """ result = {} for key, m in mapping.items(): 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/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" From eb8909c01ffb2c14682f09851b6cafa51afa8614 Mon Sep 17 00:00:00 2001 From: Yue Date: Fri, 9 Oct 2026 05:11:47 +0000 Subject: [PATCH 3/3] Drop the CHANGELOG entry for now Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yue --- CHANGELOG.rst | 1 - 1 file changed, 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7aa7e14cdc3..f287c409c71 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -48,7 +48,6 @@ Changelog *Megatron Framework (M-LM / M-Bridge)* -- Support a new model (``NemotronH_Omni_Reasoning_V3``) in the ``examples/megatron_bridge`` scripts: its MoE experts now build as grouped-GEMM (``TEGroupedMLP``), and Hugging Face export writes the quantized language model under ``language_model.`` while copying the vision tower and projector from the source checkpoint. - Add optional Top-P (nucleus) truncation to the Megatron top-k logits KD loss (now ``TopLogitsKLLoss``) via ``logit_kl_top_p`` and ``logit_kl_top_p_min_k`` in ``DistillationConfig``: after the global Top-K selection, only the smallest prefix whose cumulative teacher probability reaches ``top_p`` (with a floor of ``min_k`` entries) contributes to the KL, mirroring ``--logits-save-top-p`` / ``--logits-save-top-p-min-k`` in Megatron-LM's logits saver. - Add an end-to-end W4A4 NVFP4 PTQ and QAD tutorial for Qwen3.6-35B-A3B also covering evaluation and vLLM throughput benchmarking. See `examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/README.md `_ for details. - Add ``--mlflow `` to the ``examples/megatron_bridge`` scripts that write a checkpoint -- ``prune_minitron.py``, ``quantize.py``, ``distill.py``, ``export_quantized_megatron_to_hf.py`` and ``export_distilled_megatron_to_hf.py`` (MLflow's own ``MLFLOW_TRACKING_URI`` is honoured too). Each run records the invocation, its arguments as searchable params and its log -- a distillation records its training metrics instead of the rank-0 log -- and writes ``.experiment.json`` into the checkpoint it produced, so a pruning, a quantization, the distillation that refines its checkpoint and the export that deploys it can be traced to one another; uploading the checkpoints themselves stays off unless ``--mlflow_log_checkpoints`` is passed.