From c5a03a58452e7d2dd5fed7a22360386b92a95e86 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 28 Aug 2026 05:11:05 -0700
Subject: [PATCH 01/12] Support quantized Qwen3-VL export to unified HF from
 Megatron-Bridge

Enables the PTQ / QAD -> unified HuggingFace export path for Qwen3-VL,
and fixes a silent VLM QAD state loss found along the way.

- GPTModelExporter only unwrapped MCore LLaVAModel, so Megatron-Bridge's
  Qwen3VLModel was rejected. Unwrap any wrapper exposing .language_model.
- A VLM QAD checkpoint holds the language model only (distill_submodule),
  so load it into .language_model rather than the full VLM wrapper.
- PTQ anchors the ModelOpt state on the VLM root but QAD checkpoints only
  the language model, so the state was dropped and the export came out
  unquantized. Move it to .language_model on QAD restore.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 CHANGELOG.rst                                 |  2 ++
 examples/megatron_bridge/README.md            |  2 +-
 examples/megatron_bridge/distill.py           | 11 +++++---
 .../export_quantized_megatron_to_hf.py        |  5 ++--
 examples/megatron_bridge/quantize.py          |  4 +--
 .../torch/export/unified_export_megatron.py   | 25 ++++++++++-------
 modelopt/torch/utils/plugins/mbridge.py       | 22 +++++++++++++--
 tests/examples/megatron_bridge/test_qad.py    | 27 ++++++++++---------
 8 files changed, 66 insertions(+), 32 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 95e3f12ef58..f1c0638bc46 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -19,6 +19,7 @@ Changelog
 - Add SFT-masked data support to ``examples/megatron_bridge/distill.py``: ``--sft --sft_dataset_root <dir>`` distills on raw prompt-completion JSONL (``{"input", "output"}`` records) with the loss masked to the response tokens, using Megatron-Bridge's ``FinetuningDatasetConfig`` and the model's own HuggingFace tokenizer instead of the pretraining ``GPTDataset`` and ``NullTokenizer``.
 - Add per-expert weight quantization for Transformer Engine ``TEGroupedMLP`` (fused MoE experts): each expert now has its own ``weight_quantizer`` (a ``GroupedQuantizer`` holding one ``TensorQuantizer`` per expert) with an independent ``amax``, instead of a single shared ``amax`` across all experts. Applies to ``mtq.quantize`` calibration, HF / Megatron export, and QAD.
 - Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager.
+- Add HuggingFace unified export of quantized Qwen3-VL checkpoints (PTQ or QAD) via ``examples/megatron_bridge/export_quantized_megatron_to_hf.py``. Only the language model is quantized; the vision tower is copied from the source HuggingFace checkpoint. Other VLM architectures (e.g. Qwen3.5-VL, Gemma3-VL) are still saved in Megatron checkpoint format only.
 
 *Misc*
 
@@ -46,6 +47,7 @@ Changelog
 
 - Avoid querying CUDA/Blackwell capability when ``NVFP4QTensor.quantize`` uses its CPU path or has the optional TensorRT-LLM fast path disabled.
 - Fix NVFP4 ONNX export to quantize FP4 weights with the published FP8 block scales, matching eager ModelOpt packed weights. Block scales below ``2**-9`` are now clamped to that minimum, and non-finite or negative scales raise an error.
+- Fix Megatron-Bridge Quantization Aware Distillation of a vision-language model silently discarding the ModelOpt state, so the distilled checkpoint restored no quantizers and exported as an unquantized model. Re-run QAD to regenerate any affected checkpoint.
 - Update HuggingFace checkpoint export to use name-based tied-weight deduplication instead of the previous address-based approach. The address-based deduplication could incorrectly drop an untied weight that happened to share memory with a tied one, producing an incomplete checkpoint (observed as a false positive on MiniMax-M2.7).
 - Fix EAGLE-3 training with context parallelism (``--cp_size > 1`` in ``examples/speculative_decoding``), which failed to start on ``accelerate >= 1.13`` and then raised ``got mixed torch.Tensor and DTensor``.
 - Polygraphy minimum dependency upgraded to ``0.53.4`` to solve ONNX AutoCast failures when marking optional graph outputs.
diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md
index 9c5992ead13..919246bf9d3 100644
--- a/examples/megatron_bridge/README.md
+++ b/examples/megatron_bridge/README.md
@@ -120,7 +120,7 @@ For a vision-language model (e.g. Qwen3.5-VL, Gemma3-VL), `quantize.py` automati
 - A **text** dataset runs text-only calibration of the language model (vision tower idle).
 
 > [!NOTE]
-> HuggingFace unified export (`export_quantized_megatron_to_hf.py`) of a quantized VLM is not yet supported; the quantized VLM is saved in Megatron checkpoint format only.
+> HuggingFace unified export (`export_quantized_megatron_to_hf.py`) of a quantized VLM covers **Qwen3-VL** only -- the architecture must be in ModelOpt's Megatron export mapping. For other VLMs (Qwen3.5-VL, Gemma3-VL) the quantized model is saved in Megatron checkpoint format only.
 
 ## Distillation
 
diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py
index 7a88172f2e2..f206a43b67a 100644
--- a/examples/megatron_bridge/distill.py
+++ b/examples/megatron_bridge/distill.py
@@ -54,6 +54,7 @@
 
 import modelopt.torch.distill as mtd
 import modelopt.torch.utils.distributed as dist
+from modelopt.torch.opt.conversion import ModeloptStateManager
 from modelopt.torch.utils import print_args, print_rank_0, warn_rank_0
 from modelopt.torch.utils.plugins.mbridge import load_modelopt_megatron_checkpoint
 
@@ -420,9 +421,13 @@ def _restore_student_hook(model_chunks):
             print_rank_0(
                 f"Loading student weights from Megatron checkpoint {args.student_megatron_path}"
             )
-            load_modelopt_megatron_checkpoint(
-                [unwrap_model(model_chunks[0])], args.student_megatron_path
-            )
+            student = unwrap_model(model_chunks[0])
+            load_modelopt_megatron_checkpoint([student], args.student_megatron_path)
+            if is_vlm and student_has_modelopt_state:
+                # PTQ stores the state on the VLM root (it quantizes and saves the whole VLM), but
+                # only ``language_model`` is distilled and checkpointed here, so move it there to
+                # keep the quantizers across the QAD checkpoint's save / restore.
+                ModeloptStateManager.transfer_state_dict(student, student.language_model)
             return model_chunks
 
         distill_provider.register_pre_wrap_hook(_restore_student_hook, prepend=True)
diff --git a/examples/megatron_bridge/export_quantized_megatron_to_hf.py b/examples/megatron_bridge/export_quantized_megatron_to_hf.py
index e4e3703d8a5..738014ddd55 100644
--- a/examples/megatron_bridge/export_quantized_megatron_to_hf.py
+++ b/examples/megatron_bridge/export_quantized_megatron_to_hf.py
@@ -144,9 +144,8 @@ def main(args: argparse.Namespace):
     print_rank_0(
         f"Exporting to HuggingFace (unified) checkpoint at {args.export_unified_hf_path}..."
     )
-    # TODO (OMNIML-5366): quantized-VLM HF export. export_mcore_gpt_to_hf's per-arch mappings don't
-    # cover Qwen3.5-VL / Gemma3-VL; See if Megatron-Bridge's AutoBridge.export_hf_weights_quant can be
-    # used instead.
+    # TODO (OMNIML-5366): Qwen3-VL is the only VLM in export_mcore_gpt_to_hf's per-arch mappings;
+    # Qwen3.5-VL / Gemma3-VL are not covered yet.
     export_mcore_gpt_to_hf(
         unwrapped_model,
         args.hf_model_name_or_path,
diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py
index fedd7bf62d0..9b5d3687125 100644
--- a/examples/megatron_bridge/quantize.py
+++ b/examples/megatron_bridge/quantize.py
@@ -416,8 +416,8 @@ def forward_loop(_model=None):
     )
     if is_vlm:
         print_rank_0(
-            f"\nSaved quantized VLM to {args.export_megatron_path} in Megatron format "
-            "(HuggingFace unified export of a quantized VLM is not yet supported)."
+            f"\nSaved quantized VLM to {args.export_megatron_path} in Megatron format. To deploy this "
+            "model, convert it to a Unified HF ckpt with export_quantized_megatron_to_hf.py (Qwen3-VL only)."
         )
     else:
         print_rank_0(
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index bbbd21ad244..e750dd43616 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -133,7 +133,15 @@ def __init__(
         moe_router_dtype: str | None = None,
     ):
         """Create a GPTModel exporter instance."""
-        if not isinstance(model, (GPTModel, HybridModel, LLaVAModel)):
+        # VLM wrappers (MCore ``LLaVAModel``, Megatron-Bridge ``Qwen3VLModel``, ...) keep the decoder
+        # under ``.language_model``; only that inner model is exported, the vision tower is copied
+        # over from the HF checkpoint as-is.
+        language_model = (
+            model
+            if isinstance(model, (GPTModel, HybridModel))
+            else getattr(model, "language_model", None)
+        )
+        if not isinstance(language_model, (GPTModel, HybridModel)):
             raise ValueError("Input to GPTModelExport must be a megatron.core.models.GPTModel!")
 
         self._state_dict = OrderedDict()
@@ -153,19 +161,18 @@ def __init__(
         self._hf_text_config = getattr(self._hf_config, "text_config", self._hf_config)
 
         # Update hf_config
-        self._hf_text_config.num_hidden_layers = model.config.num_layers
-        self._hf_text_config.hidden_size = model.config.hidden_size
-        self._hf_text_config.head_dim = model.config.kv_channels
-        self._hf_text_config.num_attention_heads = model.config.num_attention_heads
-        self._hf_text_config.num_key_value_heads = model.config.num_query_groups
+        self._hf_text_config.num_hidden_layers = language_model.config.num_layers
+        self._hf_text_config.hidden_size = language_model.config.hidden_size
+        self._hf_text_config.head_dim = language_model.config.kv_channels
+        self._hf_text_config.num_attention_heads = language_model.config.num_attention_heads
+        self._hf_text_config.num_key_value_heads = language_model.config.num_query_groups
         self.is_multimodal = isinstance(model, LLaVAModel)
         if not self.is_multimodal:
-            self._hf_text_config.intermediate_size = model.config.ffn_hidden_size
+            self._hf_text_config.intermediate_size = language_model.config.ffn_hidden_size
         self._hf_quant_config: dict = {}
         self._hf_extra_config = None
         self.export_extra_modules = export_extra_modules
-        self.is_multimodal = isinstance(model, LLaVAModel)
-        self.model = model.language_model if self.is_multimodal else model
+        self.model = language_model
         self.dtype = dtype
         self.trust_remote_code = trust_remote_code
         self.arch = self._hf_config.architectures[0]
diff --git a/modelopt/torch/utils/plugins/mbridge.py b/modelopt/torch/utils/plugins/mbridge.py
index c0aa048ebcd..dcd97491a2e 100644
--- a/modelopt/torch/utils/plugins/mbridge.py
+++ b/modelopt/torch/utils/plugins/mbridge.py
@@ -30,6 +30,7 @@
 from megatron.core.models.hybrid.hybrid_model import HybridModel
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.utils import unwrap_model
+from torch.distributed.checkpoint import FileSystemReader
 from transformers import AutoTokenizer
 
 from modelopt.torch.nas.plugins.megatron import get_te_hybrid_stack_spec
@@ -118,6 +119,12 @@ def load_mbridge_model_from_hf(
     return bridge, provider, model, unwrapped_model, tokenizer
 
 
+def _has_vision_model_weights(checkpoint_path: str) -> bool:
+    """Whether a Megatron distributed checkpoint holds a VLM's vision tower (``vision_model.*``)."""
+    metadata = FileSystemReader(checkpoint_path).read_metadata()
+    return any(key.startswith("vision_model.") for key in metadata.state_dict_metadata)
+
+
 def load_modelopt_megatron_checkpoint(
     model: list[MegatronModule], megatron_path: str, restore_modelopt_state: bool = True
 ) -> None:
@@ -130,9 +137,20 @@ def load_modelopt_megatron_checkpoint(
             weights. Set ``False`` to load weights only -- e.g. to reload a full-precision distilled
             student without reconstructing the ``kd_loss`` mode (which would require a teacher model).
     """
+    # _load_model_weights_from_checkpoint does not resolve the latest iter_* directory, so resolve it explicitly
+    checkpoint_path = _get_modelopt_checkpoint_path(megatron_path)
+
+    # ``distill.py`` distills a VLM's language model only, so its checkpoint holds no ``vision_model.*``
+    # weights and must be loaded into ``.language_model`` rather than the full VLM wrapper.
+    unwrapped_model = unwrap_model(model)
+    if any(hasattr(m, "language_model") for m in unwrapped_model) and not _has_vision_model_weights(
+        checkpoint_path
+    ):
+        print_rank_0("Language-model-only checkpoint: loading into the VLM's `.language_model`.")
+        model = [getattr(m, "language_model", m) for m in unwrapped_model]
+
     # Restore the ModelOpt state before loading weights.
     # has_modelopt_state / load_modelopt_state resolves the latest iter_* directory
     if restore_modelopt_state and has_modelopt_state(megatron_path):
         load_modelopt_state(model, megatron_path)
-    # _load_model_weights_from_checkpoint does not resolve the latest iter_* directory, so resolve it explicitly
-    _load_model_weights_from_checkpoint(_get_modelopt_checkpoint_path(megatron_path), model)
+    _load_model_weights_from_checkpoint(checkpoint_path, model)
diff --git a/tests/examples/megatron_bridge/test_qad.py b/tests/examples/megatron_bridge/test_qad.py
index 6128b5ad6ae..7d982669221 100644
--- a/tests/examples/megatron_bridge/test_qad.py
+++ b/tests/examples/megatron_bridge/test_qad.py
@@ -22,14 +22,18 @@
     create_tiny_gemma3vl_dir,
     create_tiny_qwen3_5_moe_vl_dir,
     create_tiny_qwen3_dir,
+    create_tiny_qwen3vl_dir,
 )
 
 
 @pytest.mark.timeout(720)  # Multiple steps in one test hence takes longer than the default timeout
 @pytest.mark.parametrize(
-    ("create_student", "is_vlm", "is_moe"),
+    ("create_student", "exports_hf"),
     [
-        (lambda tmp_path: create_tiny_qwen3_dir(tmp_path, with_tokenizer=True), False, False),
+        (lambda tmp_path: create_tiny_qwen3_dir(tmp_path, with_tokenizer=True), True),
+        # Qwen3-VL is the only VLM architecture in the Megatron HF export mapping, so it is the one
+        # VLM case that runs the export step end-to-end.
+        (lambda tmp_path: create_tiny_qwen3vl_dir(tmp_path, with_tokenizer=True), True),
         # Dense-VLM QAD path; the MoE VLM below covers it in CI, so run this one on demand only.
         pytest.param(
             lambda tmp_path: create_tiny_gemma3vl_dir(
@@ -39,7 +43,6 @@
                 intermediate_size=128,
                 max_position_embeddings=512,
             ),
-            True,
             False,
             marks=pytest.mark.manual,
         ),
@@ -47,19 +50,19 @@
             lambda tmp_path: create_tiny_qwen3_5_moe_vl_dir(
                 tmp_path, with_processor=True, num_hidden_layers=2
             ),
-            True,
-            True,
+            False,
         ),
     ],
-    ids=["qwen3", "gemma3vl", "qwen3_5_moe_vl"],
+    ids=["qwen3", "qwen3vl", "gemma3vl", "qwen3_5_moe_vl"],
 )
-def test_qad(tmp_path: Path, num_gpus, create_student, is_vlm, is_moe):
+def test_qad(tmp_path: Path, num_gpus, create_student, exports_hf):
     """Quantize a tiny model, run QAD from the quantized student, and export the result.
 
     For VLMs only the language model is quantized and distilled (vision tower / projector untouched),
-    and a text calibration dataset infers text-only LM calibration. VLM quantized-HF export is
-    unsupported, so the VLM case stops at the distilled Megatron checkpoint and verifies the ModelOpt
-    (quantize) state survived distillation; the LLM case additionally exports a unified HF checkpoint.
+    and a text calibration dataset infers text-only LM calibration. Quantized-HF export needs the
+    architecture in the Megatron export mapping, so the cases whose architecture is missing there
+    (Gemma3-VL, Qwen3.5-VL) stop at the distilled Megatron checkpoint and only verify the ModelOpt
+    (quantize) state survived distillation.
     """
     hf_model_path = create_student(tmp_path)
     quantized_megatron_path = tmp_path / "quantized_megatron"
@@ -117,8 +120,8 @@ def test_qad(tmp_path: Path, num_gpus, create_student, is_vlm, is_moe):
         "Expected modelopt_state to be preserved in the distilled (QAD) checkpoint"
     )
 
-    if is_vlm:
-        return  # VLM quantized-HF export is unsupported; stop at the distilled Megatron checkpoint
+    if not exports_hf:
+        return  # architecture missing from the export mapping; stop at the distilled checkpoint
 
     # Step 3: export the distilled quantized checkpoint to a unified HF checkpoint. hf_quant_config.json
     # is only written for a quantized model, so its presence confirms the quantizers survived QAD.

From d1f077cbdab4fc567eb6b2916a22b05aab5c59c1 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 28 Aug 2026 07:07:45 -0700
Subject: [PATCH 02/12] Add Qwen3.5-VL export, fail loudly on unexportable MoE,
 verify exports

Builds on the Qwen3-VL enablement: data-drives the VLM export mapping,
adds Qwen3.5-VL, closes a silent expert-drop bug, and makes the tests
check exported content rather than just that files exist.

Export mapping is now table-driven
- Vision-tower passthrough prefixes move into
  all_mcore_hf_vision_passthrough_mapping; the exporter no longer
  branches on the Qwen3-VL architecture string.
- with_language_model_prefix moves to mcore_custom so any VLM mapping
  can be derived from its text-model mapping.

Qwen3.5-VL (MoE)
- GatedDeltaNetSlicing splits the fused in_proj into HF's
  in_proj_qkv/_z/_b/_a, using the module's own split sections.
- Megatron's GDN out_norm is zero-centered; add 1.0 on export, matching
  Megatron-Bridge's RMSNorm2ZeroCenteredRMSNormMapping on import.
- Emit shared_experts.gate_weight.

Silent expert drop
- The MoE dispatch had no else branch, so an architecture without an
  experts.linear_fc1 rule (e.g. Qwen3MoeForCausalLM) exported a valid
  looking checkpoint with zero routed experts. Both quantize.py and the
  exporter now raise, and --no_moe_grouped_gemm is plumbed through
  quantize.py / distill.py / the export script as the way out.

Export verification
- assert_exported_checkpoint_matches compares an exported checkpoint
  against its source: key set, shapes (accounting for NVFP4 uint8
  packing), safetensors index, and values. Wired into the two example
  tests and the unit test; it reproduces both the expert drop and the
  zero-centered-gamma bug above.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 CHANGELOG.rst                                 |   4 +-
 examples/megatron_bridge/distill.py           |  12 ++
 .../export_quantized_megatron_to_hf.py        |   9 ++
 examples/megatron_bridge/quantize.py          |  30 +++-
 modelopt/torch/export/plugins/mcore_common.py |  16 +-
 modelopt/torch/export/plugins/mcore_custom.py |  36 +++++
 .../torch/export/plugins/mcore_qwen35vl.py    |  59 +++++++
 .../torch/export/plugins/mcore_qwen3vl.py     |  40 ++---
 .../torch/export/unified_export_megatron.py   | 141 ++++++++++++++--
 .../torch/export/unified_checkpoint.py        | 150 ++++++++++++++++++
 tests/examples/megatron_bridge/test_qad.py    |  34 ++--
 .../megatron_bridge/test_quantize_export.py   |   3 +-
 .../export/test_unified_export_megatron.py    |  18 +--
 13 files changed, 483 insertions(+), 69 deletions(-)
 create mode 100644 modelopt/torch/export/plugins/mcore_qwen35vl.py
 create mode 100644 tests/_test_utils/torch/export/unified_checkpoint.py

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index f1c0638bc46..71ed7bffa9f 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -19,7 +19,8 @@ Changelog
 - Add SFT-masked data support to ``examples/megatron_bridge/distill.py``: ``--sft --sft_dataset_root <dir>`` distills on raw prompt-completion JSONL (``{"input", "output"}`` records) with the loss masked to the response tokens, using Megatron-Bridge's ``FinetuningDatasetConfig`` and the model's own HuggingFace tokenizer instead of the pretraining ``GPTDataset`` and ``NullTokenizer``.
 - Add per-expert weight quantization for Transformer Engine ``TEGroupedMLP`` (fused MoE experts): each expert now has its own ``weight_quantizer`` (a ``GroupedQuantizer`` holding one ``TensorQuantizer`` per expert) with an independent ``amax``, instead of a single shared ``amax`` across all experts. Applies to ``mtq.quantize`` calibration, HF / Megatron export, and QAD.
 - Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager.
-- Add HuggingFace unified export of quantized Qwen3-VL checkpoints (PTQ or QAD) via ``examples/megatron_bridge/export_quantized_megatron_to_hf.py``. Only the language model is quantized; the vision tower is copied from the source HuggingFace checkpoint. Other VLM architectures (e.g. Qwen3.5-VL, Gemma3-VL) are still saved in Megatron checkpoint format only.
+- Add HuggingFace unified export of quantized Qwen3-VL and Qwen3.5-VL checkpoints (PTQ or QAD) via ``examples/megatron_bridge/export_quantized_megatron_to_hf.py``. Only the language model is quantized; the vision tower is copied from the source HuggingFace checkpoint. Qwen3.5-VL additionally covers GatedDeltaNet linear-attention layers and MoE shared experts, and requires ``--no_moe_grouped_gemm``. Gemma3-VL is still saved in Megatron checkpoint format only.
+- Add ``--no_moe_grouped_gemm`` to ``examples/megatron_bridge/quantize.py``, ``distill.py`` and ``export_quantized_megatron_to_hf.py``, building MoE experts as ``SequentialMLP`` instead of the fused ``TEGroupedMLP``. The flag must match across all three, since the expert layout is baked into the Megatron checkpoint.
 
 *Misc*
 
@@ -48,6 +49,7 @@ Changelog
 - Avoid querying CUDA/Blackwell capability when ``NVFP4QTensor.quantize`` uses its CPU path or has the optional TensorRT-LLM fast path disabled.
 - Fix NVFP4 ONNX export to quantize FP4 weights with the published FP8 block scales, matching eager ModelOpt packed weights. Block scales below ``2**-9`` are now clamped to that minimum, and non-finite or negative scales raise an error.
 - Fix Megatron-Bridge Quantization Aware Distillation of a vision-language model silently discarding the ModelOpt state, so the distilled checkpoint restored no quantizers and exported as an unquantized model. Re-run QAD to regenerate any affected checkpoint.
+- Fix Megatron-Core HuggingFace export silently omitting fused (grouped GEMM) MoE experts for architectures without an ``experts.linear_fc1`` rule (e.g. ``Qwen3MoeForCausalLM``), which produced a valid-looking checkpoint containing no expert weights. ``quantize.py`` and the exporter now raise instead; pass ``--no_moe_grouped_gemm`` to export these models.
 - Update HuggingFace checkpoint export to use name-based tied-weight deduplication instead of the previous address-based approach. The address-based deduplication could incorrectly drop an untied weight that happened to share memory with a tied one, producing an incomplete checkpoint (observed as a false positive on MiniMax-M2.7).
 - Fix EAGLE-3 training with context parallelism (``--cp_size > 1`` in ``examples/speculative_decoding``), which failed to start on ``accelerate >= 1.13`` and then raised ``got mixed torch.Tensor and DTensor``.
 - Polygraphy minimum dependency upgraded to ``0.53.4`` to solve ONNX AutoCast failures when marking optional graph outputs.
diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py
index f206a43b67a..80e60546cf6 100644
--- a/examples/megatron_bridge/distill.py
+++ b/examples/megatron_bridge/distill.py
@@ -93,6 +93,15 @@ def get_args():
         help="HuggingFace model name or path for the teacher (e.g. Qwen/Qwen3-8B)",
     )
     parser.add_argument("--trust_remote_code", action="store_true", help="Trust remote code")
+    parser.add_argument(
+        "--no_moe_grouped_gemm",
+        action="store_true",
+        help=(
+            "Use SequentialMLP for MoE experts instead of the (default) efficient fused "
+            "TEGroupedMLP (grouped GEMM). Must match the checkpoint passed to "
+            "--student_megatron_path. Only affects MoE models."
+        ),
+    )
     parser.add_argument(
         "--student_megatron_path",
         type=str,
@@ -344,6 +353,9 @@ def _build_model_provider(hf_path, load_weights=True):
         provider.expert_model_parallel_size = args.ep_size
         provider.expert_tensor_parallel_size = 1  # Expert tensor parallelism is not supported
         provider.seq_length = args.seq_length
+        if (provider.num_moe_experts or 0) > 0:
+            # Must match the expert layout of --student_megatron_path (see quantize.py).
+            provider.moe_grouped_gemm = not args.no_moe_grouped_gemm
         if args.sft:
             # A response-only loss mask needs per-token reduction to combine across CP ranks.
             # Must stay in sync with ``average_in_collective=not args.sft`` on the DDP config.
diff --git a/examples/megatron_bridge/export_quantized_megatron_to_hf.py b/examples/megatron_bridge/export_quantized_megatron_to_hf.py
index 738014ddd55..018b13686a6 100644
--- a/examples/megatron_bridge/export_quantized_megatron_to_hf.py
+++ b/examples/megatron_bridge/export_quantized_megatron_to_hf.py
@@ -71,6 +71,14 @@ def get_args() -> argparse.Namespace:
         help="Directory to write the exported HuggingFace (unified) checkpoint to.",
     )
     parser.add_argument("--trust_remote_code", action="store_true")
+    parser.add_argument(
+        "--no_moe_grouped_gemm",
+        action="store_true",
+        help=(
+            "Use SequentialMLP for MoE experts instead of the (default) efficient fused "
+            "TEGroupedMLP (grouped GEMM). Only affects MoE models."
+        ),
+    )
     parser.add_argument(
         "--export_extra_modules",
         action="store_true",
@@ -108,6 +116,7 @@ def main(args: argparse.Namespace):
     _bridge, _provider, model, _unwrapped_model, _tokenizer = load_mbridge_model_from_hf(
         hf_model_name_or_path=args.hf_model_name_or_path,
         trust_remote_code=trust_remote_code,
+        moe_grouped_gemm=not args.no_moe_grouped_gemm,
         provider_overrides={
             "tensor_model_parallel_size": 1,  # Tensor parallelism is not supported
             "pipeline_model_parallel_size": args.pp_size,
diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py
index 9b5d3687125..27dd8ad630c 100644
--- a/examples/megatron_bridge/quantize.py
+++ b/examples/megatron_bridge/quantize.py
@@ -62,12 +62,13 @@
 import gc
 
 import torch
-from transformers import AutoProcessor
+from transformers import AutoConfig, AutoProcessor
 
 import modelopt.torch.quantization as mtq
 import modelopt.torch.utils.distributed as dist
 from modelopt.recipe import ModelOptPTQRecipe, load_recipe
 from modelopt.recipe.presets import KV_CACHE_NONE, KV_QUANT_CFG_CHOICES, QUANT_CFG_CHOICES
+from modelopt.torch.export.plugins.mcore_common import all_mcore_hf_export_mapping
 from modelopt.torch.utils import print_args, print_rank_0, warn_rank_0
 from modelopt.torch.utils.dataset_utils import get_supported_datasets
 from modelopt.torch.utils.plugins.mbridge import load_mbridge_model_from_hf
@@ -96,6 +97,14 @@ def get_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     parser.add_argument("--hf_model_name_or_path", type=str, required=True)
     parser.add_argument("--trust_remote_code", action="store_true")
+    parser.add_argument(
+        "--no_moe_grouped_gemm",
+        action="store_true",
+        help=(
+            "Use SequentialMLP for MoE experts instead of the (default) efficient fused "
+            "TEGroupedMLP (grouped GEMM). Only affects MoE models."
+        ),
+    )
     parser.add_argument(
         "--export_megatron_path",
         type=str,
@@ -279,6 +288,7 @@ def main(args: argparse.Namespace):
     bridge, _provider, model, unwrapped_model, tokenizer = load_mbridge_model_from_hf(
         hf_model_name_or_path=args.hf_model_name_or_path,
         trust_remote_code=args.trust_remote_code,
+        moe_grouped_gemm=not args.no_moe_grouped_gemm,
         provider_overrides={
             "tensor_model_parallel_size": args.tp_size,
             "pipeline_model_parallel_size": args.pp_size,
@@ -321,6 +331,24 @@ def main(args: argparse.Namespace):
         )
     print_rank_0(f"Using calibration dataset: {args.calib_dataset_name}")
 
+    # Fused (grouped GEMM) experts are only exportable for architectures with an
+    # "experts.linear_fc1" rule. Fail now rather than at export: the layout is baked into the
+    # checkpoint, so recovering means re-running this whole calibration.
+    if not args.no_moe_grouped_gemm and any(
+        hasattr(m, "experts") and not hasattr(m.experts, "local_experts")
+        for m in unwrapped_model.modules()
+        if type(m).__name__.endswith("MoELayer")
+    ):
+        arch = AutoConfig.from_pretrained(
+            args.hf_model_name_or_path, trust_remote_code=args.trust_remote_code
+        ).architectures[0]
+        if "experts.linear_fc1" not in all_mcore_hf_export_mapping.get(arch, {}):
+            raise NotImplementedError(
+                f"{arch} has fused (grouped GEMM) MoE experts, which "
+                "export_quantized_megatron_to_hf.py cannot export. Re-run with "
+                "--no_moe_grouped_gemm to build the experts as SequentialMLP instead."
+            )
+
     mtq_config = get_quant_config(args)
 
     # Quantize only the language model: disable quantizers on every top-level submodule that is not
diff --git a/modelopt/torch/export/plugins/mcore_common.py b/modelopt/torch/export/plugins/mcore_common.py
index 15395b7a1e5..4325a25f54e 100644
--- a/modelopt/torch/export/plugins/mcore_common.py
+++ b/modelopt/torch/export/plugins/mcore_common.py
@@ -39,7 +39,12 @@
     qwen25_causal_lm_export,
     qwen25_causal_lm_import,
 )
-from .mcore_qwen3vl import qwen3vl_causal_lm_export, qwen3vl_causal_lm_import
+from .mcore_qwen3vl import (
+    QWEN3VL_VISION_PREFIXES,
+    qwen3vl_causal_lm_export,
+    qwen3vl_causal_lm_import,
+)
+from .mcore_qwen35vl import QWEN3_5_VL_VISION_PREFIXES, qwen3_5_vl_causal_lm_export
 
 all_mcore_hf_export_mapping: dict[str, Any] = {
     "DeepseekV2ForCausalLM": deepseek_causal_lm_export,
@@ -56,6 +61,15 @@
     "Qwen2ForCausalLM": qwen25_causal_lm_export,
     "GptOssForCausalLM": gptoss_causal_lm_export,
     "Qwen3VLForConditionalGeneration": qwen3vl_causal_lm_export,
+    "Qwen3_5MoeForConditionalGeneration": qwen3_5_vl_causal_lm_export,
+}
+
+# VLM architectures whose Megatron export covers the language model only: the vision tower is copied
+# verbatim from the HF checkpoint under these key prefixes. Architectures absent here fall back to
+# ``LLAVA_VISION_PREFIXES`` when the model is an MCore ``LLaVAModel``, and to no copy otherwise.
+all_mcore_hf_vision_passthrough_mapping: dict[str, tuple[str, ...]] = {
+    "Qwen3VLForConditionalGeneration": QWEN3VL_VISION_PREFIXES,
+    "Qwen3_5MoeForConditionalGeneration": QWEN3_5_VL_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 f231da48433..014744f248c 100644
--- a/modelopt/torch/export/plugins/mcore_custom.py
+++ b/modelopt/torch/export/plugins/mcore_custom.py
@@ -16,6 +16,7 @@
 
 """Custom Megatron mapping and safetensors utility."""
 
+import copy
 import json
 import math
 import os
@@ -175,6 +176,18 @@ def __init__(self, target_name_or_prefix: str = "", func_kwargs: dict[str, Any]
         )
 
 
+class GatedDeltaNetSlicing(CustomModuleMapping):
+    """A custom module mapping that splits GatedDeltaNet's fused ``in_proj``."""
+
+    def __init__(self, target_name_or_prefix: str = "", func_kwargs: dict[str, Any] = {}):
+        """Create a custom module mapping that splits the fused GatedDeltaNet input projection."""
+        super().__init__(
+            func_name="gated_delta_net_slicing",
+            target_name_or_prefix=target_name_or_prefix,
+            func_kwargs=func_kwargs,
+        )
+
+
 class PackNameRemapping(CustomModuleMapping):
     """A custom module mapping that packs module after name remapping."""
 
@@ -223,6 +236,29 @@ def __init__(self, target_name_or_prefix: str = "", func_kwargs: dict[str, Any]
         )
 
 
+# LLaVA-style checkpoints keep the vision tower under these prefixes; see
+# ``all_mcore_hf_vision_passthrough_mapping`` for the per-architecture overrides.
+LLAVA_VISION_PREFIXES = ("multi_modal_projector", "vision_model")
+
+
+def with_language_model_prefix(
+    mapping: dict[str, CustomModuleMapping],
+) -> dict[str, CustomModuleMapping]:
+    """Derive a VLM mapping from a text-model mapping by nesting it under ``model.language_model.``.
+
+    Prefixes that do not start with ``model.`` (e.g. ``lm_head.``) are left unchanged.
+    """
+    result = {}
+    for key, m in mapping.items():
+        prefix = m.target_name_or_prefix
+        if prefix.startswith("model."):
+            prefix = "model.language_model." + prefix[len("model.") :]
+        result[key] = type(m)(
+            target_name_or_prefix=prefix, func_kwargs=copy.deepcopy(m.func_kwargs)
+        )
+    return result
+
+
 def save_safetensors(state_dict, save_directory: str | os.PathLike):
     """Save safetensors with pipeline model parallel support."""
     pp_rank = get_pipeline_model_parallel_rank()
diff --git a/modelopt/torch/export/plugins/mcore_qwen35vl.py b/modelopt/torch/export/plugins/mcore_qwen35vl.py
new file mode 100644
index 00000000000..424e5142e62
--- /dev/null
+++ b/modelopt/torch/export/plugins/mcore_qwen35vl.py
@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023-2025 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 Qwen3.5-VL Hugging Face models to Megatron Core models.
+
+Qwen3.5 interleaves GatedDeltaNet linear-attention layers -- Megatron keeps them under
+``self_attention`` with a fused ``in_proj`` that HF splits into ``in_proj_qkv`` / ``_z`` / ``_b`` /
+``_a`` -- with gated full-attention layers, and adds shared experts to the Qwen3 MoE rules. As for
+Qwen3-VL, only the language model is exported; the vision tower is copied from the HF checkpoint.
+
+Requires ``--no_moe_grouped_gemm`` (SequentialMLP experts); fused ``TEGroupedMLP`` experts have no
+gated split in the exporter. Gated full-attention layers are untested.
+"""
+
+from .mcore_custom import (
+    GatedDeltaNetSlicing,
+    GatedMLPSlicing,
+    NameRemapping,
+    with_language_model_prefix,
+)
+from .mcore_qwen import qwen3_causal_lm_export
+
+# Vision-tower weights copied straight from the HF checkpoint (never quantized).
+QWEN3_5_VL_VISION_PREFIXES = ("model.visual.",)
+
+# Qwen3.5 adds linear-attention layers and shared experts on top of the Qwen3 rules.
+_qwen3_5_extra_export: dict = {
+    # Linear attention (GatedDeltaNet). ``linear_attn`` splits the fused in_proj; the rest
+    # are plain renames of the surrounding parameters.
+    "linear_attn": GatedDeltaNetSlicing("model.layers.{}.linear_attn."),
+    "linear_attn.conv1d": NameRemapping("model.layers.{}.linear_attn.conv1d."),
+    "linear_attn.A_log": NameRemapping("model.layers.{}.linear_attn.A_log"),
+    "linear_attn.dt_bias": NameRemapping("model.layers.{}.linear_attn.dt_bias"),
+    # Megatron's GDN output norm is zero-centered; HF's RMSNorm gamma is centered on 1.
+    "linear_attn.out_norm": NameRemapping(
+        "model.layers.{}.linear_attn.norm.", {"zero_centered_gamma": True}
+    ),
+    "linear_attn.out_proj": NameRemapping("model.layers.{}.linear_attn.out_proj."),
+    # MoE shared experts (routed experts + router come from the Qwen3 rules).
+    "shared_experts.linear_fc1": GatedMLPSlicing("model.layers.{}.mlp.shared_expert."),
+    "shared_experts.linear_fc2": NameRemapping("model.layers.{}.mlp.shared_expert.down_proj."),
+    "shared_experts.gate_weight": NameRemapping("model.layers.{}.mlp.shared_expert_gate.weight"),
+}
+
+qwen3_5_vl_causal_lm_export = with_language_model_prefix(
+    {**qwen3_causal_lm_export, **_qwen3_5_extra_export}
+)
diff --git a/modelopt/torch/export/plugins/mcore_qwen3vl.py b/modelopt/torch/export/plugins/mcore_qwen3vl.py
index 1f2d3830d61..923e2d19c38 100644
--- a/modelopt/torch/export/plugins/mcore_qwen3vl.py
+++ b/modelopt/torch/export/plugins/mcore_qwen3vl.py
@@ -17,12 +17,12 @@
 
 Qwen3-VL differs from Qwen3 in one structural way: language-model weights live
 under ``model.language_model.`` instead of ``model.``, while ``lm_head.weight``
-remains at the root level.  The mappings below are derived automatically from
-the Qwen3 mappings by inserting ``language_model.`` after ``model.`` for every
-prefix that starts with ``model.``.
+remains at the root level, so the mappings below are derived from the Qwen3 ones
+with :func:`with_language_model_prefix`.
 
-Note: the visual encoder (``model.visual.*``) is intentionally excluded — this
-mapping covers only the language-model decoder used for quantization and export.
+The visual encoder (``model.visual.*``) is not mapped: only the language model is
+quantized, and the vision tower is copied verbatim from the Hugging Face checkpoint
+via ``QWEN3VL_VISION_PREFIXES``.
 
 Note: ``Qwen3VLMoeForConditionalGeneration`` is **not** supported here.  The MoE
 variant stores expert weights as 3-D tensors (``mlp.experts.gate_up_proj``,
@@ -32,31 +32,11 @@
 Reference: https://huggingface.co/Qwen/Qwen3-VL-8B-Instruct/blob/main/model.safetensors.index.json
 """
 
-import copy
-
-from .mcore_custom import CustomModuleMapping
+from .mcore_custom import with_language_model_prefix
 from .mcore_qwen import qwen3_causal_lm_export, qwen3_causal_lm_import
 
+# Vision-tower weights copied straight from the HF checkpoint (never quantized).
+QWEN3VL_VISION_PREFIXES = ("model.visual.",)
 
-def _with_language_model_prefix(
-    mapping: dict[str, CustomModuleMapping],
-) -> dict[str, CustomModuleMapping]:
-    """Derive a VL mapping from a base Qwen3 mapping.
-
-    Rewrites every ``target_name_or_prefix`` that starts with ``model.`` to
-    ``model.language_model.<rest>``.  Prefixes that do not start with
-    ``model.`` (e.g. ``lm_head.``) are left unchanged.
-    """
-    result = {}
-    for key, m in mapping.items():
-        prefix = m.target_name_or_prefix
-        if prefix.startswith("model."):
-            prefix = "model.language_model." + prefix[len("model.") :]
-        result[key] = type(m)(
-            target_name_or_prefix=prefix, func_kwargs=copy.deepcopy(m.func_kwargs)
-        )
-    return result
-
-
-qwen3vl_causal_lm_import = _with_language_model_prefix(qwen3_causal_lm_import)
-qwen3vl_causal_lm_export = _with_language_model_prefix(qwen3_causal_lm_export)
+qwen3vl_causal_lm_import = with_language_model_prefix(qwen3_causal_lm_import)
+qwen3vl_causal_lm_export = with_language_model_prefix(qwen3_causal_lm_export)
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index e750dd43616..7b832a0a2de 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -53,8 +53,12 @@
     copy_non_safetensor_files_from_ckpt,
     load_multimodal_components,
 )
-from .plugins.mcore_common import all_mcore_hf_export_mapping
+from .plugins.mcore_common import (
+    all_mcore_hf_export_mapping,
+    all_mcore_hf_vision_passthrough_mapping,
+)
 from .plugins.mcore_custom import (
+    LLAVA_VISION_PREFIXES,
     CustomModuleMapping,
     get_safetensor,
     save_safetensors_by_layer_index,
@@ -176,6 +180,11 @@ def __init__(
         self.dtype = dtype
         self.trust_remote_code = trust_remote_code
         self.arch = self._hf_config.architectures[0]
+        # A VLM's vision tower is never quantized: copy it verbatim from the HF checkpoint. ``None``
+        # means there is nothing to copy.
+        self.vision_passthrough_prefixes = all_mcore_hf_vision_passthrough_mapping.get(
+            self.arch, LLAVA_VISION_PREFIXES if self.is_multimodal else None
+        )
         # TODO: May modify this later according to what quantization exported ckpt is, currently only support BF16.
         if self.arch == "GptOssForCausalLM":
             if hasattr(self._hf_config, "quantization_config"):
@@ -400,14 +409,13 @@ def save_pretrained(
             # Merge the multimodal components into that shard so they land in a file
             # the index builder picks up (it scans shards 1..num_layers).
             first_layer_key = next(iter(layer_state_dicts))
-            if self.is_multimodal:
-                multimodal_state_dict = load_multimodal_components(pretrained_model_name_or_path)
-                layer_state_dicts[first_layer_key].update(multimodal_state_dict)
-            elif self.arch == "Qwen3VLForConditionalGeneration":
-                vision_state_dict = load_multimodal_components(
-                    pretrained_model_name_or_path, prefixes=("model.visual.",)
+            if self.vision_passthrough_prefixes is not None:
+                layer_state_dicts[first_layer_key].update(
+                    load_multimodal_components(
+                        pretrained_model_name_or_path,
+                        prefixes=self.vision_passthrough_prefixes,
+                    )
                 )
-                layer_state_dicts[first_layer_key].update(vision_state_dict)
 
         # Bracket the writer's config.json read-modify-write with barriers so peers
         # never observe a truncated file (also ensures export_dir exists).
@@ -504,8 +512,12 @@ def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False):
         if not isinstance(layer.input_layernorm, IdentityOp):
             self.rules["input_layernorm"](layer.input_layernorm, layer_id, is_mtp=is_mtp)
         else:
+            # GatedDeltaNet fuses the input layernorm into ``in_proj`` rather than ``linear_qkv``.
+            qkv_module = getattr(layer.self_attention, "linear_qkv", None)
+            if qkv_module is None:
+                qkv_module = getattr(layer.self_attention, "in_proj", None)
             fused_key, norm_weight = self._get_fused_norm_weight(
-                getattr(layer.self_attention, "linear_qkv", None),
+                qkv_module,
                 primary_key="fused_input_layernorm",
             )
             if norm_weight is not None:
@@ -538,6 +550,9 @@ def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False):
                     layer.self_attention.linear_kv_up_proj, layer_id, is_mtp=is_mtp
                 )
                 self.rules["linear_proj"](layer.self_attention.linear_proj, layer_id, is_mtp=is_mtp)
+            elif "linear_attn" in self.rules and hasattr(layer.self_attention, "in_proj"):
+                # GatedDeltaNet (Qwen3.5 linear attention): no q/k layernorm, no core_attention.
+                self._get_gated_delta_net_state_dict(layer, layer_id, is_mtp=is_mtp)
             else:
                 if layer.self_attention.q_layernorm is not None and not isinstance(
                     layer.self_attention.q_layernorm, (IdentityOp, L2Norm)
@@ -592,6 +607,13 @@ def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False):
                     self.rules["shared_experts.linear_fc2"](
                         layer.mlp.shared_experts.linear_fc2, layer_id, is_mtp=is_mtp
                     )
+                    if (
+                        "shared_experts.gate_weight" in self.rules
+                        and getattr(layer.mlp.shared_experts, "gate_weight", None) is not None
+                    ):
+                        self.rules["shared_experts.gate_weight"](
+                            layer.mlp.shared_experts.gate_weight, layer_id, is_mtp=is_mtp
+                        )
                 if hasattr(layer.mlp.experts, "local_experts"):
                     if not self.rules.get("use_packed_local_experts", False):
                         for expert_id, expert in enumerate(layer.mlp.experts.local_experts):
@@ -620,6 +642,15 @@ def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False):
                     self.rules["experts.linear_fc2"](
                         layer.mlp.experts.linear_fc2, layer_id, is_mtp=is_mtp
                     )
+                else:
+                    # Without this the routed experts are silently dropped and the exported
+                    # checkpoint looks valid but has no expert weights.
+                    raise NotImplementedError(
+                        f"No export rule for {type(layer.mlp.experts).__name__} experts of "
+                        f"{self.arch}: fused (grouped GEMM) experts need an 'experts.linear_fc1' "
+                        "rule. Re-run quantization and export with --no_moe_grouped_gemm to build "
+                        "the experts as SequentialMLP instead."
+                    )
             else:
                 self.rules["linear_fc1"](layer.mlp.linear_fc1, layer_id, is_mtp=is_mtp)
                 self.rules["linear_fc2"](layer.mlp.linear_fc2, layer_id, is_mtp=is_mtp)
@@ -735,6 +766,16 @@ def _copy_mtp_state_dict_from_pretrained(self) -> dict[str, torch.Tensor]:
             self.exclude_modules.append("mtp*")
         return mtp_state_dict
 
+    def _get_gated_delta_net_state_dict(self, layer, layer_id, is_mtp=False):
+        """Export a GatedDeltaNet (Qwen3.5 linear-attention) layer's ``self_attention``."""
+        gdn = layer.self_attention
+        self.rules["linear_attn"](gdn, layer_id, is_mtp=is_mtp)
+        self.rules["linear_attn.conv1d"](gdn.conv1d, layer_id, is_mtp=is_mtp)
+        self.rules["linear_attn.A_log"](gdn.A_log, layer_id, is_mtp=is_mtp)
+        self.rules["linear_attn.dt_bias"](gdn.dt_bias, layer_id, is_mtp=is_mtp)
+        self.rules["linear_attn.out_norm"](gdn.out_norm, layer_id, is_mtp=is_mtp)
+        self.rules["linear_attn.out_proj"](gdn.out_proj, layer_id, is_mtp=is_mtp)
+
     def _get_mamba_layer_state_dict(self, layer, layer_id, is_mtp=False):
         if not isinstance(layer.norm, IdentityOp):
             self.rules["norm"](layer.norm, layer_id, is_mtp=is_mtp)
@@ -870,6 +911,7 @@ def _custom_mapping_to_lambda(mapping):
                 "qkv_slicing": self._qkv_slicing,
                 "self_attention_scaling": self._self_attention_scaling,
                 "gated_mlp_slicing": self._gated_mlp_slicing,
+                "gated_delta_net_slicing": self._gated_delta_net_slicing,
                 "grouped_mlp_slicing": self._grouped_mlp_slicing,
                 "pack_name_remapping": self._pack_name_remapping,
                 "pack_name_remapping_gpt_oss": self._pack_name_remapping_gpt_oss,
@@ -1028,6 +1070,7 @@ def _name_remapping(
         mapping={},
         dtype: torch.dtype | None = None,
         is_mtp: bool = False,
+        zero_centered_gamma: bool = False,
     ):
         if is_mtp:
             prefix = self._mtp_prefix(prefix)
@@ -1035,13 +1078,16 @@ def _name_remapping(
             dtype = self.dtype
 
         if isinstance(module, torch.Tensor):
-            self._state_dict[prefix] = module
+            self._state_dict[prefix] = (module + 1.0) if zero_centered_gamma else module
             return
 
         name_to_value, qformat, block_size = self._get_quantized_state(module, dtype, prefix=prefix)
         self._record_layer_quant_config(prefix, qformat, block_size)
 
         weight = name_to_value.pop("weight")
+        if zero_centered_gamma:
+            # Megatron stores this norm's gamma centered on 0; HF centers it on 1.
+            weight = weight + 1.0
         weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat)
 
         if weight_scale is None:
@@ -1448,6 +1494,81 @@ def _qkv_slicing(
                 self._state_dict[k_proj_key] = val.detach().clone()
                 self._state_dict[v_proj_key] = val.detach().clone()
 
+    def _gated_delta_net_slicing(self, module, prefix, is_mtp=False):
+        """Split GatedDeltaNet's fused ``in_proj`` into the four HF projections.
+
+        Megatron-Core packs ``[query, key, value, z, beta, alpha]`` along dim 0 of a single
+        ``in_proj``; HF stores ``in_proj_qkv`` (query+key+value), ``in_proj_z``, ``in_proj_b``
+        (beta) and ``in_proj_a`` (alpha).  The sections are contiguous and in that order, so
+        the split is a plain ``torch.split`` -- sizes come from the module itself so TP
+        sharding is handled without re-deriving them here.
+        """
+        if is_mtp:
+            prefix = self._mtp_prefix(prefix)
+        in_proj = module.in_proj
+        name_to_value, qformat, block_size = self._get_quantized_state(
+            in_proj, self.dtype, prefix=prefix
+        )
+
+        sections = dict(zip(module.in_proj_split_names, module.in_proj_split_sections))
+        split_sizes = [
+            sections["query"] + sections["key"] + sections["value"],
+            sections["z"],
+            sections["beta"],
+            sections["alpha"],
+        ]
+        proj_prefixes = [
+            prefix + name + "." for name in ("in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a")
+        ]
+
+        for proj_prefix in proj_prefixes:
+            self._record_layer_quant_config(proj_prefix, qformat, block_size)
+        if qformat in (None, QUANTIZATION_NONE):
+            # Split the fused in_proj exclude entry into the per-HF-name projections.
+            self.exclude_modules = [
+                m for m in self.exclude_modules if m != prefix.removesuffix(".")
+            ]
+            for proj_prefix in proj_prefixes:
+                self._record_excluded_module(proj_prefix)
+
+        weight = name_to_value.pop("weight")
+        proj_weights = list(torch.split(weight, split_sizes, dim=0))
+        proj_keys = [p + "weight" for p in proj_prefixes]
+        weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat)
+
+        if weight_scale is None:
+            for key, proj_weight in zip(proj_keys, proj_weights):
+                self._state_dict[key] = proj_weight
+        else:
+            if len(weight_scale.shape) > 0:
+                # Per-channel / per-block scales are laid out along the same (output) dim.
+                proj_scales = list(torch.split(weight_scale, split_sizes, dim=0))
+            else:
+                proj_scales = [weight_scale.detach().clone() for _ in proj_keys]
+            for proj_weight, scale, key in zip(proj_weights, proj_scales, proj_keys):
+                self._state_dict[key] = to_quantized_weight(
+                    proj_weight, scale, qformat, weight_scale_2, block_size
+                )
+                self._state_dict[key + "_scale"] = scale
+
+        if weight_scale_2 is not None:
+            if len(weight_scale_2.shape) > 0:
+                raise ValueError("weight_scale_2 must be a scalar!")
+            for key in proj_keys:
+                self._state_dict[key + "_scale_2"] = weight_scale_2.detach().clone()
+
+        # weight and weight_scale have been popped; the rest (bias, input_scale, ...) is
+        # either split like the weight or replicated onto every projection.
+        for key, val in name_to_value.items():
+            if key == "bias":
+                for proj_bias, proj_prefix in zip(
+                    torch.split(val.detach().clone(), split_sizes, dim=0), proj_prefixes
+                ):
+                    self._state_dict[proj_prefix + key] = proj_bias
+            else:
+                for proj_prefix in proj_prefixes:
+                    self._state_dict[proj_prefix + key] = val.detach().clone()
+
     def _self_attention_scaling(
         self, module, prefix, k_scale_name="k_scale", v_scale_name="v_scale", is_mtp=False
     ):
diff --git a/tests/_test_utils/torch/export/unified_checkpoint.py b/tests/_test_utils/torch/export/unified_checkpoint.py
new file mode 100644
index 00000000000..bfe094c798f
--- /dev/null
+++ b/tests/_test_utils/torch/export/unified_checkpoint.py
@@ -0,0 +1,150 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+"""Content checks for an exported unified HuggingFace checkpoint.
+
+Existence checks (``config.json`` is there, some safetensors were written) pass even when the
+exporter silently drops a whole module family, so these helpers compare the exported tensors
+against the Hugging Face checkpoint the model came from.
+"""
+
+import json
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+
+__all__ = [
+    "assert_exported_checkpoint_matches",
+    "assert_safetensors_index_consistent",
+    "load_safetensors_dir",
+]
+
+# Per-tensor / per-channel scales the exporter adds for quantized weights. These have no
+# counterpart in the (unquantized) reference checkpoint.
+QUANT_SUFFIXES = (
+    "weight_scale",
+    "weight_scale_2",
+    "input_scale",
+    "output_scale",
+    "k_scale",
+    "v_scale",
+)
+
+
+def load_safetensors_dir(path: Path | str) -> dict[str, torch.Tensor]:
+    """Load every safetensors shard under ``path`` into one dict."""
+    state_dict: dict[str, torch.Tensor] = {}
+    for shard in sorted(Path(path).glob("*.safetensors")):
+        state_dict.update(load_file(str(shard)))
+    assert state_dict, f"No safetensors tensors found in {path}"
+    return state_dict
+
+
+def assert_safetensors_index_consistent(export_dir: Path | str) -> None:
+    """Assert ``model.safetensors.index.json`` matches the shards actually written."""
+    export_dir = Path(export_dir)
+    index_file = export_dir / "model.safetensors.index.json"
+    if not index_file.exists():  # single unsharded file: nothing to cross-check
+        return
+    weight_map = json.loads(index_file.read_text())["weight_map"]
+    missing_files = {f for f in set(weight_map.values()) if not (export_dir / f).exists()}
+    assert not missing_files, f"index.json references missing shards: {sorted(missing_files)}"
+    exported = set(load_safetensors_dir(export_dir))
+    assert set(weight_map) == exported, (
+        f"index.json disagrees with the shards: {sorted(set(weight_map) - exported)[:10]} only in "
+        f"index, {sorted(exported - set(weight_map))[:10]} only in shards"
+    )
+
+
+def _is_packed(tensor: torch.Tensor) -> bool:
+    """Whether a tensor holds sub-byte weights packed two-per-``uint8`` (NVFP4 / INT4)."""
+    return tensor.dtype == torch.uint8
+
+
+def _expected_shape(exported: torch.Tensor, reference: torch.Tensor) -> tuple[int, ...]:
+    """Reference shape, with the last dim halved when the export packs two 4-bit values per byte."""
+    shape = tuple(reference.shape)
+    if _is_packed(exported) and shape:
+        return (*shape[:-1], shape[-1] // 2)
+    return shape
+
+
+def _dequantize(key: str, exported: dict[str, torch.Tensor]) -> torch.Tensor:
+    """Undo per-tensor / per-channel weight scaling so the value can be compared to the source."""
+    weight = exported[key].to(torch.float32)
+    scale = exported.get(key.replace(".weight", ".weight_scale"))
+    if scale is None:
+        return weight
+    scale = scale.to(torch.float32)
+    return weight * (scale if scale.ndim == 0 else scale.reshape(-1, *([1] * (weight.ndim - 1))))
+
+
+def assert_exported_checkpoint_matches(
+    export_dir: Path | str,
+    ref_hf_dir: Path | str,
+    *,
+    allow_missing: tuple[str, ...] = (),
+    check_values: bool = True,
+    rtol: float = 0.15,
+) -> None:
+    """Assert an exported unified HF checkpoint reproduces the reference model it came from.
+
+    Args:
+        export_dir: Directory holding the exported unified HF checkpoint.
+        ref_hf_dir: The HuggingFace checkpoint the Megatron model was built from.
+        allow_missing: Substrings of reference keys the export is expected to omit.
+        check_values: Compare tensor values, not just names and shapes. Set ``False`` when the
+            Megatron weights are random rather than loaded from ``ref_hf_dir``.
+        rtol: Max relative error for quantized tensors (per-tensor FP8 lands well inside 0.15).
+    """
+    exported = load_safetensors_dir(export_dir)
+    reference = load_safetensors_dir(ref_hf_dir)
+    assert_safetensors_index_consistent(export_dir)
+
+    missing = {k for k in set(reference) - set(exported) if not any(a in k for a in allow_missing)}
+    assert not missing, (
+        f"{len(missing)} reference tensor(s) absent from the export, e.g. {sorted(missing)[:8]}"
+    )
+
+    # Anything extra must be a quantization scale; a stray weight means a mis-named rule.
+    unexpected = {k for k in set(exported) - set(reference) if not k.endswith(QUANT_SUFFIXES)}
+    assert not unexpected, f"Export produced unexpected tensors: {sorted(unexpected)[:8]}"
+
+    shared = sorted(set(exported) & set(reference))
+    mismatched = [
+        (k, tuple(exported[k].shape), tuple(reference[k].shape))
+        for k in shared
+        if tuple(exported[k].shape) != _expected_shape(exported[k], reference[k])
+    ]
+    assert not mismatched, f"Shape mismatches (key, exported, reference): {mismatched[:8]}"
+
+    if not check_values:
+        return
+
+    wrong = []
+    for key in shared:
+        if _is_packed(exported[key]):
+            continue  # sub-byte weights need format-specific unpacking to compare
+        got, want = _dequantize(key, exported), reference[key].to(torch.float32)
+        if exported[key].dtype == reference[key].dtype and key + "_scale" not in exported:
+            # Copied through untouched (norms, router, vision tower): must be bit-exact.
+            if not torch.equal(exported[key], reference[key]):
+                wrong.append((key, "not bit-exact"))
+        else:
+            denom = want.abs().max().clamp_min(torch.finfo(torch.float32).tiny)
+            rel = ((got - want).abs().max() / denom).item()
+            if rel > rtol:
+                wrong.append((key, f"max_rel_err={rel:.4f} > {rtol}"))
+    assert not wrong, f"{len(wrong)} tensor(s) differ from the reference, e.g. {wrong[:8]}"
diff --git a/tests/examples/megatron_bridge/test_qad.py b/tests/examples/megatron_bridge/test_qad.py
index 7d982669221..167ecce1972 100644
--- a/tests/examples/megatron_bridge/test_qad.py
+++ b/tests/examples/megatron_bridge/test_qad.py
@@ -18,6 +18,7 @@
 
 import pytest
 from _test_utils.examples.run_command import extend_cmd_parts, run_example_command
+from _test_utils.torch.export.unified_checkpoint import assert_exported_checkpoint_matches
 from _test_utils.torch.transformers_models import (
     create_tiny_gemma3vl_dir,
     create_tiny_qwen3_5_moe_vl_dir,
@@ -28,12 +29,12 @@
 
 @pytest.mark.timeout(720)  # Multiple steps in one test hence takes longer than the default timeout
 @pytest.mark.parametrize(
-    ("create_student", "exports_hf"),
+    ("create_student", "exports_hf", "no_moe_grouped_gemm"),
     [
-        (lambda tmp_path: create_tiny_qwen3_dir(tmp_path, with_tokenizer=True), True),
+        (lambda tmp_path: create_tiny_qwen3_dir(tmp_path, with_tokenizer=True), True, False),
         # Qwen3-VL is the only VLM architecture in the Megatron HF export mapping, so it is the one
         # VLM case that runs the export step end-to-end.
-        (lambda tmp_path: create_tiny_qwen3vl_dir(tmp_path, with_tokenizer=True), True),
+        (lambda tmp_path: create_tiny_qwen3vl_dir(tmp_path, with_tokenizer=True), True, False),
         # Dense-VLM QAD path; the MoE VLM below covers it in CI, so run this one on demand only.
         pytest.param(
             lambda tmp_path: create_tiny_gemma3vl_dir(
@@ -44,27 +45,30 @@
                 max_position_embeddings=512,
             ),
             False,
+            False,
             marks=pytest.mark.manual,
         ),
         pytest.param(
             lambda tmp_path: create_tiny_qwen3_5_moe_vl_dir(
                 tmp_path, with_processor=True, num_hidden_layers=2
             ),
-            False,
+            True,
+            # Gated MoE experts are only exportable as SequentialMLP; grouped GEMM raises.
+            True,
         ),
     ],
     ids=["qwen3", "qwen3vl", "gemma3vl", "qwen3_5_moe_vl"],
 )
-def test_qad(tmp_path: Path, num_gpus, create_student, exports_hf):
+def test_qad(tmp_path: Path, num_gpus, create_student, exports_hf, no_moe_grouped_gemm):
     """Quantize a tiny model, run QAD from the quantized student, and export the result.
 
     For VLMs only the language model is quantized and distilled (vision tower / projector untouched),
     and a text calibration dataset infers text-only LM calibration. Quantized-HF export needs the
-    architecture in the Megatron export mapping, so the cases whose architecture is missing there
-    (Gemma3-VL, Qwen3.5-VL) stop at the distilled Megatron checkpoint and only verify the ModelOpt
-    (quantize) state survived distillation.
+    architecture in the Megatron export mapping, so a case missing there (Gemma3-VL) stops at the
+    distilled Megatron checkpoint and only verifies the ModelOpt (quantize) state survived.
     """
     hf_model_path = create_student(tmp_path)
+    moe_flag = ["--no_moe_grouped_gemm"] if no_moe_grouped_gemm else []
     quantized_megatron_path = tmp_path / "quantized_megatron"
     distill_output_dir = tmp_path / "qad_output"
     train_iters = 3
@@ -72,7 +76,7 @@ def test_qad(tmp_path: Path, num_gpus, create_student, exports_hf):
 
     # Step 1: PTQ the (language) model to FP8 and save a Megatron checkpoint carrying the ModelOpt state.
     quantize_cmd = extend_cmd_parts(
-        ["torchrun", f"--nproc_per_node={num_gpus}", "quantize.py", "--skip_generate"],
+        ["torchrun", f"--nproc_per_node={num_gpus}", "quantize.py", "--skip_generate", *moe_flag],
         hf_model_name_or_path=hf_model_path,
         recipe="general/ptq/fp8_default-kv_fp8",
         tp_size=num_gpus,
@@ -92,7 +96,7 @@ def test_qad(tmp_path: Path, num_gpus, create_student, exports_hf):
     # quantizers) and distill from the (unquantized) HF teacher. The distilled checkpoint must keep the
     # ModelOpt state so the quantizers survive distillation.
     distill_cmd = extend_cmd_parts(
-        ["torchrun", f"--nproc_per_node={num_gpus}", "distill.py", "--use_mock_data"],
+        ["torchrun", f"--nproc_per_node={num_gpus}", "distill.py", "--use_mock_data", *moe_flag],
         student_hf_path=hf_model_path,
         student_megatron_path=quantized_megatron_path,
         teacher_hf_path=hf_model_path,
@@ -127,7 +131,12 @@ def test_qad(tmp_path: Path, num_gpus, create_student, exports_hf):
     # is only written for a quantized model, so its presence confirms the quantizers survived QAD.
     hf_export_path = tmp_path / "qad_fp8_hf"
     export_cmd = extend_cmd_parts(
-        ["torchrun", f"--nproc_per_node={num_gpus}", "export_quantized_megatron_to_hf.py"],
+        [
+            "torchrun",
+            f"--nproc_per_node={num_gpus}",
+            "export_quantized_megatron_to_hf.py",
+            *moe_flag,
+        ],
         hf_model_name_or_path=hf_model_path,
         megatron_path=distilled_megatron_path,
         export_unified_hf_path=hf_export_path,
@@ -136,4 +145,5 @@ def test_qad(tmp_path: Path, num_gpus, create_student, exports_hf):
     run_example_command(export_cmd, example_path="megatron_bridge", setup_free_port=True)
     assert (hf_export_path / "config.json").exists()
     assert (hf_export_path / "hf_quant_config.json").exists()
-    assert list(hf_export_path.glob("*.safetensors")), "Expected exported safetensors weights"
+    # QAD trains the student, so weights drift from the reference: names/shapes only.
+    assert_exported_checkpoint_matches(hf_export_path, hf_model_path, check_values=False)
diff --git a/tests/examples/megatron_bridge/test_quantize_export.py b/tests/examples/megatron_bridge/test_quantize_export.py
index 55a97b49d00..90817f91c10 100644
--- a/tests/examples/megatron_bridge/test_quantize_export.py
+++ b/tests/examples/megatron_bridge/test_quantize_export.py
@@ -17,6 +17,7 @@
 from pathlib import Path
 
 from _test_utils.examples.run_command import extend_cmd_parts, run_example_command
+from _test_utils.torch.export.unified_checkpoint import assert_exported_checkpoint_matches
 from _test_utils.torch.transformers_models import create_tiny_qwen3_dir
 
 
@@ -65,7 +66,7 @@ def test_quantize_and_export(tmp_path: Path, num_gpus):
     run_example_command(export_cmd, example_path="megatron_bridge", setup_free_port=True)
     assert (hf_export_path / "config.json").exists()
     assert (hf_export_path / "hf_quant_config.json").exists()
-    assert list(hf_export_path.glob("*.safetensors")), "Expected exported safetensors weights"
+    assert_exported_checkpoint_matches(hf_export_path, hf_model_path)
 
     # The exported unified checkpoint should be loadable and runnable by vLLM. The deployment check below
     # is disabled because it takes too long in CI (likely because of first run)
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 30d592e94ba..322f771af08 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -21,6 +21,7 @@
 import pytest
 import torch
 import transformers
+from _test_utils.torch.export.unified_checkpoint import assert_exported_checkpoint_matches
 from _test_utils.torch.megatron.models import get_mcore_gpt_model
 from _test_utils.torch.megatron.utils import get_forward
 from _test_utils.torch.transformers_models import (
@@ -166,20 +167,11 @@ def _test_unified_export_megatron(
     if quant_config:
         _verify_model_quant_config(tmp_export_dir, quant_config, kv_cache_quant_cfg)
 
+    if rank == 0 and extra_module is None:
+        # Names / shapes only: these Megatron weights are random, not loaded from model_dir.
+        assert_exported_checkpoint_matches(tmp_export_dir, model_dir, check_values=False)
+
     if model_type == "qwen3vl" and rank == 0:
-        # sanity check that vision weights were merged by export_mcore_gpt_to_hf
-        keys = []
-        for sf in sorted(tmp_export_dir.glob("*.safetensors")):
-            with safe_open(str(sf), framework="pt", device="cpu") as f:
-                keys.extend(f.keys())
-        # every decoder layer should be present, not just some
-        for i in range(num_layers):
-            assert any(k.startswith(f"model.language_model.layers.{i}.") for k in keys), (
-                f"language model layer {i} keys missing from export"
-            )
-        assert any(k.startswith("model.visual.") for k in keys), (
-            "vision encoder keys missing from export"
-        )
         # try to load the model and run a forward pass
         vl_model = Qwen3VLForConditionalGeneration.from_pretrained(
             tmp_export_dir, torch_dtype=torch.bfloat16

From ef4761dd03e8e2939ad2450285bf6556f0f4b3a5 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 28 Aug 2026 08:07:44 -0700
Subject: [PATCH 03/12] Cover gated attention and MoE export, fixing the gated
 QKV split

Widening export coverage turned up a real bug: Qwen3.5's gated
full-attention layers exported a wrong q/k/v split.

- Gated attention packs a per-head output gate next to every query head,
  so a query group is [q, gate, k, v]. _qkv_slicing assumed [q, k, v]
  and split 192 rows as 96/48/48 instead of 128/32/32. It now derives
  the group stride from config.attention_output_gate and concatenates
  the gate into q, matching Megatron-Bridge's split_qkv_weights. The
  non-gated path is unchanged.
- test_qad's qwen3_5_moe_vl case pins layer_types so it covers both
  decoder kinds; auto-generated types are all linear-attention at this
  depth. Layer count is unchanged, so CI cost is not.
- Add Qwen3-MoE to the export matrix -- the architecture whose routed
  experts were silently dropped had no export test at all.
- assert_exported_checkpoint_matches grows allow_unexpected for tensors
  the Megatron test fixture adds but tiny HF configs lack.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .../torch/export/unified_export_megatron.py   | 39 ++++++++++++++-----
 .../torch/export/unified_checkpoint.py        |  9 ++++-
 tests/examples/megatron_bridge/test_qad.py    |  7 +++-
 .../export/test_unified_export_megatron.py    | 31 ++++++++++++++-
 4 files changed, 74 insertions(+), 12 deletions(-)

diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index 7b832a0a2de..977746a6270 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -1391,7 +1391,11 @@ def _qkv_slicing(
         head_num = config.num_attention_heads
         head_size = config.kv_channels
         heads_per_group = head_num // num_query_groups
-        qkv_total_dim = head_num + 2 * num_query_groups
+        # Gated attention (e.g. Qwen3.5) packs a per-head output gate next to every query head, so a
+        # group holds [q, gate, k, v] instead of [q, k, v]. HF keeps the gate inside ``q_proj``.
+        output_gate = getattr(config, "attention_output_gate", False)
+        group_dim = (2 * heads_per_group if output_gate else heads_per_group) + 2
+        qkv_total_dim = num_query_groups * group_dim
 
         weight = name_to_value.pop("weight")
 
@@ -1413,12 +1417,24 @@ def _qkv_slicing(
 
         q_slice = torch.cat(
             [
-                torch.arange((heads_per_group + 2) * i, (heads_per_group + 2) * i + heads_per_group)
+                torch.arange(group_dim * i, group_dim * i + heads_per_group)
                 for i in range(num_query_groups_local)
             ]
         )
-        k_slice = torch.arange(heads_per_group, per_rank_qkv_dim, (heads_per_group + 2))
-        v_slice = torch.arange(heads_per_group + 1, per_rank_qkv_dim, (heads_per_group + 2))
+        gate_slice = (
+            torch.cat(
+                [
+                    torch.arange(
+                        group_dim * i + heads_per_group, group_dim * i + 2 * heads_per_group
+                    )
+                    for i in range(num_query_groups_local)
+                ]
+            )
+            if output_gate
+            else None
+        )
+        k_slice = torch.arange(group_dim - 2, per_rank_qkv_dim, group_dim)
+        v_slice = torch.arange(group_dim - 1, per_rank_qkv_dim, group_dim)
         ## Example of slices
         ## 7b: num_query_groups = head_num = 32,
         ## q_slice = [0, 3, 6, 9 , ... 90, 93]
@@ -1431,7 +1447,14 @@ def _qkv_slicing(
         slices = [q_slice, k_slice, v_slice]
         prefixes = [q_proj_prefix, k_proj_prefix, v_proj_prefix]
 
-        proj_weights = [weight[s].reshape(-1, hidden_size) for s in slices]
+        def _take(tensor, index, last_dim):
+            """Gather ``index`` heads; for q under gated attention also append the gate heads."""
+            taken = tensor[index]
+            if output_gate and index is q_slice:
+                taken = torch.cat([taken, tensor[gate_slice]], dim=1)
+            return taken.reshape(-1, last_dim)
+
+        proj_weights = [_take(weight, s, hidden_size) for s in slices]
         proj_keys = [p + "weight" for p in prefixes]
 
         if weight_scale is None:
@@ -1446,9 +1469,7 @@ def _qkv_slicing(
                     [per_rank_qkv_dim, head_size, weight_scale_hidden_size]
                 )
                 proj_weight_scales = [
-                    weight_scale[s]
-                    .reshape(-1, weight_scale_hidden_size)
-                    .to(dtype=weight_scale_dtype)
+                    _take(weight_scale, s, weight_scale_hidden_size).to(dtype=weight_scale_dtype)
                     for s in slices
                 ]
             else:
@@ -1485,7 +1506,7 @@ def _qkv_slicing(
                 # Slice bias similar to weight
                 bias = val.detach().clone()
                 bias = bias.reshape([per_rank_qkv_dim, head_size])
-                proj_biases = [bias[s].reshape(-1) for s in slices]
+                proj_biases = [_take(bias, s, 1).reshape(-1) for s in slices]
                 proj_bias_keys = [q_proj_prefix + key, k_proj_prefix + key, v_proj_prefix + key]
                 for bias_tensor, bias_key in zip(proj_biases, proj_bias_keys):
                     self._state_dict[bias_key] = bias_tensor
diff --git a/tests/_test_utils/torch/export/unified_checkpoint.py b/tests/_test_utils/torch/export/unified_checkpoint.py
index bfe094c798f..076c785ca40 100644
--- a/tests/_test_utils/torch/export/unified_checkpoint.py
+++ b/tests/_test_utils/torch/export/unified_checkpoint.py
@@ -96,6 +96,7 @@ def assert_exported_checkpoint_matches(
     ref_hf_dir: Path | str,
     *,
     allow_missing: tuple[str, ...] = (),
+    allow_unexpected: tuple[str, ...] = (),
     check_values: bool = True,
     rtol: float = 0.15,
 ) -> None:
@@ -105,6 +106,8 @@ def assert_exported_checkpoint_matches(
         export_dir: Directory holding the exported unified HF checkpoint.
         ref_hf_dir: The HuggingFace checkpoint the Megatron model was built from.
         allow_missing: Substrings of reference keys the export is expected to omit.
+        allow_unexpected: Substrings of exported keys with no reference counterpart, beyond the
+            quantization scales that are always allowed.
         check_values: Compare tensor values, not just names and shapes. Set ``False`` when the
             Megatron weights are random rather than loaded from ``ref_hf_dir``.
         rtol: Max relative error for quantized tensors (per-tensor FP8 lands well inside 0.15).
@@ -119,7 +122,11 @@ def assert_exported_checkpoint_matches(
     )
 
     # Anything extra must be a quantization scale; a stray weight means a mis-named rule.
-    unexpected = {k for k in set(exported) - set(reference) if not k.endswith(QUANT_SUFFIXES)}
+    unexpected = {
+        k
+        for k in set(exported) - set(reference)
+        if not k.endswith(QUANT_SUFFIXES) and not any(a in k for a in allow_unexpected)
+    }
     assert not unexpected, f"Export produced unexpected tensors: {sorted(unexpected)[:8]}"
 
     shared = sorted(set(exported) & set(reference))
diff --git a/tests/examples/megatron_bridge/test_qad.py b/tests/examples/megatron_bridge/test_qad.py
index 167ecce1972..39ace9c0f97 100644
--- a/tests/examples/megatron_bridge/test_qad.py
+++ b/tests/examples/megatron_bridge/test_qad.py
@@ -50,7 +50,12 @@
         ),
         pytest.param(
             lambda tmp_path: create_tiny_qwen3_5_moe_vl_dir(
-                tmp_path, with_processor=True, num_hidden_layers=2
+                tmp_path,
+                with_processor=True,
+                # Cover both Qwen3.5 decoder kinds at the same layer count: auto-generated
+                # layer_types would give linear attention only at this depth.
+                num_hidden_layers=2,
+                layer_types=["linear_attention", "full_attention"],
             ),
             True,
             # Gated MoE experts are only exportable as SequentialMLP; grouped GEMM raises.
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 322f771af08..8b0b77a1f36 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -27,6 +27,7 @@
 from _test_utils.torch.transformers_models import (
     create_tiny_llama_dir,
     create_tiny_nemotron_dir,
+    create_tiny_qwen3_moe_dir,
     create_tiny_qwen3vl_dir,
 )
 from safetensors import safe_open
@@ -99,6 +100,23 @@ def _test_unified_export_megatron(
         max_sequence_length = text_cfg.max_position_embeddings
         vocab_size = text_cfg.vocab_size
         extra_kwargs = {"kv_channels": text_cfg.head_dim, "qk_layernorm": True}
+    elif model_type == "qwen3_moe":
+        config = transformers.AutoConfig.from_pretrained(model_dir)
+        num_layers = config.num_hidden_layers
+        hidden_size = config.hidden_size
+        num_attention_heads = config.num_attention_heads
+        num_query_groups = config.num_key_value_heads
+        ffn_hidden_size = config.intermediate_size
+        max_sequence_length = config.max_position_embeddings
+        vocab_size = config.vocab_size
+        # SequentialMLP: the Qwen3 rules only cover per-expert (``local_experts``) MoE.
+        extra_kwargs = {
+            "kv_channels": config.hidden_size // config.num_attention_heads,
+            "qk_layernorm": True,
+            "num_moe_experts": config.num_experts,
+            "moe_ffn_hidden_size": config.moe_intermediate_size,
+            "moe_grouped_gemm": False,
+        }
     elif model_type in {"llama", "nemotron"}:
         config = transformers.AutoConfig.from_pretrained(model_dir)
         num_layers = config.num_hidden_layers
@@ -169,7 +187,13 @@ def _test_unified_export_megatron(
 
     if rank == 0 and extra_module is None:
         # Names / shapes only: these Megatron weights are random, not loaded from model_dir.
-        assert_exported_checkpoint_matches(tmp_export_dir, model_dir, check_values=False)
+        assert_exported_checkpoint_matches(
+            tmp_export_dir,
+            model_dir,
+            check_values=False,
+            # get_mcore_gpt_model always enables the MoE router bias; tiny HF configs have none.
+            allow_unexpected=("mlp.gate.expert_bias",),
+        )
 
     if model_type == "qwen3vl" and rank == 0:
         # try to load the model and run a forward pass
@@ -197,6 +221,9 @@ def _test_unified_export_megatron(
         ("llama", "medusa", None, None),
         ("qwen3vl", None, None, None),
         ("qwen3vl", None, "FP8_DEFAULT_CFG", None),
+        # Regression guard: routed experts used to be dropped silently from the export.
+        ("qwen3_moe", None, None, None),
+        ("qwen3_moe", None, "FP8_DEFAULT_CFG", None),
     ],
 )
 def test_unified_export_megatron(
@@ -208,6 +235,8 @@ def test_unified_export_megatron(
         model_dir = create_tiny_qwen3vl_dir(tmp_path)
     elif model_type == "nemotron":
         model_dir = create_tiny_nemotron_dir(tmp_path)
+    elif model_type == "qwen3_moe":
+        model_dir = create_tiny_qwen3_moe_dir(tmp_path)
     else:
         raise ValueError(f"Unsupported model_type: {model_type}")
     # TODO: Fix TP>1 failures

From 12e4fcb7e6c1feadec11c2f6b02fcaac40b7dfc2 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 28 Aug 2026 08:43:59 -0700
Subject: [PATCH 04/12] Verify checkpoint contents the existing assertions
 could not see

Three blind spots, each of which a real bug had already slipped through.

- rglob("modelopt_state") passes on an empty state, which is exactly what
  a dropped-state checkpoint looks like. assert_has_modelopt_state uses
  Megatron-Bridge's has_modelopt_state so the state must actually hold a
  restorable mode.
- QAD exports ran check_values=False because training drifts the weights,
  which also skipped the vision tower -- never trained, so it must stay
  bit-exact. bit_exact_prefixes checks it regardless.
- Sub-byte weights were skipped entirely. NVFP4 now unpacks (E2M1 codes,
  per-block E4M3 scale, global scale) with its own tolerance: the 8-level
  grid lands near 0.2 while a wrong split is off by ~1.0, verified by
  flipping a q_proj and watching it fail at 1.74.

Also trims docstrings across the preceding commits.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 modelopt/torch/export/plugins/mcore_custom.py |  5 +-
 .../torch/export/plugins/mcore_qwen35vl.py    | 11 ++--
 .../torch/export/plugins/mcore_qwen3vl.py     | 19 ++----
 .../torch/export/unified_export_megatron.py   |  9 +--
 .../torch/export/unified_checkpoint.py        | 65 ++++++++++++++-----
 .../torch/megatron/modelopt_state.py          | 35 ++++++++++
 tests/examples/megatron_bridge/test_qad.py    | 20 +++---
 .../megatron_bridge/test_quantize_export.py   |  5 +-
 8 files changed, 109 insertions(+), 60 deletions(-)
 create mode 100644 tests/_test_utils/torch/megatron/modelopt_state.py

diff --git a/modelopt/torch/export/plugins/mcore_custom.py b/modelopt/torch/export/plugins/mcore_custom.py
index 014744f248c..a0144a44c7f 100644
--- a/modelopt/torch/export/plugins/mcore_custom.py
+++ b/modelopt/torch/export/plugins/mcore_custom.py
@@ -244,10 +244,7 @@ def __init__(self, target_name_or_prefix: str = "", func_kwargs: dict[str, Any]
 def with_language_model_prefix(
     mapping: dict[str, CustomModuleMapping],
 ) -> dict[str, CustomModuleMapping]:
-    """Derive a VLM mapping from a text-model mapping by nesting it under ``model.language_model.``.
-
-    Prefixes that do not start with ``model.`` (e.g. ``lm_head.``) are left unchanged.
-    """
+    """Nest a text-model mapping under ``model.language_model.``; other prefixes are unchanged."""
     result = {}
     for key, m in mapping.items():
         prefix = m.target_name_or_prefix
diff --git a/modelopt/torch/export/plugins/mcore_qwen35vl.py b/modelopt/torch/export/plugins/mcore_qwen35vl.py
index 424e5142e62..a3e225cba84 100644
--- a/modelopt/torch/export/plugins/mcore_qwen35vl.py
+++ b/modelopt/torch/export/plugins/mcore_qwen35vl.py
@@ -15,13 +15,10 @@
 
 """Custom mapping from Qwen3.5-VL Hugging Face models to Megatron Core models.
 
-Qwen3.5 interleaves GatedDeltaNet linear-attention layers -- Megatron keeps them under
-``self_attention`` with a fused ``in_proj`` that HF splits into ``in_proj_qkv`` / ``_z`` / ``_b`` /
-``_a`` -- with gated full-attention layers, and adds shared experts to the Qwen3 MoE rules. As for
-Qwen3-VL, only the language model is exported; the vision tower is copied from the HF checkpoint.
-
-Requires ``--no_moe_grouped_gemm`` (SequentialMLP experts); fused ``TEGroupedMLP`` experts have no
-gated split in the exporter. Gated full-attention layers are untested.
+Qwen3.5 interleaves GatedDeltaNet linear-attention layers (fused ``in_proj``, split here into
+HF's ``in_proj_qkv`` / ``_z`` / ``_b`` / ``_a``) with gated full-attention layers, and adds MoE
+shared experts. Only the language model is exported; the vision tower is copied from HF.
+Requires ``--no_moe_grouped_gemm``: fused ``TEGroupedMLP`` experts have no gated split.
 """
 
 from .mcore_custom import (
diff --git a/modelopt/torch/export/plugins/mcore_qwen3vl.py b/modelopt/torch/export/plugins/mcore_qwen3vl.py
index 923e2d19c38..6312be1abb5 100644
--- a/modelopt/torch/export/plugins/mcore_qwen3vl.py
+++ b/modelopt/torch/export/plugins/mcore_qwen3vl.py
@@ -15,21 +15,10 @@
 
 """Custom mapping from Qwen3-VL Hugging Face models to Megatron Core models.
 
-Qwen3-VL differs from Qwen3 in one structural way: language-model weights live
-under ``model.language_model.`` instead of ``model.``, while ``lm_head.weight``
-remains at the root level, so the mappings below are derived from the Qwen3 ones
-with :func:`with_language_model_prefix`.
-
-The visual encoder (``model.visual.*``) is not mapped: only the language model is
-quantized, and the vision tower is copied verbatim from the Hugging Face checkpoint
-via ``QWEN3VL_VISION_PREFIXES``.
-
-Note: ``Qwen3VLMoeForConditionalGeneration`` is **not** supported here.  The MoE
-variant stores expert weights as 3-D tensors (``mlp.experts.gate_up_proj``,
-``mlp.experts.down_proj``) that require a dedicated fused-expert mapping and
-cannot reuse the dense Qwen3 rules.
-
-Reference: https://huggingface.co/Qwen/Qwen3-VL-8B-Instruct/blob/main/model.safetensors.index.json
+Qwen3-VL nests the language model under ``model.language_model.`` while ``lm_head`` stays at the
+root, so the mappings are derived from Qwen3's. The visual encoder is copied verbatim from HF via
+``QWEN3VL_VISION_PREFIXES`` rather than mapped. ``Qwen3VLMoeForConditionalGeneration`` is not
+supported: its 3-D fused expert weights cannot reuse the dense Qwen3 rules.
 """
 
 from .mcore_custom import with_language_model_prefix
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index 977746a6270..a57daa3b6ad 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -1516,13 +1516,10 @@ def _take(tensor, index, last_dim):
                 self._state_dict[v_proj_key] = val.detach().clone()
 
     def _gated_delta_net_slicing(self, module, prefix, is_mtp=False):
-        """Split GatedDeltaNet's fused ``in_proj`` into the four HF projections.
+        """Split GatedDeltaNet's fused ``in_proj`` into HF's qkv / z / b / a projections.
 
-        Megatron-Core packs ``[query, key, value, z, beta, alpha]`` along dim 0 of a single
-        ``in_proj``; HF stores ``in_proj_qkv`` (query+key+value), ``in_proj_z``, ``in_proj_b``
-        (beta) and ``in_proj_a`` (alpha).  The sections are contiguous and in that order, so
-        the split is a plain ``torch.split`` -- sizes come from the module itself so TP
-        sharding is handled without re-deriving them here.
+        Megatron packs ``[query, key, value, z, beta, alpha]``; sizes come from the module so TP
+        sharding needs no re-derivation.
         """
         if is_mtp:
             prefix = self._mtp_prefix(prefix)
diff --git a/tests/_test_utils/torch/export/unified_checkpoint.py b/tests/_test_utils/torch/export/unified_checkpoint.py
index 076c785ca40..9b00d7cc9d9 100644
--- a/tests/_test_utils/torch/export/unified_checkpoint.py
+++ b/tests/_test_utils/torch/export/unified_checkpoint.py
@@ -14,9 +14,8 @@
 # limitations under the License.
 """Content checks for an exported unified HuggingFace checkpoint.
 
-Existence checks (``config.json`` is there, some safetensors were written) pass even when the
-exporter silently drops a whole module family, so these helpers compare the exported tensors
-against the Hugging Face checkpoint the model came from.
+Existence checks pass even when the exporter silently drops a whole module family, so compare
+the exported tensors against the HuggingFace checkpoint the model came from.
 """
 
 import json
@@ -81,10 +80,29 @@ def _expected_shape(exported: torch.Tensor, reference: torch.Tensor) -> tuple[in
     return shape
 
 
+def _unpack_nvfp4(packed: torch.Tensor) -> torch.Tensor:
+    """Expand two E2M1 values per ``uint8`` back into a float tensor (low nibble first)."""
+    from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor
+
+    codes = torch.empty(
+        (*packed.shape[:-1], packed.shape[-1] * 2), dtype=torch.long, device=packed.device
+    )
+    codes[..., 0::2] = packed & 0x0F
+    codes[..., 1::2] = packed >> 4
+    return NVFP4QTensor.get_e2m1_values(packed.device).to(torch.float32)[codes]
+
+
 def _dequantize(key: str, exported: dict[str, torch.Tensor]) -> torch.Tensor:
-    """Undo per-tensor / per-channel weight scaling so the value can be compared to the source."""
-    weight = exported[key].to(torch.float32)
+    """Undo the exporter's weight scaling so the value can be compared to the source."""
     scale = exported.get(key.replace(".weight", ".weight_scale"))
+    if _is_packed(exported[key]):
+        # NVFP4: E2M1 codes, an E4M3 scale per block along the last dim, and a global FP32 scale.
+        weight = _unpack_nvfp4(exported[key])
+        block_size = weight.shape[-1] // scale.shape[-1]
+        weight = weight * scale.to(torch.float32).repeat_interleave(block_size, dim=-1)
+        scale_2 = exported.get(key.replace(".weight", ".weight_scale_2"))
+        return weight if scale_2 is None else weight * scale_2.to(torch.float32)
+    weight = exported[key].to(torch.float32)
     if scale is None:
         return weight
     scale = scale.to(torch.float32)
@@ -98,19 +116,21 @@ def assert_exported_checkpoint_matches(
     allow_missing: tuple[str, ...] = (),
     allow_unexpected: tuple[str, ...] = (),
     check_values: bool = True,
+    bit_exact_prefixes: tuple[str, ...] = (),
     rtol: float = 0.15,
+    packed_rtol: float = 0.3,
 ) -> None:
     """Assert an exported unified HF checkpoint reproduces the reference model it came from.
 
     Args:
-        export_dir: Directory holding the exported unified HF checkpoint.
-        ref_hf_dir: The HuggingFace checkpoint the Megatron model was built from.
-        allow_missing: Substrings of reference keys the export is expected to omit.
-        allow_unexpected: Substrings of exported keys with no reference counterpart, beyond the
-            quantization scales that are always allowed.
-        check_values: Compare tensor values, not just names and shapes. Set ``False`` when the
-            Megatron weights are random rather than loaded from ``ref_hf_dir``.
-        rtol: Max relative error for quantized tensors (per-tensor FP8 lands well inside 0.15).
+        export_dir: Exported unified HF checkpoint.
+        ref_hf_dir: HuggingFace checkpoint the Megatron model was built from.
+        allow_missing: Substrings of reference keys the export may omit.
+        allow_unexpected: Substrings of extra exported keys, beyond quantization scales.
+        check_values: Compare values too; ``False`` when the Megatron weights are random.
+        bit_exact_prefixes: Keys that must match exactly even when ``check_values`` is off.
+        rtol: Max relative error for quantized tensors.
+        packed_rtol: Same for sub-byte formats, whose grid is coarser.
     """
     exported = load_safetensors_dir(export_dir)
     reference = load_safetensors_dir(ref_hf_dir)
@@ -137,13 +157,23 @@ def assert_exported_checkpoint_matches(
     ]
     assert not mismatched, f"Shape mismatches (key, exported, reference): {mismatched[:8]}"
 
+    if bit_exact_prefixes:
+        drifted = [
+            k
+            for k in shared
+            if k.startswith(bit_exact_prefixes) and not torch.equal(exported[k], reference[k])
+        ]
+        assert not drifted, (
+            f"{len(drifted)} copied-through tensor(s) differ from the reference: {drifted[:8]}"
+        )
+        checked = [k for k in shared if k.startswith(bit_exact_prefixes)]
+        assert checked, f"bit_exact_prefixes {bit_exact_prefixes} matched no exported tensor"
+
     if not check_values:
         return
 
     wrong = []
     for key in shared:
-        if _is_packed(exported[key]):
-            continue  # sub-byte weights need format-specific unpacking to compare
         got, want = _dequantize(key, exported), reference[key].to(torch.float32)
         if exported[key].dtype == reference[key].dtype and key + "_scale" not in exported:
             # Copied through untouched (norms, router, vision tower): must be bit-exact.
@@ -152,6 +182,7 @@ def assert_exported_checkpoint_matches(
         else:
             denom = want.abs().max().clamp_min(torch.finfo(torch.float32).tiny)
             rel = ((got - want).abs().max() / denom).item()
-            if rel > rtol:
-                wrong.append((key, f"max_rel_err={rel:.4f} > {rtol}"))
+            tol = packed_rtol if _is_packed(exported[key]) else rtol
+            if rel > tol:
+                wrong.append((key, f"max_rel_err={rel:.4f} > {tol}"))
     assert not wrong, f"{len(wrong)} tensor(s) differ from the reference, e.g. {wrong[:8]}"
diff --git a/tests/_test_utils/torch/megatron/modelopt_state.py b/tests/_test_utils/torch/megatron/modelopt_state.py
new file mode 100644
index 00000000000..a21faee551d
--- /dev/null
+++ b/tests/_test_utils/torch/megatron/modelopt_state.py
@@ -0,0 +1,35 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+"""ModelOpt-state checks for Megatron distributed checkpoints."""
+
+from pathlib import Path
+
+from megatron.bridge.training.post_training.checkpointing import has_modelopt_state
+
+__all__ = ["assert_has_modelopt_state"]
+
+
+def assert_has_modelopt_state(megatron_path: Path | str) -> None:
+    """Assert a Megatron checkpoint carries restorable ModelOpt state.
+
+    ``rglob("modelopt_state")`` passes on an empty state, which exports unquantized.
+    """
+
+    state_dirs = list(Path(megatron_path).rglob("modelopt_state"))
+    assert state_dirs, f"No modelopt_state directory under {megatron_path}"
+    assert has_modelopt_state(str(megatron_path)), (
+        f"modelopt_state under {megatron_path} holds no restorable mode (only 'kd_loss' or "
+        "empty), so the quantizers would not survive a reload"
+    )
diff --git a/tests/examples/megatron_bridge/test_qad.py b/tests/examples/megatron_bridge/test_qad.py
index 39ace9c0f97..8b7d43868a0 100644
--- a/tests/examples/megatron_bridge/test_qad.py
+++ b/tests/examples/megatron_bridge/test_qad.py
@@ -19,6 +19,7 @@
 import pytest
 from _test_utils.examples.run_command import extend_cmd_parts, run_example_command
 from _test_utils.torch.export.unified_checkpoint import assert_exported_checkpoint_matches
+from _test_utils.torch.megatron.modelopt_state import assert_has_modelopt_state
 from _test_utils.torch.transformers_models import (
     create_tiny_gemma3vl_dir,
     create_tiny_qwen3_5_moe_vl_dir,
@@ -73,6 +74,7 @@ def test_qad(tmp_path: Path, num_gpus, create_student, exports_hf, no_moe_groupe
     distilled Megatron checkpoint and only verifies the ModelOpt (quantize) state survived.
     """
     hf_model_path = create_student(tmp_path)
+    is_vlm = "vision_config" in (hf_model_path / "config.json").read_text()
     moe_flag = ["--no_moe_grouped_gemm"] if no_moe_grouped_gemm else []
     quantized_megatron_path = tmp_path / "quantized_megatron"
     distill_output_dir = tmp_path / "qad_output"
@@ -93,9 +95,7 @@ def test_qad(tmp_path: Path, num_gpus, create_student, exports_hf, no_moe_groupe
         export_megatron_path=quantized_megatron_path,
     )
     run_example_command(quantize_cmd, example_path="megatron_bridge", setup_free_port=True)
-    assert list(quantized_megatron_path.rglob("modelopt_state")), (
-        "Expected modelopt_state in the quantized Megatron checkpoint"
-    )
+    assert_has_modelopt_state(quantized_megatron_path)
 
     # Step 2: QAD -- load the quantized student from the Megatron checkpoint (restoring the ModelOpt
     # quantizers) and distill from the (unquantized) HF teacher. The distilled checkpoint must keep the
@@ -125,9 +125,7 @@ def test_qad(tmp_path: Path, num_gpus, create_student, exports_hf, no_moe_groupe
     tracker = distilled_megatron_path / "latest_checkpointed_iteration.txt"
     assert tracker.read_text(encoding="utf-8").strip() == str(early_exit_iter)
     assert (distilled_megatron_path / "iter_0000001").is_dir()
-    assert list(distilled_megatron_path.rglob("modelopt_state")), (
-        "Expected modelopt_state to be preserved in the distilled (QAD) checkpoint"
-    )
+    assert_has_modelopt_state(distilled_megatron_path)
 
     if not exports_hf:
         return  # architecture missing from the export mapping; stop at the distilled checkpoint
@@ -150,5 +148,11 @@ def test_qad(tmp_path: Path, num_gpus, create_student, exports_hf, no_moe_groupe
     run_example_command(export_cmd, example_path="megatron_bridge", setup_free_port=True)
     assert (hf_export_path / "config.json").exists()
     assert (hf_export_path / "hf_quant_config.json").exists()
-    # QAD trains the student, so weights drift from the reference: names/shapes only.
-    assert_exported_checkpoint_matches(hf_export_path, hf_model_path, check_values=False)
+    # QAD trains the student, so language-model weights drift from the reference; the vision
+    # tower is never trained and must still come through byte for byte.
+    assert_exported_checkpoint_matches(
+        hf_export_path,
+        hf_model_path,
+        check_values=False,
+        bit_exact_prefixes=("model.visual.",) if is_vlm else (),
+    )
diff --git a/tests/examples/megatron_bridge/test_quantize_export.py b/tests/examples/megatron_bridge/test_quantize_export.py
index 90817f91c10..f4cc7098805 100644
--- a/tests/examples/megatron_bridge/test_quantize_export.py
+++ b/tests/examples/megatron_bridge/test_quantize_export.py
@@ -18,6 +18,7 @@
 
 from _test_utils.examples.run_command import extend_cmd_parts, run_example_command
 from _test_utils.torch.export.unified_checkpoint import assert_exported_checkpoint_matches
+from _test_utils.torch.megatron.modelopt_state import assert_has_modelopt_state
 from _test_utils.torch.transformers_models import create_tiny_qwen3_dir
 
 
@@ -52,9 +53,7 @@ def test_quantize_and_export(tmp_path: Path, num_gpus):
     )
     run_example_command(quantize_cmd, example_path="megatron_bridge", setup_free_port=True)
     assert (megatron_path / "latest_checkpointed_iteration.txt").exists()
-    assert list(megatron_path.rglob("modelopt_state")), (
-        "Expected modelopt_state in the Megatron checkpoint"
-    )
+    assert_has_modelopt_state(megatron_path)
 
     # Step 2: export to HF
     export_cmd = extend_cmd_parts(

From 4c0082c3c36edfd86a2073cc9e1d579e642e4634 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 28 Aug 2026 09:13:06 -0700
Subject: [PATCH 05/12] Fix GatedDeltaNet quantizer exclusions on Megatron-Core

The disabled-quantizer patterns are written against HuggingFace module
names, so they silently no-op wherever Megatron names things differently.
Auditing them against Megatron paths found one gap and one that no
pattern can express.

- '*linear_attn.conv1d*' never matched: Megatron calls the GatedDeltaNet
  block self_attention, so the conv1d was calibrated. '*mixer.conv1d*'
  works only because MCore and HF agree on "mixer" for Mamba. Added a
  self_attention alias to the PTQ and AutoQuantize units.
- '*linear_attn.in_proj_a/b*' cannot work at all: Megatron fuses all six
  GDN sections behind one quantizer, so the alpha/beta gates are only
  separable at export. They now go out in BF16 and are recorded in
  exclude_modules.
- assert_no_quantizers_matching makes a future name drift fail the test;
  nothing checked that an exclusion had excluded anything.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 CHANGELOG.rst                                 |  1 +
 .../torch/export/unified_export_megatron.py   | 29 ++++++++++++++-----
 .../units/base_disabled_layers.yaml           |  2 ++
 .../units/default_disabled_quantizers.yaml    |  3 ++
 .../torch/export/unified_checkpoint.py        |  4 +--
 .../torch/megatron/modelopt_state.py          | 19 +++++++++++-
 tests/examples/megatron_bridge/test_qad.py    |  7 ++++-
 7 files changed, 54 insertions(+), 11 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 71ed7bffa9f..9157edc13a0 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -50,6 +50,7 @@ Changelog
 - Fix NVFP4 ONNX export to quantize FP4 weights with the published FP8 block scales, matching eager ModelOpt packed weights. Block scales below ``2**-9`` are now clamped to that minimum, and non-finite or negative scales raise an error.
 - Fix Megatron-Bridge Quantization Aware Distillation of a vision-language model silently discarding the ModelOpt state, so the distilled checkpoint restored no quantizers and exported as an unquantized model. Re-run QAD to regenerate any affected checkpoint.
 - Fix Megatron-Core HuggingFace export silently omitting fused (grouped GEMM) MoE experts for architectures without an ``experts.linear_fc1`` rule (e.g. ``Qwen3MoeForCausalLM``), which produced a valid-looking checkpoint containing no expert weights. ``quantize.py`` and the exporter now raise instead; pass ``--no_moe_grouped_gemm`` to export these models.
+- Fix GatedDeltaNet (Qwen3.5) quantizer exclusions on Megatron-Core: the recipe patterns name the HuggingFace ``linear_attn`` module, so the ``conv1d`` was calibrated and the alpha / beta gate projections were exported in FP8. ``conv1d`` now has a ``self_attention`` alias in the default disabled-quantizer units, and the alpha / beta projections are exported in BF16 (they share Megatron's fused ``in_proj`` quantizer and cannot be disabled by name).
 - Update HuggingFace checkpoint export to use name-based tied-weight deduplication instead of the previous address-based approach. The address-based deduplication could incorrectly drop an untied weight that happened to share memory with a tied one, producing an incomplete checkpoint (observed as a false positive on MiniMax-M2.7).
 - Fix EAGLE-3 training with context parallelism (``--cp_size > 1`` in ``examples/speculative_decoding``), which failed to start on ``accelerate >= 1.13`` and then raised ``got mixed torch.Tensor and DTensor``.
 - Polygraphy minimum dependency upgraded to ``0.53.4`` to solve ONNX AutoCast failures when marking optional graph outputs.
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index a57daa3b6ad..14de63efa1d 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -1535,12 +1535,19 @@ def _gated_delta_net_slicing(self, module, prefix, is_mtp=False):
             sections["beta"],
             sections["alpha"],
         ]
-        proj_prefixes = [
-            prefix + name + "." for name in ("in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a")
-        ]
+        proj_names = ("in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a")
+        proj_prefixes = [prefix + name + "." for name in proj_names]
+        # The recipes keep the alpha / beta gates in BF16, but Megatron fuses all six sections
+        # behind one quantizer, so they can only be dropped here rather than by a quantizer_name.
+        keep_bf16 = {
+            p for p, n in zip(proj_prefixes, proj_names) if n in ("in_proj_a", "in_proj_b")
+        }
 
         for proj_prefix in proj_prefixes:
-            self._record_layer_quant_config(proj_prefix, qformat, block_size)
+            if proj_prefix in keep_bf16:
+                self._record_excluded_module(proj_prefix)
+            else:
+                self._record_layer_quant_config(proj_prefix, qformat, block_size)
         if qformat in (None, QUANTIZATION_NONE):
             # Split the fused in_proj exclude entry into the per-HF-name projections.
             self.exclude_modules = [
@@ -1563,7 +1570,12 @@ def _gated_delta_net_slicing(self, module, prefix, is_mtp=False):
                 proj_scales = list(torch.split(weight_scale, split_sizes, dim=0))
             else:
                 proj_scales = [weight_scale.detach().clone() for _ in proj_keys]
-            for proj_weight, scale, key in zip(proj_weights, proj_scales, proj_keys):
+            for proj_prefix, proj_weight, scale, key in zip(
+                proj_prefixes, proj_weights, proj_scales, proj_keys
+            ):
+                if proj_prefix in keep_bf16:
+                    self._state_dict[key] = proj_weight
+                    continue
                 self._state_dict[key] = to_quantized_weight(
                     proj_weight, scale, qformat, weight_scale_2, block_size
                 )
@@ -1572,8 +1584,9 @@ def _gated_delta_net_slicing(self, module, prefix, is_mtp=False):
         if weight_scale_2 is not None:
             if len(weight_scale_2.shape) > 0:
                 raise ValueError("weight_scale_2 must be a scalar!")
-            for key in proj_keys:
-                self._state_dict[key + "_scale_2"] = weight_scale_2.detach().clone()
+            for proj_prefix, key in zip(proj_prefixes, proj_keys):
+                if proj_prefix not in keep_bf16:
+                    self._state_dict[key + "_scale_2"] = weight_scale_2.detach().clone()
 
         # weight and weight_scale have been popped; the rest (bias, input_scale, ...) is
         # either split like the weight or replicated onto every projection.
@@ -1585,6 +1598,8 @@ def _gated_delta_net_slicing(self, module, prefix, is_mtp=False):
                     self._state_dict[proj_prefix + key] = proj_bias
             else:
                 for proj_prefix in proj_prefixes:
+                    if proj_prefix in keep_bf16:
+                        continue
                     self._state_dict[proj_prefix + key] = val.detach().clone()
 
     def _self_attention_scaling(
diff --git a/modelopt_recipes/configs/auto_quantize/units/base_disabled_layers.yaml b/modelopt_recipes/configs/auto_quantize/units/base_disabled_layers.yaml
index fd27046d608..fa34eda4a44 100644
--- a/modelopt_recipes/configs/auto_quantize/units/base_disabled_layers.yaml
+++ b/modelopt_recipes/configs/auto_quantize/units/base_disabled_layers.yaml
@@ -23,6 +23,8 @@
   - "*linear_attn.in_proj_a*"
   - "*linear_attn.in_proj_b*"
   - "*mixer.conv1d*"
+  # Megatron-Core names the GatedDeltaNet block ``self_attention``, HF ``linear_attn``.
+  - "*self_attention.conv1d*"
   - "*mlp.gate.*"
   - "*mlp.shared_expert_gate.*"
   - "*output_layer*"
diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml
index 3aadadd289c..e819c82ab4e 100644
--- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml
+++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml
@@ -28,6 +28,9 @@
     enable: false
   - quantizer_name: '*mixer.conv1d*'
     enable: false
+  # Megatron-Core names the GatedDeltaNet block ``self_attention``; HF calls it ``linear_attn``.
+  - quantizer_name: '*self_attention.conv1d*'
+    enable: false
   - quantizer_name: '*mlp.gate.*'
     enable: false
   - quantizer_name: '*mlp.shared_expert_gate.*'
diff --git a/tests/_test_utils/torch/export/unified_checkpoint.py b/tests/_test_utils/torch/export/unified_checkpoint.py
index 9b00d7cc9d9..e087f5f738d 100644
--- a/tests/_test_utils/torch/export/unified_checkpoint.py
+++ b/tests/_test_utils/torch/export/unified_checkpoint.py
@@ -24,6 +24,8 @@
 import torch
 from safetensors.torch import load_file
 
+from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor
+
 __all__ = [
     "assert_exported_checkpoint_matches",
     "assert_safetensors_index_consistent",
@@ -82,8 +84,6 @@ def _expected_shape(exported: torch.Tensor, reference: torch.Tensor) -> tuple[in
 
 def _unpack_nvfp4(packed: torch.Tensor) -> torch.Tensor:
     """Expand two E2M1 values per ``uint8`` back into a float tensor (low nibble first)."""
-    from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor
-
     codes = torch.empty(
         (*packed.shape[:-1], packed.shape[-1] * 2), dtype=torch.long, device=packed.device
     )
diff --git a/tests/_test_utils/torch/megatron/modelopt_state.py b/tests/_test_utils/torch/megatron/modelopt_state.py
index a21faee551d..2c5f22bd5d5 100644
--- a/tests/_test_utils/torch/megatron/modelopt_state.py
+++ b/tests/_test_utils/torch/megatron/modelopt_state.py
@@ -17,8 +17,9 @@
 from pathlib import Path
 
 from megatron.bridge.training.post_training.checkpointing import has_modelopt_state
+from torch.distributed.checkpoint import FileSystemReader
 
-__all__ = ["assert_has_modelopt_state"]
+__all__ = ["assert_has_modelopt_state", "assert_no_quantizers_matching"]
 
 
 def assert_has_modelopt_state(megatron_path: Path | str) -> None:
@@ -33,3 +34,19 @@ def assert_has_modelopt_state(megatron_path: Path | str) -> None:
         f"modelopt_state under {megatron_path} holds no restorable mode (only 'kd_loss' or "
         "empty), so the quantizers would not survive a reload"
     )
+
+
+def assert_no_quantizers_matching(megatron_path: Path | str, *substrings: str) -> None:
+    """Assert no calibrated quantizer under ``megatron_path`` matches ``substrings``.
+
+    Disabled-quantizer patterns are written against HuggingFace names, so they silently
+    no-op wherever Megatron names the module differently.
+    """
+    iter_dirs = sorted(Path(megatron_path).glob("iter_*"))
+    assert iter_dirs, f"No iter_* checkpoint under {megatron_path}"
+    keys = FileSystemReader(str(iter_dirs[-1])).read_metadata().state_dict_metadata
+    quantizers = [k for k in keys if "_quantizer." in k]
+    assert quantizers, f"No quantizers at all under {megatron_path}; was the model quantized?"
+    for substring in substrings:
+        hits = sorted(k for k in quantizers if substring in k)
+        assert not hits, f"Expected no quantizer matching {substring!r}, found: {hits[:4]}"
diff --git a/tests/examples/megatron_bridge/test_qad.py b/tests/examples/megatron_bridge/test_qad.py
index 8b7d43868a0..79fe0758d21 100644
--- a/tests/examples/megatron_bridge/test_qad.py
+++ b/tests/examples/megatron_bridge/test_qad.py
@@ -19,7 +19,10 @@
 import pytest
 from _test_utils.examples.run_command import extend_cmd_parts, run_example_command
 from _test_utils.torch.export.unified_checkpoint import assert_exported_checkpoint_matches
-from _test_utils.torch.megatron.modelopt_state import assert_has_modelopt_state
+from _test_utils.torch.megatron.modelopt_state import (
+    assert_has_modelopt_state,
+    assert_no_quantizers_matching,
+)
 from _test_utils.torch.transformers_models import (
     create_tiny_gemma3vl_dir,
     create_tiny_qwen3_5_moe_vl_dir,
@@ -96,6 +99,8 @@ def test_qad(tmp_path: Path, num_gpus, create_student, exports_hf, no_moe_groupe
     )
     run_example_command(quantize_cmd, example_path="megatron_bridge", setup_free_port=True)
     assert_has_modelopt_state(quantized_megatron_path)
+    # Megatron names these differently from HF, so the recipe's patterns must have aliases.
+    assert_no_quantizers_matching(quantized_megatron_path, "conv1d", "mlp.router", "output_layer")
 
     # Step 2: QAD -- load the quantized student from the Megatron checkpoint (restoring the ModelOpt
     # quantizers) and distill from the (unquantized) HF teacher. The distilled checkpoint must keep the

From 18db85af8afb0ffc536e7e976eb4c8433a530a71 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 28 Aug 2026 09:45:01 -0700
Subject: [PATCH 06/12] Make the exporter verify its own output

A missing export rule emits nothing rather than failing, so any
architecture can lose a whole module family and still write a checkpoint
that loads. The MoE guard only covers the case we happened to find.

save_pretrained now compares the exported key set against the source
checkpoint and raises on tensors that went missing, accounting for
depth-pruned models and tied embeddings. Hub ids are skipped rather than
downloaded during export.

Verified by disabling the MoE guard and re-exporting Qwen3-MoE: the check
independently reports all 24 dropped expert tensors. No false positives
across llama, nemotron, qwen3, qwen3-moe, qwen3vl and qwen3.5-vl,
including the eagle, medusa and MTP variants.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 CHANGELOG.rst                                 |  1 +
 .../torch/export/unified_export_megatron.py   | 54 +++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 9157edc13a0..bb6371c05d9 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -51,6 +51,7 @@ Changelog
 - Fix Megatron-Bridge Quantization Aware Distillation of a vision-language model silently discarding the ModelOpt state, so the distilled checkpoint restored no quantizers and exported as an unquantized model. Re-run QAD to regenerate any affected checkpoint.
 - Fix Megatron-Core HuggingFace export silently omitting fused (grouped GEMM) MoE experts for architectures without an ``experts.linear_fc1`` rule (e.g. ``Qwen3MoeForCausalLM``), which produced a valid-looking checkpoint containing no expert weights. ``quantize.py`` and the exporter now raise instead; pass ``--no_moe_grouped_gemm`` to export these models.
 - Fix GatedDeltaNet (Qwen3.5) quantizer exclusions on Megatron-Core: the recipe patterns name the HuggingFace ``linear_attn`` module, so the ``conv1d`` was calibrated and the alpha / beta gate projections were exported in FP8. ``conv1d`` now has a ``self_attention`` alias in the default disabled-quantizer units, and the alpha / beta projections are exported in BF16 (they share Megatron's fused ``in_proj`` quantizer and cannot be disabled by name).
+- Megatron-Core HuggingFace export now verifies its own output: if the exported checkpoint is missing tensors the source checkpoint has (an architecture whose module has no export rule), it raises instead of writing a valid-looking checkpoint. Depth-pruned models and tied embeddings are accounted for.
 - Update HuggingFace checkpoint export to use name-based tied-weight deduplication instead of the previous address-based approach. The address-based deduplication could incorrectly drop an untied weight that happened to share memory with a tied one, producing an incomplete checkpoint (observed as a false positive on MiniMax-M2.7).
 - Fix EAGLE-3 training with context parallelism (``--cp_size > 1`` in ``examples/speculative_decoding``), which failed to start on ``accelerate >= 1.13`` and then raised ``got mixed torch.Tensor and DTensor``.
 - Polygraphy minimum dependency upgraded to ``0.53.4`` to solve ONNX AutoCast failures when marking optional graph outputs.
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index 14de63efa1d..395822df176 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -21,6 +21,7 @@
 import io
 import json
 import os
+import re
 import tempfile
 from collections import OrderedDict
 from pathlib import Path
@@ -439,6 +440,45 @@ def save_pretrained(
             name_template="model-{:05d}-of-{:05d}",
         )
 
+        # Every rank has written its shards; one rank now checks nothing was dropped.
+        torch.distributed.barrier()
+        if is_writer_rank:
+            self._verify_exported_keys(save_directory, pretrained_model_name_or_path)
+
+    def _verify_exported_keys(self, save_directory, pretrained_model_name_or_path) -> None:
+        """Raise if the export dropped tensors the source checkpoint has.
+
+        A missing export rule emits nothing rather than failing, so the checkpoint looks valid.
+        """
+        if pretrained_model_name_or_path is None or not os.path.isdir(
+            str(pretrained_model_name_or_path)
+        ):
+            return  # hub id: not worth a download inside export
+        index_file = Path(save_directory) / "model.safetensors.index.json"
+        if not index_file.exists():
+            return
+        with open(index_file) as f:
+            exported = set(json.load(f)["weight_map"])
+        source = _read_checkpoint_keys(pretrained_model_name_or_path)
+        if not source:
+            return
+
+        num_layers = self.model.config.num_layers
+        missing = set()
+        for key in source - exported:
+            layer = re.search(r"\.layers\.(\d+)\.", key)
+            if layer is not None and int(layer.group(1)) >= num_layers:
+                continue  # depth-pruned model: the source has layers this export does not
+            if key == "lm_head.weight" and self.model.share_embeddings_and_output_weights:
+                continue  # tied embeddings: no separate output layer to export
+            missing.add(key)
+        if missing:
+            raise RuntimeError(
+                f"Export dropped {len(missing)} tensor(s) present in "
+                f"{pretrained_model_name_or_path}, e.g. {sorted(missing)[:8]}. This usually means "
+                "an architecture has no export rule for one of its modules."
+            )
+
     @property
     def state_dict(self):
         """Return the real quantized state_dict of the base model."""
@@ -1845,6 +1885,20 @@ def _gather_kv_cache_dtype(self):
         return None
 
 
+def _read_checkpoint_keys(checkpoint_dir) -> set[str]:
+    """Tensor names in a local HuggingFace checkpoint, from its index or single safetensors file."""
+    directory = Path(checkpoint_dir)
+    index_file = directory / "model.safetensors.index.json"
+    if index_file.exists():
+        with open(index_file) as f:
+            return set(json.load(f)["weight_map"])
+    single_file = directory / "model.safetensors"
+    if single_file.exists():
+        with safe_open(str(single_file), framework="pt", device="cpu") as f:
+            return set(f.keys())
+    return set()
+
+
 def export_mcore_gpt_to_hf(
     model: torch.nn.Module,
     pretrained_model_name_or_path: str | os.PathLike,

From 2baab5ece9e1794dd7eab91a3cf5ef278fa51679 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 28 Aug 2026 10:05:52 -0700
Subject: [PATCH 07/12] Unify VLM detection and fail loudly on dropped ModelOpt
 state

Five call sites answered "is this a VLM" three different ways -- by
probing .language_model, by looking for vision_config, and by isinstance
against LLaVAModel. They can disagree, and a disagreement silently
quantizes the vision tower or skips the language model.

get_language_model (model-side) and is_vlm_config (config-side, for
callers that run before the Megatron model exists) replace all of them,
documented as having to agree.

Loading a checkpoint whose quantizer tensors have no restorable state now
raises. That state loss is what the VLM QAD bug produced, and the loader
ignores the leftover amax tensors, so the model would otherwise come back
unquantized with no error.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 CHANGELOG.rst                                 |  1 +
 examples/megatron_bridge/distill.py           |  9 +--
 .../export_distilled_megatron_to_hf.py        |  6 +-
 examples/megatron_bridge/prune_minitron.py    |  5 +-
 examples/megatron_bridge/quantize.py          |  5 +-
 modelopt/torch/utils/plugins/mbridge.py       | 56 +++++++++++++++----
 6 files changed, 55 insertions(+), 27 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index bb6371c05d9..bf448c786e6 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -52,6 +52,7 @@ Changelog
 - Fix Megatron-Core HuggingFace export silently omitting fused (grouped GEMM) MoE experts for architectures without an ``experts.linear_fc1`` rule (e.g. ``Qwen3MoeForCausalLM``), which produced a valid-looking checkpoint containing no expert weights. ``quantize.py`` and the exporter now raise instead; pass ``--no_moe_grouped_gemm`` to export these models.
 - Fix GatedDeltaNet (Qwen3.5) quantizer exclusions on Megatron-Core: the recipe patterns name the HuggingFace ``linear_attn`` module, so the ``conv1d`` was calibrated and the alpha / beta gate projections were exported in FP8. ``conv1d`` now has a ``self_attention`` alias in the default disabled-quantizer units, and the alpha / beta projections are exported in BF16 (they share Megatron's fused ``in_proj`` quantizer and cannot be disabled by name).
 - Megatron-Core HuggingFace export now verifies its own output: if the exported checkpoint is missing tensors the source checkpoint has (an architecture whose module has no export rule), it raises instead of writing a valid-looking checkpoint. Depth-pruned models and tied embeddings are accounted for.
+- Loading a Megatron checkpoint that holds quantizer tensors but no restorable ModelOpt state now raises instead of silently loading the model unquantized.
 - Update HuggingFace checkpoint export to use name-based tied-weight deduplication instead of the previous address-based approach. The address-based deduplication could incorrectly drop an untied weight that happened to share memory with a tied one, producing an incomplete checkpoint (observed as a false positive on MiniMax-M2.7).
 - Fix EAGLE-3 training with context parallelism (``--cp_size > 1`` in ``examples/speculative_decoding``), which failed to start on ``accelerate >= 1.13`` and then raised ``got mixed torch.Tensor and DTensor``.
 - Polygraphy minimum dependency upgraded to ``0.53.4`` to solve ONNX AutoCast failures when marking optional graph outputs.
diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py
index 80e60546cf6..0151921c4df 100644
--- a/examples/megatron_bridge/distill.py
+++ b/examples/megatron_bridge/distill.py
@@ -50,13 +50,13 @@
 from megatron.core.datasets.utils import get_blend_from_list
 from megatron.core.distributed import DistributedDataParallelConfig
 from megatron.core.utils import unwrap_model
-from transformers import AutoConfig, AutoTokenizer
+from transformers import AutoTokenizer
 
 import modelopt.torch.distill as mtd
 import modelopt.torch.utils.distributed as dist
 from modelopt.torch.opt.conversion import ModeloptStateManager
 from modelopt.torch.utils import print_args, print_rank_0, warn_rank_0
-from modelopt.torch.utils.plugins.mbridge import load_modelopt_megatron_checkpoint
+from modelopt.torch.utils.plugins.mbridge import is_vlm_config, load_modelopt_megatron_checkpoint
 
 with contextlib.suppress(ModuleNotFoundError):
     import modelopt.torch.puzzletron.plugins.mbridge  # noqa: F401
@@ -403,10 +403,7 @@ def _build_model_provider(hf_path, load_weights=True):
 
     # HF VLM configs expose ``vision_config``; Megatron-Bridge nests the text model under
     # ``language_model`` (used as ``distill_submodule`` below).
-    is_vlm = hasattr(
-        AutoConfig.from_pretrained(args.student_hf_path, trust_remote_code=args.trust_remote_code),
-        "vision_config",
-    )
+    is_vlm = is_vlm_config(args.student_hf_path, trust_remote_code=args.trust_remote_code)
 
     if is_vlm:
         warn_rank_0(
diff --git a/examples/megatron_bridge/export_distilled_megatron_to_hf.py b/examples/megatron_bridge/export_distilled_megatron_to_hf.py
index d5e95ff9d18..f70df8381ef 100644
--- a/examples/megatron_bridge/export_distilled_megatron_to_hf.py
+++ b/examples/megatron_bridge/export_distilled_megatron_to_hf.py
@@ -65,6 +65,7 @@
 from modelopt.torch.export import copy_hf_ckpt_remote_code
 from modelopt.torch.utils import print_args, print_rank_0
 from modelopt.torch.utils.plugins.mbridge import (
+    is_vlm_config,
     load_mbridge_model_from_hf,
     load_modelopt_megatron_checkpoint,
 )
@@ -226,10 +227,7 @@ def get_args() -> argparse.Namespace:
 
 def main(args: argparse.Namespace):
     checkpoint_export_paths: list[tuple[Path, Path]] = _get_checkpoint_export_paths(args)
-    is_vlm = hasattr(
-        AutoConfig.from_pretrained(args.student_hf_path, trust_remote_code=args.trust_remote_code),
-        "vision_config",
-    )
+    is_vlm = is_vlm_config(args.student_hf_path, trust_remote_code=args.trust_remote_code)
 
     if is_vlm:
         # Build the full VLM (vision tower / projector + original LM from HF), then overwrite the LM
diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py
index 42cb78f2a7c..974de2192a8 100644
--- a/examples/megatron_bridge/prune_minitron.py
+++ b/examples/megatron_bridge/prune_minitron.py
@@ -66,7 +66,7 @@
     print_rank_0,
     warn_rank_0,
 )
-from modelopt.torch.utils.plugins.mbridge import load_mbridge_model_from_hf
+from modelopt.torch.utils.plugins.mbridge import get_language_model, load_mbridge_model_from_hf
 from modelopt.torch.utils.plugins.megatron_calibration import (
     get_megatron_calibration_forward_loop,
     get_megatron_vlm_calibration_forward_loop,
@@ -430,8 +430,7 @@ def main(args: argparse.Namespace):
 
     # For VLMs (e.g. Qwen3-VL), only the language model is pruned; the vision tower is left intact.
     # hidden_size is shared with the vision->LM projector, so it is skipped
-    language_model = getattr(unwrapped_model, "language_model", unwrapped_model)
-    is_vlm = language_model is not unwrapped_model
+    language_model, is_vlm = get_language_model(unwrapped_model)
     if is_vlm:
         warn_rank_0(
             "VLM detected: pruning model.language_model only; all non-language-model components "
diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py
index 27dd8ad630c..295a8f1cc16 100644
--- a/examples/megatron_bridge/quantize.py
+++ b/examples/megatron_bridge/quantize.py
@@ -71,7 +71,7 @@
 from modelopt.torch.export.plugins.mcore_common import all_mcore_hf_export_mapping
 from modelopt.torch.utils import print_args, print_rank_0, warn_rank_0
 from modelopt.torch.utils.dataset_utils import get_supported_datasets
-from modelopt.torch.utils.plugins.mbridge import load_mbridge_model_from_hf
+from modelopt.torch.utils.plugins.mbridge import get_language_model, load_mbridge_model_from_hf
 from modelopt.torch.utils.plugins.megatron_calibration import (
     get_megatron_calibration_forward_loop,
     get_megatron_vlm_calibration_forward_loop,
@@ -303,8 +303,7 @@ def main(args: argparse.Namespace):
     )
 
     # Only the language model is quantized (vision tower + projector stay full precision)
-    language_model = getattr(unwrapped_model, "language_model", unwrapped_model)
-    is_vlm = language_model is not unwrapped_model
+    language_model, is_vlm = get_language_model(unwrapped_model)
     if is_vlm:
         warn_rank_0(
             "VLM detected: quantizing `model.language_model` only (vision tower left in full precision)."
diff --git a/modelopt/torch/utils/plugins/mbridge.py b/modelopt/torch/utils/plugins/mbridge.py
index dcd97491a2e..c9120269ba1 100644
--- a/modelopt/torch/utils/plugins/mbridge.py
+++ b/modelopt/torch/utils/plugins/mbridge.py
@@ -31,12 +31,33 @@
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.utils import unwrap_model
 from torch.distributed.checkpoint import FileSystemReader
-from transformers import AutoTokenizer
+from transformers import AutoConfig, AutoTokenizer
 
 from modelopt.torch.nas.plugins.megatron import get_te_hybrid_stack_spec
 from modelopt.torch.utils import print_rank_0
 
-__all__ = ["load_mbridge_model_from_hf", "load_modelopt_megatron_checkpoint"]
+__all__ = [
+    "get_language_model",
+    "is_vlm_config",
+    "load_mbridge_model_from_hf",
+    "load_modelopt_megatron_checkpoint",
+]
+
+
+def get_language_model(model: MegatronModule) -> tuple[MegatronModule, bool]:
+    """Return ``(language_model, is_vlm)``; VLM wrappers nest it under ``.language_model``.
+
+    Must agree with :func:`is_vlm_config`, which answers the same question before the
+    Megatron model exists.
+    """
+    language_model = getattr(model, "language_model", None)
+    return (model, False) if language_model is None else (language_model, True)
+
+
+def is_vlm_config(hf_model_name_or_path: str, trust_remote_code: bool = False) -> bool:
+    """Whether a HuggingFace checkpoint describes a VLM, from its config alone."""
+    config = AutoConfig.from_pretrained(hf_model_name_or_path, trust_remote_code=trust_remote_code)
+    return hasattr(config, "vision_config")
 
 
 def load_mbridge_model_from_hf(
@@ -100,9 +121,9 @@ def load_mbridge_model_from_hf(
     model = provider.provide_distributed_model(wrap_with_ddp=False)
     assert len(model) == 1
     unwrapped_model = unwrap_model(model[0])
-    # VLMs (e.g. Qwen3-VL) wrap the language model as ``.language_model``; the pruning target is the
-    # inner GPTModel/HybridModel, but we still return the full wrapper so callers can save the VLM.
-    language_model = getattr(unwrapped_model, "language_model", unwrapped_model)
+    # The optimization target is the inner GPTModel/HybridModel, but callers get the full
+    # wrapper back so they can save the whole VLM.
+    language_model, _ = get_language_model(unwrapped_model)
     assert isinstance(language_model, GPTModel | HybridModel), (
         f"Expected a GPTModel/HybridModel (optionally wrapped as `.language_model`), "
         f"got {type(unwrapped_model)}"
@@ -119,10 +140,14 @@ def load_mbridge_model_from_hf(
     return bridge, provider, model, unwrapped_model, tokenizer
 
 
+def _checkpoint_keys(checkpoint_path: str) -> list[str]:
+    """Tensor names in a Megatron distributed checkpoint."""
+    return list(FileSystemReader(checkpoint_path).read_metadata().state_dict_metadata)
+
+
 def _has_vision_model_weights(checkpoint_path: str) -> bool:
     """Whether a Megatron distributed checkpoint holds a VLM's vision tower (``vision_model.*``)."""
-    metadata = FileSystemReader(checkpoint_path).read_metadata()
-    return any(key.startswith("vision_model.") for key in metadata.state_dict_metadata)
+    return any(key.startswith("vision_model.") for key in _checkpoint_keys(checkpoint_path))
 
 
 def load_modelopt_megatron_checkpoint(
@@ -143,14 +168,23 @@ def load_modelopt_megatron_checkpoint(
     # ``distill.py`` distills a VLM's language model only, so its checkpoint holds no ``vision_model.*``
     # weights and must be loaded into ``.language_model`` rather than the full VLM wrapper.
     unwrapped_model = unwrap_model(model)
-    if any(hasattr(m, "language_model") for m in unwrapped_model) and not _has_vision_model_weights(
+    if any(get_language_model(m)[1] for m in unwrapped_model) and not _has_vision_model_weights(
         checkpoint_path
     ):
         print_rank_0("Language-model-only checkpoint: loading into the VLM's `.language_model`.")
-        model = [getattr(m, "language_model", m) for m in unwrapped_model]
+        model = [get_language_model(m)[0] for m in unwrapped_model]
 
     # Restore the ModelOpt state before loading weights.
     # has_modelopt_state / load_modelopt_state resolves the latest iter_* directory
-    if restore_modelopt_state and has_modelopt_state(megatron_path):
-        load_modelopt_state(model, megatron_path)
+    if restore_modelopt_state:
+        if has_modelopt_state(megatron_path):
+            load_modelopt_state(model, megatron_path)
+        elif any("_quantizer." in key for key in _checkpoint_keys(checkpoint_path)):
+            # The quantizers cannot be rebuilt without the state, and the loader ignores the
+            # leftover amax tensors, so the model would silently load unquantized.
+            raise RuntimeError(
+                f"{megatron_path} holds quantizer tensors but no restorable ModelOpt state. "
+                "The state was dropped when the checkpoint was written -- re-run the step that "
+                "produced it."
+            )
     _load_model_weights_from_checkpoint(checkpoint_path, model)

From a2f0fe0154f42cc6e5637a752c334ab739533761 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 28 Aug 2026 14:25:42 -0700
Subject: [PATCH 08/12] Address review, and pick the MoE expert layout
 automatically

Review fixes
- _verify_exported_keys compared tensor names, so an already-quantized HF
  source failed export outright: DeepSeek carries weight_scale_inv and
  GPT-OSS carries *_blocks / *_scales, neither of which has an export
  counterpart. It now compares module prefixes within decoder layers,
  which still catches a dropped module family while tolerating both
  source-side quantization names and per-architecture top-level naming.
- Assert the GDN out_norm's config agrees before shifting gamma by 1.0.
  Deriving the offset from layernorm_zero_centered_gamma would be wrong:
  that flag is model-wide on Qwen3.5, but only this norm needs it.
- export_distilled_megatron_to_hf.py never got the MoE layout option, so
  it rebuilt a mismatched model.
- Transferring ModelOpt state assumed the state was on the VLM root, and
  asserted when QAD resumed from a language-model-only checkpoint. The
  loader now returns what it loaded into.
- Route trust_remote_code through is_safe_repo in quantize.py.
- Hoist set_moe_expert_layout out of load_mbridge_model_from_hf so
  distill.py stops duplicating it: hybrid providers also need their stack
  spec rebuilt, which the copy in distill.py was missing.
- Drop the stale "Qwen3-VL only" notes.

MoE expert layout
Only Nemotron-H can export fused grouped-GEMM experts, so every other MoE
architecture had to be run with --no_moe_grouped_gemm on all four scripts
or hit a wall at export. The scripts now derive the layout from the model
config: grouped GEMM unless it would not be exportable, SequentialMLP
otherwise, with --no_moe_grouped_gemm forcing the latter. Being a pure
function of the config, all four agree without threading a flag.

Note this changes MoE activation scales from one shared scale to
per-expert for the affected architectures.

Test coverage
Per-architecture export mappings are covered in-process by
tests/gpu_megatron; the example tests are slow because each step spawns
torchrun, so they now cover script wiring only. QAD keeps one LLM and one
VLM case (its unique property is that ModelOpt state survives distill),
while qwen3vl moves to quantize+export and qwen3_moe and nemotron_h join
it -- nemotron_h being the one architecture that keeps grouped GEMM.

The tiny NemotronH fixture also saved its embedding under the legacy
singular name, disagreeing with every released checkpoint. That rename is
transformers <= 5.15 behaviour, dropped upstream in 5.16; the fixture now
removes it so saved tiny models match real ones.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 CHANGELOG.rst                                 |   6 +-
 examples/megatron_bridge/README.md            |   2 +-
 examples/megatron_bridge/distill.py           |  31 ++++--
 .../export_distilled_megatron_to_hf.py        |  15 +++
 .../export_quantized_megatron_to_hf.py        |  15 ++-
 examples/megatron_bridge/quantize.py          |  53 +++++-----
 .../torch/export/plugins/mcore_qwen35vl.py    |   3 +-
 .../torch/export/unified_export_megatron.py   |  63 ++++++-----
 modelopt/torch/utils/plugins/mbridge.py       | 100 ++++++++++++++----
 .../_test_utils/torch/transformers_models.py  |  27 +++++
 tests/examples/megatron_bridge/test_qad.py    |  47 ++------
 .../megatron_bridge/test_quantize_export.py   |  55 +++++++---
 .../export/plugins/test_moe_layout_choice.py  |  40 +++++++
 13 files changed, 306 insertions(+), 151 deletions(-)
 create mode 100644 tests/gpu_megatron/torch/export/plugins/test_moe_layout_choice.py

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index bf448c786e6..16ec41ae920 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -19,8 +19,8 @@ Changelog
 - Add SFT-masked data support to ``examples/megatron_bridge/distill.py``: ``--sft --sft_dataset_root <dir>`` distills on raw prompt-completion JSONL (``{"input", "output"}`` records) with the loss masked to the response tokens, using Megatron-Bridge's ``FinetuningDatasetConfig`` and the model's own HuggingFace tokenizer instead of the pretraining ``GPTDataset`` and ``NullTokenizer``.
 - Add per-expert weight quantization for Transformer Engine ``TEGroupedMLP`` (fused MoE experts): each expert now has its own ``weight_quantizer`` (a ``GroupedQuantizer`` holding one ``TensorQuantizer`` per expert) with an independent ``amax``, instead of a single shared ``amax`` across all experts. Applies to ``mtq.quantize`` calibration, HF / Megatron export, and QAD.
 - Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager.
-- Add HuggingFace unified export of quantized Qwen3-VL and Qwen3.5-VL checkpoints (PTQ or QAD) via ``examples/megatron_bridge/export_quantized_megatron_to_hf.py``. Only the language model is quantized; the vision tower is copied from the source HuggingFace checkpoint. Qwen3.5-VL additionally covers GatedDeltaNet linear-attention layers and MoE shared experts, and requires ``--no_moe_grouped_gemm``. Gemma3-VL is still saved in Megatron checkpoint format only.
-- Add ``--no_moe_grouped_gemm`` to ``examples/megatron_bridge/quantize.py``, ``distill.py`` and ``export_quantized_megatron_to_hf.py``, building MoE experts as ``SequentialMLP`` instead of the fused ``TEGroupedMLP``. The flag must match across all three, since the expert layout is baked into the Megatron checkpoint.
+- Add HuggingFace unified export of quantized Qwen3-VL and Qwen3.5-VL checkpoints (PTQ or QAD) via ``examples/megatron_bridge/export_quantized_megatron_to_hf.py``, Qwen3.5-VL additionally covering GatedDeltaNet linear-attention layers and MoE shared experts. Only the language model is quantized; the vision tower is copied from the source HuggingFace checkpoint.
+- Megatron-Bridge scripts now choose the MoE expert layout automatically from the model config: the faster fused ``TEGroupedMLP`` (grouped GEMM) unless the architecture cannot export it to HuggingFace, in which case ``SequentialMLP`` keeps the checkpoint exportable and ``--no_moe_grouped_gemm`` forces it explicitly. For the affected architectures this changes MoE activation scales from one shared scale to per-expert.
 
 *Misc*
 
@@ -49,7 +49,7 @@ Changelog
 - Avoid querying CUDA/Blackwell capability when ``NVFP4QTensor.quantize`` uses its CPU path or has the optional TensorRT-LLM fast path disabled.
 - Fix NVFP4 ONNX export to quantize FP4 weights with the published FP8 block scales, matching eager ModelOpt packed weights. Block scales below ``2**-9`` are now clamped to that minimum, and non-finite or negative scales raise an error.
 - Fix Megatron-Bridge Quantization Aware Distillation of a vision-language model silently discarding the ModelOpt state, so the distilled checkpoint restored no quantizers and exported as an unquantized model. Re-run QAD to regenerate any affected checkpoint.
-- Fix Megatron-Core HuggingFace export silently omitting fused (grouped GEMM) MoE experts for architectures without an ``experts.linear_fc1`` rule (e.g. ``Qwen3MoeForCausalLM``), which produced a valid-looking checkpoint containing no expert weights. ``quantize.py`` and the exporter now raise instead; pass ``--no_moe_grouped_gemm`` to export these models.
+- Fix Megatron-Core HuggingFace export silently omitting fused (grouped GEMM) MoE experts for architectures without an ``experts.linear_fc1`` rule (e.g. ``Qwen3MoeForCausalLM``), which produced a valid-looking checkpoint containing no expert weights. The exporter now raises instead of writing that checkpoint; the scripts also avoid the situation by selecting ``SequentialMLP`` for those architectures.
 - Fix GatedDeltaNet (Qwen3.5) quantizer exclusions on Megatron-Core: the recipe patterns name the HuggingFace ``linear_attn`` module, so the ``conv1d`` was calibrated and the alpha / beta gate projections were exported in FP8. ``conv1d`` now has a ``self_attention`` alias in the default disabled-quantizer units, and the alpha / beta projections are exported in BF16 (they share Megatron's fused ``in_proj`` quantizer and cannot be disabled by name).
 - Megatron-Core HuggingFace export now verifies its own output: if the exported checkpoint is missing tensors the source checkpoint has (an architecture whose module has no export rule), it raises instead of writing a valid-looking checkpoint. Depth-pruned models and tied embeddings are accounted for.
 - Loading a Megatron checkpoint that holds quantizer tensors but no restorable ModelOpt state now raises instead of silently loading the model unquantized.
diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md
index 919246bf9d3..06b02d8e804 100644
--- a/examples/megatron_bridge/README.md
+++ b/examples/megatron_bridge/README.md
@@ -120,7 +120,7 @@ For a vision-language model (e.g. Qwen3.5-VL, Gemma3-VL), `quantize.py` automati
 - A **text** dataset runs text-only calibration of the language model (vision tower idle).
 
 > [!NOTE]
-> HuggingFace unified export (`export_quantized_megatron_to_hf.py`) of a quantized VLM covers **Qwen3-VL** only -- the architecture must be in ModelOpt's Megatron export mapping. For other VLMs (Qwen3.5-VL, Gemma3-VL) the quantized model is saved in Megatron checkpoint format only.
+> HuggingFace unified export (`export_quantized_megatron_to_hf.py`) of a quantized VLM covers **Qwen3-VL** and **Qwen3.5-VL**. Other VLMs such as Gemma3-VL are saved in Megatron checkpoint format only.
 
 ## Distillation
 
diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py
index 0151921c4df..fcb83ea91ca 100644
--- a/examples/megatron_bridge/distill.py
+++ b/examples/megatron_bridge/distill.py
@@ -56,7 +56,12 @@
 import modelopt.torch.utils.distributed as dist
 from modelopt.torch.opt.conversion import ModeloptStateManager
 from modelopt.torch.utils import print_args, print_rank_0, warn_rank_0
-from modelopt.torch.utils.plugins.mbridge import is_vlm_config, load_modelopt_megatron_checkpoint
+from modelopt.torch.utils.plugins.mbridge import (
+    is_vlm_config,
+    load_modelopt_megatron_checkpoint,
+    set_moe_expert_layout,
+    use_moe_grouped_gemm,
+)
 
 with contextlib.suppress(ModuleNotFoundError):
     import modelopt.torch.puzzletron.plugins.mbridge  # noqa: F401
@@ -97,9 +102,9 @@ def get_args():
         "--no_moe_grouped_gemm",
         action="store_true",
         help=(
-            "Use SequentialMLP for MoE experts instead of the (default) efficient fused "
-            "TEGroupedMLP (grouped GEMM). Must match the checkpoint passed to "
-            "--student_megatron_path. Only affects MoE models."
+            "Force SequentialMLP for MoE experts instead of the fused TEGroupedMLP (grouped GEMM). "
+            "By default grouped GEMM is used unless the architecture cannot export it to "
+            "HuggingFace, in which case SequentialMLP is selected automatically."
         ),
     )
     parser.add_argument(
@@ -336,6 +341,12 @@ def _tokenizer_prepends_bos(args) -> bool:
 
 
 def main(args: argparse.Namespace):
+    # Same layout choice as quantize.py -- it must match --student_megatron_path.
+    moe_grouped_gemm = use_moe_grouped_gemm(
+        args.student_hf_path,
+        trust_remote_code=args.trust_remote_code,
+        force_sequential=args.no_moe_grouped_gemm,
+    )
     checkpoint_dir = os.path.join(args.output_dir, "checkpoints")
     tensorboard_dir = os.path.join(args.output_dir, "tb_logs")
 
@@ -353,9 +364,8 @@ def _build_model_provider(hf_path, load_weights=True):
         provider.expert_model_parallel_size = args.ep_size
         provider.expert_tensor_parallel_size = 1  # Expert tensor parallelism is not supported
         provider.seq_length = args.seq_length
-        if (provider.num_moe_experts or 0) > 0:
-            # Must match the expert layout of --student_megatron_path (see quantize.py).
-            provider.moe_grouped_gemm = not args.no_moe_grouped_gemm
+        # Must match the expert layout of --student_megatron_path (see quantize.py).
+        set_moe_expert_layout(provider, moe_grouped_gemm)
         if args.sft:
             # A response-only loss mask needs per-token reduction to combine across CP ranks.
             # Must stay in sync with ``average_in_collective=not args.sft`` on the DDP config.
@@ -431,11 +441,12 @@ def _restore_student_hook(model_chunks):
                 f"Loading student weights from Megatron checkpoint {args.student_megatron_path}"
             )
             student = unwrap_model(model_chunks[0])
-            load_modelopt_megatron_checkpoint([student], args.student_megatron_path)
-            if is_vlm and student_has_modelopt_state:
+            loaded = load_modelopt_megatron_checkpoint([student], args.student_megatron_path)
+            if is_vlm and student_has_modelopt_state and loaded[0] is student:
                 # PTQ stores the state on the VLM root (it quantizes and saves the whole VLM), but
                 # only ``language_model`` is distilled and checkpointed here, so move it there to
-                # keep the quantizers across the QAD checkpoint's save / restore.
+                # keep the quantizers across the QAD checkpoint's save / restore. Resuming from a
+                # language-model-only checkpoint already restores it there.
                 ModeloptStateManager.transfer_state_dict(student, student.language_model)
             return model_chunks
 
diff --git a/examples/megatron_bridge/export_distilled_megatron_to_hf.py b/examples/megatron_bridge/export_distilled_megatron_to_hf.py
index f70df8381ef..414efb2eda1 100644
--- a/examples/megatron_bridge/export_distilled_megatron_to_hf.py
+++ b/examples/megatron_bridge/export_distilled_megatron_to_hf.py
@@ -68,6 +68,7 @@
     is_vlm_config,
     load_mbridge_model_from_hf,
     load_modelopt_megatron_checkpoint,
+    use_moe_grouped_gemm,
 )
 
 # Megatron-Bridge checkpoint iteration directories use names like ``iter_0000100``.
@@ -214,6 +215,15 @@ def get_args() -> argparse.Namespace:
         "correct for homogeneous students; unused for VLMs.",
     )
     parser.add_argument("--trust_remote_code", action="store_true", help="Trust remote code")
+    parser.add_argument(
+        "--no_moe_grouped_gemm",
+        action="store_true",
+        help=(
+            "Force SequentialMLP for MoE experts instead of the fused TEGroupedMLP (grouped GEMM). "
+            "By default grouped GEMM is used unless the architecture cannot export it to "
+            "HuggingFace, in which case SequentialMLP is selected automatically."
+        ),
+    )
     parser.add_argument("--tp_size", type=int, default=1, help="Tensor parallel size")
     parser.add_argument("--pp_size", type=int, default=1, help="Pipeline parallel size")
     parser.add_argument("--ep_size", type=int, default=1, help="Expert parallel size")
@@ -236,6 +246,11 @@ def main(args: argparse.Namespace):
         _bridge, _provider, _model, full_model, _tokenizer = load_mbridge_model_from_hf(
             hf_model_name_or_path=args.student_hf_path,
             trust_remote_code=args.trust_remote_code,
+            moe_grouped_gemm=use_moe_grouped_gemm(
+                args.student_hf_path,
+                trust_remote_code=args.trust_remote_code,
+                force_sequential=args.no_moe_grouped_gemm,
+            ),
             provider_overrides={
                 "tensor_model_parallel_size": args.tp_size,
                 "pipeline_model_parallel_size": args.pp_size,
diff --git a/examples/megatron_bridge/export_quantized_megatron_to_hf.py b/examples/megatron_bridge/export_quantized_megatron_to_hf.py
index 018b13686a6..8889e1983ae 100644
--- a/examples/megatron_bridge/export_quantized_megatron_to_hf.py
+++ b/examples/megatron_bridge/export_quantized_megatron_to_hf.py
@@ -47,6 +47,7 @@
 from modelopt.torch.utils.plugins.mbridge import (
     load_mbridge_model_from_hf,
     load_modelopt_megatron_checkpoint,
+    use_moe_grouped_gemm,
 )
 
 
@@ -75,8 +76,9 @@ def get_args() -> argparse.Namespace:
         "--no_moe_grouped_gemm",
         action="store_true",
         help=(
-            "Use SequentialMLP for MoE experts instead of the (default) efficient fused "
-            "TEGroupedMLP (grouped GEMM). Only affects MoE models."
+            "Force SequentialMLP for MoE experts instead of the fused TEGroupedMLP (grouped GEMM). "
+            "By default grouped GEMM is used unless the architecture cannot export it to "
+            "HuggingFace, in which case SequentialMLP is selected automatically."
         ),
     )
     parser.add_argument(
@@ -116,7 +118,11 @@ def main(args: argparse.Namespace):
     _bridge, _provider, model, _unwrapped_model, _tokenizer = load_mbridge_model_from_hf(
         hf_model_name_or_path=args.hf_model_name_or_path,
         trust_remote_code=trust_remote_code,
-        moe_grouped_gemm=not args.no_moe_grouped_gemm,
+        moe_grouped_gemm=use_moe_grouped_gemm(
+            args.hf_model_name_or_path,
+            trust_remote_code=trust_remote_code,
+            force_sequential=args.no_moe_grouped_gemm,
+        ),
         provider_overrides={
             "tensor_model_parallel_size": 1,  # Tensor parallelism is not supported
             "pipeline_model_parallel_size": args.pp_size,
@@ -153,8 +159,7 @@ def main(args: argparse.Namespace):
     print_rank_0(
         f"Exporting to HuggingFace (unified) checkpoint at {args.export_unified_hf_path}..."
     )
-    # TODO (OMNIML-5366): Qwen3-VL is the only VLM in export_mcore_gpt_to_hf's per-arch mappings;
-    # Qwen3.5-VL / Gemma3-VL are not covered yet.
+    # TODO: Gemma3-VL is not in export_mcore_gpt_to_hf's per-arch mappings yet.
     export_mcore_gpt_to_hf(
         unwrapped_model,
         args.hf_model_name_or_path,
diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py
index 295a8f1cc16..954fa0a4702 100644
--- a/examples/megatron_bridge/quantize.py
+++ b/examples/megatron_bridge/quantize.py
@@ -62,16 +62,20 @@
 import gc
 
 import torch
-from transformers import AutoConfig, AutoProcessor
+from megatron.bridge.models.hf_pretrained.utils import is_safe_repo
+from transformers import AutoProcessor
 
 import modelopt.torch.quantization as mtq
 import modelopt.torch.utils.distributed as dist
 from modelopt.recipe import ModelOptPTQRecipe, load_recipe
 from modelopt.recipe.presets import KV_CACHE_NONE, KV_QUANT_CFG_CHOICES, QUANT_CFG_CHOICES
-from modelopt.torch.export.plugins.mcore_common import all_mcore_hf_export_mapping
 from modelopt.torch.utils import print_args, print_rank_0, warn_rank_0
 from modelopt.torch.utils.dataset_utils import get_supported_datasets
-from modelopt.torch.utils.plugins.mbridge import get_language_model, load_mbridge_model_from_hf
+from modelopt.torch.utils.plugins.mbridge import (
+    get_language_model,
+    load_mbridge_model_from_hf,
+    use_moe_grouped_gemm,
+)
 from modelopt.torch.utils.plugins.megatron_calibration import (
     get_megatron_calibration_forward_loop,
     get_megatron_vlm_calibration_forward_loop,
@@ -101,8 +105,9 @@ def get_args() -> argparse.Namespace:
         "--no_moe_grouped_gemm",
         action="store_true",
         help=(
-            "Use SequentialMLP for MoE experts instead of the (default) efficient fused "
-            "TEGroupedMLP (grouped GEMM). Only affects MoE models."
+            "Force SequentialMLP for MoE experts instead of the fused TEGroupedMLP (grouped GEMM). "
+            "By default grouped GEMM is used unless the architecture cannot export it to "
+            "HuggingFace, in which case SequentialMLP is selected automatically."
         ),
     )
     parser.add_argument(
@@ -285,10 +290,20 @@ def get_quant_config(args: argparse.Namespace) -> dict:
 
 
 def main(args: argparse.Namespace):
+    trust_remote_code = is_safe_repo(
+        trust_remote_code=args.trust_remote_code, hf_path=args.hf_model_name_or_path
+    )
+
+    moe_grouped_gemm = use_moe_grouped_gemm(
+        args.hf_model_name_or_path,
+        trust_remote_code=trust_remote_code,
+        force_sequential=args.no_moe_grouped_gemm,
+    )
+
     bridge, _provider, model, unwrapped_model, tokenizer = load_mbridge_model_from_hf(
         hf_model_name_or_path=args.hf_model_name_or_path,
-        trust_remote_code=args.trust_remote_code,
-        moe_grouped_gemm=not args.no_moe_grouped_gemm,
+        trust_remote_code=trust_remote_code,
+        moe_grouped_gemm=moe_grouped_gemm,
         provider_overrides={
             "tensor_model_parallel_size": args.tp_size,
             "pipeline_model_parallel_size": args.pp_size,
@@ -330,24 +345,6 @@ def main(args: argparse.Namespace):
         )
     print_rank_0(f"Using calibration dataset: {args.calib_dataset_name}")
 
-    # Fused (grouped GEMM) experts are only exportable for architectures with an
-    # "experts.linear_fc1" rule. Fail now rather than at export: the layout is baked into the
-    # checkpoint, so recovering means re-running this whole calibration.
-    if not args.no_moe_grouped_gemm and any(
-        hasattr(m, "experts") and not hasattr(m.experts, "local_experts")
-        for m in unwrapped_model.modules()
-        if type(m).__name__.endswith("MoELayer")
-    ):
-        arch = AutoConfig.from_pretrained(
-            args.hf_model_name_or_path, trust_remote_code=args.trust_remote_code
-        ).architectures[0]
-        if "experts.linear_fc1" not in all_mcore_hf_export_mapping.get(arch, {}):
-            raise NotImplementedError(
-                f"{arch} has fused (grouped GEMM) MoE experts, which "
-                "export_quantized_megatron_to_hf.py cannot export. Re-run with "
-                "--no_moe_grouped_gemm to build the experts as SequentialMLP instead."
-            )
-
     mtq_config = get_quant_config(args)
 
     # Quantize only the language model: disable quantizers on every top-level submodule that is not
@@ -404,7 +401,7 @@ def forward_loop(_model=None):
         # VLMs: drive the full VLM forward on image-text pairs so the language model's quantizers
         # see vision-conditioned activations (we still quantize the LM only).
         processor = AutoProcessor.from_pretrained(
-            args.hf_model_name_or_path, trust_remote_code=args.trust_remote_code
+            args.hf_model_name_or_path, trust_remote_code=trust_remote_code
         )
         forward_loop = get_megatron_vlm_calibration_forward_loop(
             unwrapped_model,  # full VLM (vision encoder + projector + language model)
@@ -439,12 +436,12 @@ def forward_loop(_model=None):
         model,
         args.export_megatron_path,
         hf_tokenizer_path=args.hf_model_name_or_path,
-        hf_tokenizer_kwargs={"trust_remote_code": args.trust_remote_code},
+        hf_tokenizer_kwargs={"trust_remote_code": trust_remote_code},
     )
     if is_vlm:
         print_rank_0(
             f"\nSaved quantized VLM to {args.export_megatron_path} in Megatron format. To deploy this "
-            "model, convert it to a Unified HF ckpt with export_quantized_megatron_to_hf.py (Qwen3-VL only)."
+            "model, convert it to a Unified HF ckpt with export_quantized_megatron_to_hf.py."
         )
     else:
         print_rank_0(
diff --git a/modelopt/torch/export/plugins/mcore_qwen35vl.py b/modelopt/torch/export/plugins/mcore_qwen35vl.py
index a3e225cba84..435bf5d143b 100644
--- a/modelopt/torch/export/plugins/mcore_qwen35vl.py
+++ b/modelopt/torch/export/plugins/mcore_qwen35vl.py
@@ -18,7 +18,8 @@
 Qwen3.5 interleaves GatedDeltaNet linear-attention layers (fused ``in_proj``, split here into
 HF's ``in_proj_qkv`` / ``_z`` / ``_b`` / ``_a``) with gated full-attention layers, and adds MoE
 shared experts. Only the language model is exported; the vision tower is copied from HF.
-Requires ``--no_moe_grouped_gemm``: fused ``TEGroupedMLP`` experts have no gated split.
+Fused ``TEGroupedMLP`` experts have no gated split, so the scripts build these experts as
+``SequentialMLP`` automatically.
 """
 
 from .mcore_custom import (
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index 395822df176..0cc98a116b5 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -138,9 +138,8 @@ def __init__(
         moe_router_dtype: str | None = None,
     ):
         """Create a GPTModel exporter instance."""
-        # VLM wrappers (MCore ``LLaVAModel``, Megatron-Bridge ``Qwen3VLModel``, ...) keep the decoder
-        # under ``.language_model``; only that inner model is exported, the vision tower is copied
-        # over from the HF checkpoint as-is.
+        # VLM wrappers keep the decoder under ``.language_model``; only that is exported, the
+        # vision tower being copied from the HF checkpoint as-is.
         language_model = (
             model
             if isinstance(model, (GPTModel, HybridModel))
@@ -181,8 +180,7 @@ def __init__(
         self.dtype = dtype
         self.trust_remote_code = trust_remote_code
         self.arch = self._hf_config.architectures[0]
-        # A VLM's vision tower is never quantized: copy it verbatim from the HF checkpoint. ``None``
-        # means there is nothing to copy.
+        # ``None`` when there is no vision tower to copy through.
         self.vision_passthrough_prefixes = all_mcore_hf_vision_passthrough_mapping.get(
             self.arch, LLAVA_VISION_PREFIXES if self.is_multimodal else None
         )
@@ -446,10 +444,7 @@ def save_pretrained(
             self._verify_exported_keys(save_directory, pretrained_model_name_or_path)
 
     def _verify_exported_keys(self, save_directory, pretrained_model_name_or_path) -> None:
-        """Raise if the export dropped tensors the source checkpoint has.
-
-        A missing export rule emits nothing rather than failing, so the checkpoint looks valid.
-        """
+        """Raise if the export dropped tensors the source has: a missing rule emits nothing."""
         if pretrained_model_name_or_path is None or not os.path.isdir(
             str(pretrained_model_name_or_path)
         ):
@@ -463,20 +458,29 @@ def _verify_exported_keys(self, save_directory, pretrained_model_name_or_path) -
         if not source:
             return
 
+        # Narrow on purpose: compare module prefixes, not tensor names, since a quantized source
+        # carries extras with no export counterpart, and only inside decoder layers, whose naming
+        # is stable. A dropped decoder module is the case that loads fine and produces garbage.
         num_layers = self.model.config.num_layers
+        exported_modules = {key.rsplit(".", 1)[0] for key in exported}
         missing = set()
         for key in source - exported:
             layer = re.search(r"\.layers\.(\d+)\.", key)
-            if layer is not None and int(layer.group(1)) >= num_layers:
+            if layer is None:
+                continue  # see the note above: decoder layers only
+            if int(layer.group(1)) >= num_layers:
                 continue  # depth-pruned model: the source has layers this export does not
-            if key == "lm_head.weight" and self.model.share_embeddings_and_output_weights:
-                continue  # tied embeddings: no separate output layer to export
+            if key.rsplit(".", 1)[0] in exported_modules:
+                continue  # module is exported; this name is a source-side quantization artifact
+            if "rotary_emb" in key:
+                continue  # non-persistent buffer some conversions still ship
             missing.add(key)
         if missing:
             raise RuntimeError(
                 f"Export dropped {len(missing)} tensor(s) present in "
-                f"{pretrained_model_name_or_path}, e.g. {sorted(missing)[:8]}. This usually means "
-                "an architecture has no export rule for one of its modules."
+                f"{pretrained_model_name_or_path}, e.g. {sorted(missing)[:8]}. The checkpoint "
+                f"written to {save_directory} is incomplete -- the architecture has no export "
+                "rule for one of its decoder modules."
             )
 
     @property
@@ -683,8 +687,7 @@ def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False):
                         layer.mlp.experts.linear_fc2, layer_id, is_mtp=is_mtp
                     )
                 else:
-                    # Without this the routed experts are silently dropped and the exported
-                    # checkpoint looks valid but has no expert weights.
+                    # Otherwise the routed experts are dropped and the checkpoint looks valid.
                     raise NotImplementedError(
                         f"No export rule for {type(layer.mlp.experts).__name__} experts of "
                         f"{self.arch}: fused (grouped GEMM) experts need an 'experts.linear_fc1' "
@@ -1126,7 +1129,13 @@ def _name_remapping(
 
         weight = name_to_value.pop("weight")
         if zero_centered_gamma:
-            # Megatron stores this norm's gamma centered on 0; HF centers it on 1.
+            # Megatron centres this gamma on 0, HF on 1. Assert, don't derive: the config flag
+            # is model-wide while this one is per-norm.
+            module_config = getattr(module, "config", None)
+            assert getattr(module_config, "layernorm_zero_centered_gamma", True), (
+                f"{prefix} is mapped as zero-centered gamma but its config disables it; "
+                "exporting would shift the weights by 1.0"
+            )
             weight = weight + 1.0
         weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat)
 
@@ -1431,8 +1440,8 @@ def _qkv_slicing(
         head_num = config.num_attention_heads
         head_size = config.kv_channels
         heads_per_group = head_num // num_query_groups
-        # Gated attention (e.g. Qwen3.5) packs a per-head output gate next to every query head, so a
-        # group holds [q, gate, k, v] instead of [q, k, v]. HF keeps the gate inside ``q_proj``.
+        # Gated attention (Qwen3.5) packs a gate beside every query head, so a group holds
+        # [q, gate, k, v]; HF keeps the gate inside ``q_proj``.
         output_gate = getattr(config, "attention_output_gate", False)
         group_dim = (2 * heads_per_group if output_gate else heads_per_group) + 2
         qkv_total_dim = num_query_groups * group_dim
@@ -1487,14 +1496,16 @@ def _qkv_slicing(
         slices = [q_slice, k_slice, v_slice]
         prefixes = [q_proj_prefix, k_proj_prefix, v_proj_prefix]
 
-        def _take(tensor, index, last_dim):
-            """Gather ``index`` heads; for q under gated attention also append the gate heads."""
+        def _take(tensor, index, last_dim, with_gate=False):
+            """Gather ``index`` heads, appending the gate heads for q under gated attention."""
             taken = tensor[index]
-            if output_gate and index is q_slice:
+            if with_gate:
                 taken = torch.cat([taken, tensor[gate_slice]], dim=1)
             return taken.reshape(-1, last_dim)
 
-        proj_weights = [_take(weight, s, hidden_size) for s in slices]
+        gated = [output_gate, False, False]  # q carries the gate; k and v do not
+
+        proj_weights = [_take(weight, s, hidden_size, g) for s, g in zip(slices, gated)]
         proj_keys = [p + "weight" for p in prefixes]
 
         if weight_scale is None:
@@ -1509,8 +1520,8 @@ def _take(tensor, index, last_dim):
                     [per_rank_qkv_dim, head_size, weight_scale_hidden_size]
                 )
                 proj_weight_scales = [
-                    _take(weight_scale, s, weight_scale_hidden_size).to(dtype=weight_scale_dtype)
-                    for s in slices
+                    _take(weight_scale, s, weight_scale_hidden_size, g).to(dtype=weight_scale_dtype)
+                    for s, g in zip(slices, gated)
                 ]
             else:
                 # per-tensor scaling
@@ -1546,7 +1557,7 @@ def _take(tensor, index, last_dim):
                 # Slice bias similar to weight
                 bias = val.detach().clone()
                 bias = bias.reshape([per_rank_qkv_dim, head_size])
-                proj_biases = [_take(bias, s, 1).reshape(-1) for s in slices]
+                proj_biases = [_take(bias, s, 1, g).reshape(-1) for s, g in zip(slices, gated)]
                 proj_bias_keys = [q_proj_prefix + key, k_proj_prefix + key, v_proj_prefix + key]
                 for bias_tensor, bias_key in zip(proj_biases, proj_bias_keys):
                     self._state_dict[bias_key] = bias_tensor
diff --git a/modelopt/torch/utils/plugins/mbridge.py b/modelopt/torch/utils/plugins/mbridge.py
index c9120269ba1..7fd36c906ee 100644
--- a/modelopt/torch/utils/plugins/mbridge.py
+++ b/modelopt/torch/utils/plugins/mbridge.py
@@ -14,6 +14,7 @@
 # limitations under the License.
 """Megatron-Bridge plugins for using with Model-Optimizer."""
 
+from functools import cache
 from typing import Any
 
 from megatron.bridge import AutoBridge
@@ -33,31 +34,91 @@
 from torch.distributed.checkpoint import FileSystemReader
 from transformers import AutoConfig, AutoTokenizer
 
+from modelopt.torch.export.plugins.mcore_common import all_mcore_hf_export_mapping
 from modelopt.torch.nas.plugins.megatron import get_te_hybrid_stack_spec
-from modelopt.torch.utils import print_rank_0
+from modelopt.torch.utils import print_rank_0, warn_rank_0
 
 __all__ = [
     "get_language_model",
     "is_vlm_config",
     "load_mbridge_model_from_hf",
     "load_modelopt_megatron_checkpoint",
+    "set_moe_expert_layout",
+    "use_moe_grouped_gemm",
 ]
 
 
 def get_language_model(model: MegatronModule) -> tuple[MegatronModule, bool]:
     """Return ``(language_model, is_vlm)``; VLM wrappers nest it under ``.language_model``.
 
-    Must agree with :func:`is_vlm_config`, which answers the same question before the
-    Megatron model exists.
+    Must agree with :func:`is_vlm_config`, which answers the same question before the Megatron
+    model exists, and with ``GPTModelExporter``, which repeats this because ``modelopt.torch.export``
+    cannot import this module.
     """
+    if isinstance(model, GPTModel | HybridModel):
+        return model, False
     language_model = getattr(model, "language_model", None)
     return (model, False) if language_model is None else (language_model, True)
 
 
+@cache
+def _hf_config(hf_model_name_or_path: str, trust_remote_code: bool = False):
+    """Cached ``AutoConfig`` for the helpers below; a script asks the same questions repeatedly.
+
+    Shared, so treat the result as read-only.
+    """
+    return AutoConfig.from_pretrained(hf_model_name_or_path, trust_remote_code=trust_remote_code)
+
+
 def is_vlm_config(hf_model_name_or_path: str, trust_remote_code: bool = False) -> bool:
     """Whether a HuggingFace checkpoint describes a VLM, from its config alone."""
-    config = AutoConfig.from_pretrained(hf_model_name_or_path, trust_remote_code=trust_remote_code)
-    return hasattr(config, "vision_config")
+    return hasattr(_hf_config(hf_model_name_or_path, trust_remote_code), "vision_config")
+
+
+def use_moe_grouped_gemm(
+    hf_model_name_or_path: str,
+    trust_remote_code: bool = False,
+    force_sequential: bool = False,
+) -> bool:
+    """Pick the MoE expert layout: grouped GEMM unless that would not be HF-exportable.
+
+    Grouped GEMM calibrates faster, but only architectures with an ``experts.linear_fc1`` export
+    rule can be converted to HuggingFace from it. Every script that builds the model must agree,
+    since the layout is baked into the Megatron checkpoint -- hence a pure function of the config.
+    """
+    if force_sequential:
+        return False
+    config = _hf_config(hf_model_name_or_path, trust_remote_code)
+    text_config = getattr(config, "text_config", config)
+    is_moe = any(
+        getattr(text_config, name, None)
+        for name in ("num_experts", "num_local_experts", "n_routed_experts")
+    )
+    if not is_moe:
+        return True  # ignored for dense models
+    architectures = getattr(config, "architectures", None) or [""]
+    exportable = "experts.linear_fc1" in all_mcore_hf_export_mapping.get(architectures[0], {})
+    if not exportable:
+        warn_rank_0(
+            f"{architectures[0]} has no export rule for fused (grouped GEMM) MoE experts; "
+            "building them as SequentialMLP so the checkpoint stays exportable."
+        )
+    return exportable
+
+
+def set_moe_expert_layout(provider, moe_grouped_gemm: bool) -> None:
+    """Apply the MoE expert layout to a provider, hybrid stack spec included.
+
+    Set ``moe_grouped_gemm`` on the provider (the bridge's native, possibly custom/hybrid spec
+    reads it at build time) rather than replacing the whole layer spec -- overwriting it would
+    drop custom layers (e.g. Qwen3.5's GatedDeltaNet or Gemma3's custom spec). A hybrid provider
+    additionally needs its stack spec rebuilt, since the native one pins ``TEGroupedMLP``.
+    """
+    if isinstance(provider, HybridModelProvider):
+        provider.hybrid_stack_spec = get_te_hybrid_stack_spec(moe_grouped_gemm=moe_grouped_gemm)
+        provider.moe_grouped_gemm = moe_grouped_gemm
+    elif (provider.num_moe_experts or 0) > 0:
+        provider.moe_grouped_gemm = moe_grouped_gemm
 
 
 def load_mbridge_model_from_hf(
@@ -106,14 +167,7 @@ def load_mbridge_model_from_hf(
             assert hasattr(provider, key), f"{type(provider)} does not have attribute {key}"
             setattr(provider, key, value)
 
-    # Set moe_grouped_gemm on the provider (the bridge's native, possibly custom/hybrid spec reads
-    # it at build time) rather than replacing the whole layer spec -- overwriting it would drop
-    # custom layers (e.g. Qwen3.5's GatedDeltaNet + gated-attention or Gemma3's custom spec).
-    if isinstance(provider, HybridModelProvider):
-        provider.hybrid_stack_spec = get_te_hybrid_stack_spec(moe_grouped_gemm=moe_grouped_gemm)
-        provider.moe_grouped_gemm = moe_grouped_gemm
-    elif (provider.num_moe_experts or 0) > 0:
-        provider.moe_grouped_gemm = moe_grouped_gemm
+    set_moe_expert_layout(provider, moe_grouped_gemm)
     provider.finalize()
     if init_model_parallel:
         provider.initialize_model_parallel(seed=0)
@@ -145,14 +199,9 @@ def _checkpoint_keys(checkpoint_path: str) -> list[str]:
     return list(FileSystemReader(checkpoint_path).read_metadata().state_dict_metadata)
 
 
-def _has_vision_model_weights(checkpoint_path: str) -> bool:
-    """Whether a Megatron distributed checkpoint holds a VLM's vision tower (``vision_model.*``)."""
-    return any(key.startswith("vision_model.") for key in _checkpoint_keys(checkpoint_path))
-
-
 def load_modelopt_megatron_checkpoint(
     model: list[MegatronModule], megatron_path: str, restore_modelopt_state: bool = True
-) -> None:
+) -> list[MegatronModule]:
     """Load Megatron checkpoint weights (with modelopt_state).
 
     Args:
@@ -161,15 +210,21 @@ def load_modelopt_megatron_checkpoint(
         restore_modelopt_state: Whether to restore the ModelOpt state (e.g. quantizers) before loading
             weights. Set ``False`` to load weights only -- e.g. to reload a full-precision distilled
             student without reconstructing the ``kd_loss`` mode (which would require a teacher model).
+
+    Returns:
+        The modules loaded into: ``.language_model`` for a language-model-only VLM checkpoint,
+        otherwise the modules passed in. Callers that then move the ModelOpt state need this to
+        know where it landed.
     """
     # _load_model_weights_from_checkpoint does not resolve the latest iter_* directory, so resolve it explicitly
     checkpoint_path = _get_modelopt_checkpoint_path(megatron_path)
 
     # ``distill.py`` distills a VLM's language model only, so its checkpoint holds no ``vision_model.*``
     # weights and must be loaded into ``.language_model`` rather than the full VLM wrapper.
+    checkpoint_keys = _checkpoint_keys(checkpoint_path)
     unwrapped_model = unwrap_model(model)
-    if any(get_language_model(m)[1] for m in unwrapped_model) and not _has_vision_model_weights(
-        checkpoint_path
+    if any(get_language_model(m)[1] for m in unwrapped_model) and not any(
+        key.startswith("vision_model.") for key in checkpoint_keys
     ):
         print_rank_0("Language-model-only checkpoint: loading into the VLM's `.language_model`.")
         model = [get_language_model(m)[0] for m in unwrapped_model]
@@ -179,7 +234,7 @@ def load_modelopt_megatron_checkpoint(
     if restore_modelopt_state:
         if has_modelopt_state(megatron_path):
             load_modelopt_state(model, megatron_path)
-        elif any("_quantizer." in key for key in _checkpoint_keys(checkpoint_path)):
+        elif any("_quantizer." in key for key in checkpoint_keys):
             # The quantizers cannot be rebuilt without the state, and the loader ignores the
             # leftover amax tensors, so the model would silently load unquantized.
             raise RuntimeError(
@@ -188,3 +243,4 @@ def load_modelopt_megatron_checkpoint(
                 "produced it."
             )
     _load_model_weights_from_checkpoint(checkpoint_path, model)
+    return model
diff --git a/tests/_test_utils/torch/transformers_models.py b/tests/_test_utils/torch/transformers_models.py
index cf75e50e107..5f8ee749aab 100644
--- a/tests/_test_utils/torch/transformers_models.py
+++ b/tests/_test_utils/torch/transformers_models.py
@@ -447,6 +447,31 @@ def create_tiny_nemotron_dir(
 
 
 ##### NEMOTRON-H (Mamba + Attention + MoE/MLP hybrid) #####
+def _match_released_nemotron_h_embedding_name() -> None:
+    """Save NemotronH's embedding as ``backbone.embeddings`` like released checkpoints do.
+
+    transformers <= 5.15 renames it to the legacy singular on save, so a saved fixture disagrees
+    with every real checkpoint on one tensor. Dropped upstream in 5.16, where this is a no-op
+    since there is no such rule left to remove.
+    """
+    # Optional dependency: conversion_mapping is transformers 5.x only, while NemotronH also
+    # exists on the 4.57 floor of the support matrix.
+    try:
+        from transformers.conversion_mapping import (
+            get_checkpoint_conversion_mapping,
+            register_checkpoint_conversion_mapping,
+        )
+    except ImportError:
+        return
+
+    mapping = get_checkpoint_conversion_mapping("nemotron_h")
+    if mapping is None:
+        return
+    kept = [m for m in mapping if getattr(m, "source_patterns", None) != ["embedding.weight"]]
+    if len(kept) != len(mapping):
+        register_checkpoint_conversion_mapping("nemotron_h", kept, overwrite=True)
+
+
 def get_tiny_nemotron_h(**config_kwargs) -> PreTrainedModel:
     set_seed(SEED)
 
@@ -454,6 +479,8 @@ def get_tiny_nemotron_h(**config_kwargs) -> PreTrainedModel:
     # module is imported broadly (including by the min-transformers CI job).
     from transformers import NemotronHConfig
 
+    _match_released_nemotron_h_embedding_name()
+
     # Tiny NemotronH hybrid. hybrid_override_pattern letters: M=Mamba, E=MoE/FFN, *=Attention.
     # "ME*E" matches the NemotronH default and exercises Mamba + MoE + attention layers.
     kwargs = {
diff --git a/tests/examples/megatron_bridge/test_qad.py b/tests/examples/megatron_bridge/test_qad.py
index 79fe0758d21..11907936210 100644
--- a/tests/examples/megatron_bridge/test_qad.py
+++ b/tests/examples/megatron_bridge/test_qad.py
@@ -24,61 +24,36 @@
     assert_no_quantizers_matching,
 )
 from _test_utils.torch.transformers_models import (
-    create_tiny_gemma3vl_dir,
     create_tiny_qwen3_5_moe_vl_dir,
     create_tiny_qwen3_dir,
-    create_tiny_qwen3vl_dir,
 )
 
 
 @pytest.mark.timeout(720)  # Multiple steps in one test hence takes longer than the default timeout
 @pytest.mark.parametrize(
-    ("create_student", "exports_hf", "no_moe_grouped_gemm"),
+    "create_student",
     [
-        (lambda tmp_path: create_tiny_qwen3_dir(tmp_path, with_tokenizer=True), True, False),
-        # Qwen3-VL is the only VLM architecture in the Megatron HF export mapping, so it is the one
-        # VLM case that runs the export step end-to-end.
-        (lambda tmp_path: create_tiny_qwen3vl_dir(tmp_path, with_tokenizer=True), True, False),
-        # Dense-VLM QAD path; the MoE VLM below covers it in CI, so run this one on demand only.
-        pytest.param(
-            lambda tmp_path: create_tiny_gemma3vl_dir(
-                tmp_path,
-                with_processor=True,
-                num_hidden_layers=2,
-                intermediate_size=128,
-                max_position_embeddings=512,
-            ),
-            False,
-            False,
-            marks=pytest.mark.manual,
-        ),
+        lambda tmp_path: create_tiny_qwen3_dir(tmp_path, with_tokenizer=True),
         pytest.param(
             lambda tmp_path: create_tiny_qwen3_5_moe_vl_dir(
                 tmp_path,
                 with_processor=True,
-                # Cover both Qwen3.5 decoder kinds at the same layer count: auto-generated
-                # layer_types would give linear attention only at this depth.
+                # Cover both Qwen3.5 decoder kinds at the same layer count
                 num_hidden_layers=2,
                 layer_types=["linear_attention", "full_attention"],
             ),
-            True,
-            # Gated MoE experts are only exportable as SequentialMLP; grouped GEMM raises.
-            True,
         ),
     ],
-    ids=["qwen3", "qwen3vl", "gemma3vl", "qwen3_5_moe_vl"],
+    ids=["qwen3", "qwen3_5_moe_vl"],
 )
-def test_qad(tmp_path: Path, num_gpus, create_student, exports_hf, no_moe_grouped_gemm):
+def test_qad(tmp_path: Path, num_gpus, create_student):
     """Quantize a tiny model, run QAD from the quantized student, and export the result.
 
-    For VLMs only the language model is quantized and distilled (vision tower / projector untouched),
-    and a text calibration dataset infers text-only LM calibration. Quantized-HF export needs the
-    architecture in the Megatron export mapping, so a case missing there (Gemma3-VL) stops at the
-    distilled Megatron checkpoint and only verifies the ModelOpt (quantize) state survived.
+    Covers what only QAD exercises: that the ModelOpt state survives distillation. Per-architecture
+    export is covered more cheaply by test_quantize_export.py, so keep this to one LLM and one VLM.
     """
     hf_model_path = create_student(tmp_path)
     is_vlm = "vision_config" in (hf_model_path / "config.json").read_text()
-    moe_flag = ["--no_moe_grouped_gemm"] if no_moe_grouped_gemm else []
     quantized_megatron_path = tmp_path / "quantized_megatron"
     distill_output_dir = tmp_path / "qad_output"
     train_iters = 3
@@ -86,7 +61,7 @@ def test_qad(tmp_path: Path, num_gpus, create_student, exports_hf, no_moe_groupe
 
     # Step 1: PTQ the (language) model to FP8 and save a Megatron checkpoint carrying the ModelOpt state.
     quantize_cmd = extend_cmd_parts(
-        ["torchrun", f"--nproc_per_node={num_gpus}", "quantize.py", "--skip_generate", *moe_flag],
+        ["torchrun", f"--nproc_per_node={num_gpus}", "quantize.py", "--skip_generate"],
         hf_model_name_or_path=hf_model_path,
         recipe="general/ptq/fp8_default-kv_fp8",
         tp_size=num_gpus,
@@ -106,7 +81,7 @@ def test_qad(tmp_path: Path, num_gpus, create_student, exports_hf, no_moe_groupe
     # quantizers) and distill from the (unquantized) HF teacher. The distilled checkpoint must keep the
     # ModelOpt state so the quantizers survive distillation.
     distill_cmd = extend_cmd_parts(
-        ["torchrun", f"--nproc_per_node={num_gpus}", "distill.py", "--use_mock_data", *moe_flag],
+        ["torchrun", f"--nproc_per_node={num_gpus}", "distill.py", "--use_mock_data"],
         student_hf_path=hf_model_path,
         student_megatron_path=quantized_megatron_path,
         teacher_hf_path=hf_model_path,
@@ -132,9 +107,6 @@ def test_qad(tmp_path: Path, num_gpus, create_student, exports_hf, no_moe_groupe
     assert (distilled_megatron_path / "iter_0000001").is_dir()
     assert_has_modelopt_state(distilled_megatron_path)
 
-    if not exports_hf:
-        return  # architecture missing from the export mapping; stop at the distilled checkpoint
-
     # Step 3: export the distilled quantized checkpoint to a unified HF checkpoint. hf_quant_config.json
     # is only written for a quantized model, so its presence confirms the quantizers survived QAD.
     hf_export_path = tmp_path / "qad_fp8_hf"
@@ -143,7 +115,6 @@ def test_qad(tmp_path: Path, num_gpus, create_student, exports_hf, no_moe_groupe
             "torchrun",
             f"--nproc_per_node={num_gpus}",
             "export_quantized_megatron_to_hf.py",
-            *moe_flag,
         ],
         hf_model_name_or_path=hf_model_path,
         megatron_path=distilled_megatron_path,
diff --git a/tests/examples/megatron_bridge/test_quantize_export.py b/tests/examples/megatron_bridge/test_quantize_export.py
index f4cc7098805..61f258c3759 100644
--- a/tests/examples/megatron_bridge/test_quantize_export.py
+++ b/tests/examples/megatron_bridge/test_quantize_export.py
@@ -16,28 +16,49 @@
 
 from pathlib import Path
 
+import pytest
 from _test_utils.examples.run_command import extend_cmd_parts, run_example_command
 from _test_utils.torch.export.unified_checkpoint import assert_exported_checkpoint_matches
 from _test_utils.torch.megatron.modelopt_state import assert_has_modelopt_state
-from _test_utils.torch.transformers_models import create_tiny_qwen3_dir
+from _test_utils.torch.transformers_models import (
+    create_tiny_nemotron_h_dir,
+    create_tiny_qwen3_dir,
+    create_tiny_qwen3_moe_dir,
+    create_tiny_qwen3vl_dir,
+)
 
+# Per-architecture export *mappings* are covered in-process by
+# tests/gpu_megatron/torch/export/test_unified_export_megatron.py; these cases cover the script
+# wiring (CLI, recipe, checkpoint hand-off) that only running quantize.py + the exporter exercises.
+# Use a vLLM-friendly head_dim (64): the default tiny config (head_dim=2) is unsupported.
+_DENSE_KWARGS = {
+    "hidden_size": 128,
+    "num_attention_heads": 2,
+    "num_key_value_heads": 2,
+    "num_hidden_layers": 2,
+    "intermediate_size": 256,
+    "max_position_embeddings": 512,
+}
 
-# NOTE: Qwen3.5-VL covered by test_qad.py
-def test_quantize_and_export(tmp_path: Path, num_gpus):
-    """Quantize a tiny Qwen3 via a YAML recipe and export it to a unified HF checkpoint."""
-    # Use a vLLM-friendly head_dim (64) since the default tiny config (head_dim=2) is unsupported.
-    hf_model_path = create_tiny_qwen3_dir(
-        tmp_path,
-        with_tokenizer=True,
-        hidden_size=128,
-        num_attention_heads=2,
-        num_key_value_heads=2,
-        num_hidden_layers=2,
-        intermediate_size=256,
-        max_position_embeddings=512,
-    )
-    megatron_path = tmp_path / "qwen3_fp8_megatron"
-    hf_export_path = tmp_path / "qwen3_fp8_hf"
+
+@pytest.mark.parametrize(
+    ("create_model", "model_kwargs"),
+    [
+        (create_tiny_qwen3_dir, _DENSE_KWARGS),
+        # MoE: routed experts used to be dropped silently from the export.
+        (create_tiny_qwen3_moe_dir, _DENSE_KWARGS),
+        # Dense VLM: only the language model is quantized, vision is copied through.
+        (create_tiny_qwen3vl_dir, {}),
+        # Mamba hybrid + MoE, and the one architecture that keeps grouped-GEMM experts.
+        (create_tiny_nemotron_h_dir, {}),
+    ],
+    ids=["qwen3", "qwen3_moe", "qwen3vl", "nemotron_h"],
+)
+def test_quantize_and_export(tmp_path: Path, num_gpus, create_model, model_kwargs):
+    """Quantize a tiny model via a YAML recipe and export it to a unified HF checkpoint."""
+    hf_model_path = create_model(tmp_path, with_tokenizer=True, **model_kwargs)
+    megatron_path = tmp_path / "fp8_megatron"
+    hf_export_path = tmp_path / "fp8_hf"
 
     # Step 1: quantize and save a Megatron checkpoint
     quantize_cmd = extend_cmd_parts(
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
new file mode 100644
index 00000000000..11421b53e80
--- /dev/null
+++ b/tests/gpu_megatron/torch/export/plugins/test_moe_layout_choice.py
@@ -0,0 +1,40 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+"""The MoE expert layout must be chosen identically by every script that builds the model."""
+
+import pytest
+
+from modelopt.torch.export.plugins.mcore_common import all_mcore_hf_export_mapping
+
+
+@pytest.mark.parametrize(
+    ("arch", "grouped_is_exportable"),
+    [
+        # Only Nemotron-H maps fused grouped-GEMM experts, so it keeps the faster layout.
+        ("NemotronHForCausalLM", True),
+        ("Qwen3MoeForCausalLM", False),
+        ("DeepseekV3ForCausalLM", False),
+        ("GptOssForCausalLM", False),
+        ("Llama4ForConditionalGeneration", False),
+        ("Qwen3_5MoeForConditionalGeneration", False),
+    ],
+)
+def test_grouped_gemm_exportability(arch, grouped_is_exportable):
+    assert arch in all_mcore_hf_export_mapping, f"{arch} is no longer an exported architecture"
+    has_rule = "experts.linear_fc1" in all_mcore_hf_export_mapping[arch]
+    assert has_rule is grouped_is_exportable, (
+        f"{arch} grouped-GEMM exportability changed; use_moe_grouped_gemm would now pick the "
+        "other expert layout, which silently breaks checkpoints written by the previous default"
+    )

From 7f745d9302fca67fb4dd177e5fde353ac94611ca Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Sat, 29 Aug 2026 10:54:19 -0700
Subject: [PATCH 09/12] Fix Qwen3.5 export: packed experts, NVFP4 packing, MTP
 names, repo ids

Validating Qwen3.5 export against real checkpoints surfaced four exporter bugs.

Packed routed experts. Real Qwen3.5 stores experts packed as [num_experts, out, in];
the mapping emitted per-expert names, so every routed expert was dropped. Adds a
`transpose` kwarg to `_pack_name_remapping` (Qwen3.5 keeps Megatron's orientation) and
a `GroupedMLPPacking` rule so fused TEGroupedMLP reaches the same packed tensors. This
also lets `use_moe_grouped_gemm` keep grouped GEMM for Qwen3.5, which measured
22.1 GB/GPU vs 38.9 GB/GPU for SequentialMLP on a 20-layer 256-expert model.

NVFP4 grouped packing. `_grouped_mlp_packing` max-merged `weight_scale`, but NVFP4
needs each expert's per-block scales stacked with only the global `weight_scale_2`
merged; it also dequantized packed uint8 against per-block scales and passed
`block_size=None`. `_grouped_mlp_slicing` gains `quantize=False` to emit raw weights
plus `(qformat, block_size)`, and packing now quantizes once over the stack, matching
`_pack_name_remapping`. Scale suffixes are aligned to `_weight_scale` so both packed
paths agree.

MTP names. `_mtp_prefix` replaced every occurrence of "model", so a VLM prefix
`model.language_model.layers.{}` became `mtp.language_mtp.layers.0.*` -- tensors
present and correct, under names nothing loads. Only the root segment is rewritten now.
`_get_mtp_state_dict` also assumed `mtp_model_layer.layers`, which Qwen3.5 does not have.

Repo ids. `load_multimodal_components` rejected HF repo ids, so `quantize.py` accepted
`Qwen/Qwen3.5-0.8B` but the documented export step failed with "It should be a
directory". It now resolves them via `snapshot_download`, as its sibling in the same
file already did. This affected every VLM export.

Registers `Qwen3_5ForConditionalGeneration` (dense) for export and vision passthrough,
and fixes `with_language_model_prefix` crashing on non-mapping flags.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .../export/plugins/hf_checkpoint_utils.py     |  20 +++-
 modelopt/torch/export/plugins/mcore_common.py |   2 +
 modelopt/torch/export/plugins/mcore_custom.py |  15 +++
 .../torch/export/plugins/mcore_qwen35vl.py    |  18 ++-
 .../torch/export/unified_export_megatron.py   | 108 +++++++++++++++---
 5 files changed, 139 insertions(+), 24 deletions(-)

diff --git a/modelopt/torch/export/plugins/hf_checkpoint_utils.py b/modelopt/torch/export/plugins/hf_checkpoint_utils.py
index 54058c742db..b2046fe081a 100644
--- a/modelopt/torch/export/plugins/hf_checkpoint_utils.py
+++ b/modelopt/torch/export/plugins/hf_checkpoint_utils.py
@@ -205,7 +205,7 @@ def load_multimodal_components(
     """Load multimodal components from safetensors file.
 
     Args:
-        pretrained_model_path: Path to the pretrained model.
+        pretrained_model_path: Directory or HuggingFace repo id of the pretrained model.
         prefixes: Tensor key prefixes to select.  Defaults to the LLaVA-style
             ``multi_modal_projector`` / ``vision_model`` prefixes.  Pass
             ``("model.visual.",)`` for Qwen3-VL checkpoints.
@@ -215,9 +215,21 @@ def load_multimodal_components(
     """
     hf_checkpoint_path = Path(pretrained_model_path)
     if not hf_checkpoint_path.is_dir():
-        raise ValueError(
-            f"Invalid pretrained model path: {pretrained_model_path}. It should be a directory."
-        )
+        # Also accept a repo id, which is what the example scripts pass to quantize.py.
+        local_files_only = _is_hf_hub_offline()
+        try:
+            hf_checkpoint_path = Path(
+                snapshot_download(
+                    repo_id=str(pretrained_model_path),
+                    allow_patterns=["*.safetensors", "*.safetensors.index.json"],
+                    local_files_only=local_files_only,
+                )
+            )
+        except (LocalEntryNotFoundError, OSError, ValueError) as exc:
+            raise ValueError(
+                f"Invalid pretrained model path: {pretrained_model_path}. It should be a "
+                "directory or an available HuggingFace repo id."
+            ) from exc
 
     safetensors_file = Path(hf_checkpoint_path) / "model.safetensors"
     safetensors_index_file = Path(hf_checkpoint_path) / "model.safetensors.index.json"
diff --git a/modelopt/torch/export/plugins/mcore_common.py b/modelopt/torch/export/plugins/mcore_common.py
index 4325a25f54e..d275aee4c0a 100644
--- a/modelopt/torch/export/plugins/mcore_common.py
+++ b/modelopt/torch/export/plugins/mcore_common.py
@@ -61,6 +61,7 @@
     "Qwen2ForCausalLM": qwen25_causal_lm_export,
     "GptOssForCausalLM": gptoss_causal_lm_export,
     "Qwen3VLForConditionalGeneration": qwen3vl_causal_lm_export,
+    "Qwen3_5ForConditionalGeneration": qwen3_5_vl_causal_lm_export,
     "Qwen3_5MoeForConditionalGeneration": qwen3_5_vl_causal_lm_export,
 }
 
@@ -69,6 +70,7 @@
 # ``LLAVA_VISION_PREFIXES`` when the model is an MCore ``LLaVAModel``, and to no copy otherwise.
 all_mcore_hf_vision_passthrough_mapping: dict[str, tuple[str, ...]] = {
     "Qwen3VLForConditionalGeneration": QWEN3VL_VISION_PREFIXES,
+    "Qwen3_5ForConditionalGeneration": QWEN3_5_VL_VISION_PREFIXES,
     "Qwen3_5MoeForConditionalGeneration": QWEN3_5_VL_VISION_PREFIXES,
 }
 
diff --git a/modelopt/torch/export/plugins/mcore_custom.py b/modelopt/torch/export/plugins/mcore_custom.py
index a0144a44c7f..87b696e641f 100644
--- a/modelopt/torch/export/plugins/mcore_custom.py
+++ b/modelopt/torch/export/plugins/mcore_custom.py
@@ -128,6 +128,18 @@ def __init__(self, target_name_or_prefix: str = "", func_kwargs: dict[str, Any]
         )
 
 
+class GroupedMLPPacking(CustomModuleMapping):
+    """A custom module mapping that packs grouped MoE experts into one stacked tensor."""
+
+    def __init__(self, target_name_or_prefix: str = "", func_kwargs: dict[str, Any] = {}):
+        """Create a custom module mapping that packs grouped experts."""
+        super().__init__(
+            func_name="grouped_mlp_packing",
+            target_name_or_prefix=target_name_or_prefix,
+            func_kwargs=func_kwargs,
+        )
+
+
 class GatedMLPMerging(CustomModuleMapping):
     """A custom module mapping that merges gate_proj and up_proj."""
 
@@ -247,6 +259,9 @@ def with_language_model_prefix(
     """Nest a text-model mapping under ``model.language_model.``; other prefixes 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.") :]
diff --git a/modelopt/torch/export/plugins/mcore_qwen35vl.py b/modelopt/torch/export/plugins/mcore_qwen35vl.py
index 435bf5d143b..737ffdebb22 100644
--- a/modelopt/torch/export/plugins/mcore_qwen35vl.py
+++ b/modelopt/torch/export/plugins/mcore_qwen35vl.py
@@ -18,14 +18,15 @@
 Qwen3.5 interleaves GatedDeltaNet linear-attention layers (fused ``in_proj``, split here into
 HF's ``in_proj_qkv`` / ``_z`` / ``_b`` / ``_a``) with gated full-attention layers, and adds MoE
 shared experts. Only the language model is exported; the vision tower is copied from HF.
-Fused ``TEGroupedMLP`` experts have no gated split, so the scripts build these experts as
-``SequentialMLP`` automatically.
+Routed experts are packed as ``[num_experts, out, in]``, exported from either expert layout.
 """
 
 from .mcore_custom import (
     GatedDeltaNetSlicing,
     GatedMLPSlicing,
+    GroupedMLPPacking,
     NameRemapping,
+    PackNameRemapping,
     with_language_model_prefix,
 )
 from .mcore_qwen import qwen3_causal_lm_export
@@ -46,6 +47,19 @@
         "model.layers.{}.linear_attn.norm.", {"zero_centered_gamma": True}
     ),
     "linear_attn.out_proj": NameRemapping("model.layers.{}.linear_attn.out_proj."),
+    # Routed experts are stored packed: [num_experts, out, in], keeping Megatron's orientation.
+    "use_packed_local_experts": True,
+    "local_experts.linear_fc1": PackNameRemapping(
+        "model.layers.{}.mlp.experts.gate_up_proj",
+        {"layer_type": "linear_fc1", "transpose": False},
+    ),
+    "local_experts.linear_fc2": PackNameRemapping(
+        "model.layers.{}.mlp.experts.down_proj",
+        {"layer_type": "linear_fc2", "transpose": False},
+    ),
+    # Same packed layout from fused TEGroupedMLP, so grouped GEMM stays usable.
+    "experts.linear_fc1": GroupedMLPPacking("model.layers.{}.mlp.experts.gate_up_proj"),
+    "experts.linear_fc2": GroupedMLPPacking("model.layers.{}.mlp.experts.down_proj"),
     # MoE shared experts (routed experts + router come from the Qwen3 rules).
     "shared_experts.linear_fc1": GatedMLPSlicing("model.layers.{}.mlp.shared_expert."),
     "shared_experts.linear_fc2": NameRemapping("model.layers.{}.mlp.shared_expert.down_proj."),
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index 0cc98a116b5..281271a210e 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -710,7 +710,9 @@ def _get_mtp_state_dict(self) -> dict[str, torch.Tensor]:
         self._state_dict = OrderedDict()
         try:
             for mtp_layer in mtp.layers:
-                inner_layers = mtp_layer.mtp_model_layer.layers
+                # Some architectures (Qwen3.5) put a single TransformerLayer here, not a container.
+                inner = mtp_layer.mtp_model_layer
+                inner_layers = getattr(inner, "layers", None) or [inner]
                 first_id = inner_layers[0].layer_number - 1
                 last_id = inner_layers[-1].layer_number - 1
 
@@ -956,6 +958,7 @@ def _custom_mapping_to_lambda(mapping):
                 "gated_mlp_slicing": self._gated_mlp_slicing,
                 "gated_delta_net_slicing": self._gated_delta_net_slicing,
                 "grouped_mlp_slicing": self._grouped_mlp_slicing,
+                "grouped_mlp_packing": self._grouped_mlp_packing,
                 "pack_name_remapping": self._pack_name_remapping,
                 "pack_name_remapping_gpt_oss": self._pack_name_remapping_gpt_oss,
             }
@@ -1102,8 +1105,11 @@ def _record_excluded_module(self, prefix: str):
     def _mtp_prefix(prefix: str) -> str:
         """Rewrite a base-model target prefix (backbone/model root) to its MTP counterpart."""
         if "backbone" in prefix:
-            return prefix.replace("backbone", "mtp")
-        return prefix.replace("model", "mtp")
+            return prefix.replace("backbone", "mtp", 1)
+        # Replace the root only: a VLM's "model.language_model." must not become "mtp.language_mtp.".
+        if prefix.startswith("model.language_model."):
+            return "mtp." + prefix[len("model.language_model.") :]
+        return prefix.replace("model", "mtp", 1)
 
     def _name_remapping(
         self,
@@ -1227,9 +1233,70 @@ def _gated_mlp_slicing(
                 self._state_dict[gate_proj_key] = val.detach().clone()
                 self._state_dict[up_proj_key] = val.detach().clone()
 
-    def _grouped_mlp_slicing(self, module, prefix, parallel_config=None, is_mtp=False):
+    def _grouped_mlp_packing(self, module, prefix, parallel_config=None, is_mtp=False):
+        """Pack TEGroupedMLP experts into one ``[num_experts, out, in]`` tensor (Qwen3.5 layout).
+
+        Reuses the per-expert path for the EP gather and per-expert quantizers, then stacks by
+        global expert id and re-quantizes once with the max scale, as ``_pack_name_remapping`` does.
+        """
+        if is_mtp:
+            prefix = self._mtp_prefix(prefix)
+        marker = "\x00pack\x00"
+        saved_state_dict = self._state_dict
+        self._state_dict = OrderedDict()
+        try:
+            qformat, block_size = self._grouped_mlp_slicing(
+                module, marker + "{}", parallel_config=parallel_config, is_mtp=False, quantize=False
+            )
+            per_expert = self._state_dict
+        finally:
+            self._state_dict = saved_state_dict
+
+        def collect(suffix):
+            found = {}
+            for key, value in per_expert.items():
+                if not key.startswith(marker) or not key.endswith(suffix):
+                    continue
+                found[int(key[len(marker) :].split(".", 1)[0])] = value
+            return [found[i] for i in sorted(found)]
+
+        weights = collect(".weight")
+        if not weights:
+            return
+        scales, scales_2 = collect(".weight_scale"), collect(".weight_scale_2")
+        input_scales = collect(".input_scale")
+
+        # Quantize once over the stack, exactly as _pack_name_remapping does.
+        merged_weight = torch.stack(weights, dim=0)
+        if not scales:
+            self._state_dict[prefix] = merged_weight
+        else:
+            if scales_2:
+                # NVFP4 keeps each expert's block scales; only the global scale is merged.
+                merged_scale = torch.stack(scales, dim=0)
+                merged_scale_2 = torch.max(torch.stack(scales_2, dim=0), dim=0)[0]
+            else:
+                merged_scale, merged_scale_2 = torch.max(torch.stack(scales, dim=0), dim=0)[0], None
+            self._state_dict[prefix] = to_quantized_weight(
+                merged_weight, merged_scale, qformat, merged_scale_2, block_size
+            )
+            # Same suffixes as _pack_name_remapping so both packed paths agree.
+            self._state_dict[prefix + "_weight_scale"] = merged_scale
+            if merged_scale_2 is not None:
+                self._state_dict[prefix + "_weight_scale_2"] = merged_scale_2
+        if input_scales:
+            self._state_dict[prefix + "_input_scale"] = torch.max(
+                torch.stack(input_scales, dim=0), dim=0
+            )[0]
+
+    def _grouped_mlp_slicing(
+        self, module, prefix, parallel_config=None, is_mtp=False, quantize=True
+    ):
         """Export TEGroupedMLP weight0..weight{N-1} as one HF-style entry per expert.
 
+        ``quantize=False`` emits unquantized weights alongside the scales, which
+        ``_grouped_mlp_packing`` needs so it can quantize once over the stacked tensor.
+
         At EP>1, local ids are mapped to global via ``module.local_expert_indices``
         and per-expert state is ``all_gather_object``-ed across the EP group. All EP ranks
         MUST enter this method for the same layer in lockstep or the gather hangs.
@@ -1347,12 +1414,16 @@ def _grouped_mlp_slicing(self, module, prefix, parallel_config=None, is_mtp=Fals
                 if weight_scale_cpu is None:
                     local_expert_state[expert_prefix + "weight"] = weight
                 else:
-                    local_expert_state[expert_prefix + "weight"] = to_quantized_weight(
-                        weight,
-                        weight_scale_cpu,
-                        qformat,
-                        weight_scale_2_cpu,
-                        block_size,
+                    local_expert_state[expert_prefix + "weight"] = (
+                        weight
+                        if not quantize
+                        else to_quantized_weight(
+                            weight,
+                            weight_scale_cpu,
+                            qformat,
+                            weight_scale_2_cpu,
+                            block_size,
+                        )
                     )
                     local_expert_state[expert_prefix + "weight_scale"] = weight_scale_cpu.clone()
 
@@ -1404,6 +1475,7 @@ def _grouped_mlp_slicing(self, module, prefix, parallel_config=None, is_mtp=Fals
             del gathered_bytes
         else:
             self._state_dict.update(local_expert_state)
+        return seen_qformat, seen_block_size
 
     def _qkv_slicing(
         self,
@@ -1672,8 +1744,8 @@ def _self_attention_scaling(
                 # FP8 KV Cache is supported in VLLM; NVFP4 supported in TRTLLM
                 self.kv_cache_dtype = kv_cache_dtype
 
-    def _pack_name_remapping(self, module, prefix, layer_type=None, is_mtp=False):
-        """Pack name remapping into one tensor."""
+    def _pack_name_remapping(self, module, prefix, layer_type=None, is_mtp=False, transpose=True):
+        """Pack per-expert weights into one tensor; ``transpose`` for HF [E, in, out] layouts."""
         if is_mtp:
             prefix = self._mtp_prefix(prefix)
         weight_list = []
@@ -1700,10 +1772,10 @@ def _pack_name_remapping(self, module, prefix, layer_type=None, is_mtp=False):
 
         merged_weight = torch.stack(weight_list, dim=0)
 
-        # Transpose the last two dimensions to match HuggingFace format
-        # Megatron format: [num_experts, out_features, in_features]
-        # HF format: [num_experts, in_features, out_features]
-        merged_weight = merged_weight.transpose(-2, -1).contiguous()
+        # Megatron is [num_experts, out, in]; most HF layouts want [num_experts, in, out], but
+        # Qwen3.5 keeps Megatron's orientation.
+        if transpose:
+            merged_weight = merged_weight.transpose(-2, -1).contiguous()
 
         if weight_scale_2_list[0] is None:
             merged_weight_scale_2 = None
@@ -1715,8 +1787,8 @@ def _pack_name_remapping(self, module, prefix, layer_type=None, is_mtp=False):
             # NVFP4
             merged_weight_scale_2 = torch.max(torch.stack(weight_scale_2_list, dim=0), dim=0)[0]
             merged_weight_scale = torch.stack(weight_scale_list, dim=0)
-            # Transpose the scaling factors to match the transposed weights
-            merged_weight_scale = merged_weight_scale.transpose(-2, -1).contiguous()
+            if transpose:
+                merged_weight_scale = merged_weight_scale.transpose(-2, -1).contiguous()
 
         if input_scale_list[0] is not None:
             merged_input_scale = torch.max(torch.stack(input_scale_list, dim=0), dim=0)[0]

From c262cc6c9132baf0b3e5f076ec10285387aab345 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Sat, 29 Aug 2026 10:54:44 -0700
Subject: [PATCH 10/12] Keep grouped GEMM for unquantized distillation
 students, guard QAD export

The MoE layout must match the checkpoint being loaded, but only a quantized student
is pinned by it. distill.py applied `use_moe_grouped_gemm` unconditionally, forcing
SequentialMLP onto pruned BF16 students that Megatron-Bridge can load in either
layout -- giving up grouped GEMM's memory saving for no reason. The layout choice is
now made only when the student carries ModelOpt state.

export_distilled_megatron_to_hf.py drops quantization by design, so pointing it at a
QAD checkpoint silently produced an unquantized export. It now fails with a pointer to
export_quantized_megatron_to_hf.py. `has_modelopt_state` ignores `kd_loss`, so plain
distillation is unaffected.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 examples/megatron_bridge/distill.py           | 22 ++++++++++++-------
 .../export_distilled_megatron_to_hf.py        |  9 ++++++++
 modelopt/torch/utils/plugins/mbridge.py       |  1 -
 3 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py
index fcb83ea91ca..7df5489159a 100644
--- a/examples/megatron_bridge/distill.py
+++ b/examples/megatron_bridge/distill.py
@@ -341,11 +341,20 @@ def _tokenizer_prepends_bos(args) -> bool:
 
 
 def main(args: argparse.Namespace):
-    # Same layout choice as quantize.py -- it must match --student_megatron_path.
-    moe_grouped_gemm = use_moe_grouped_gemm(
-        args.student_hf_path,
-        trust_remote_code=args.trust_remote_code,
-        force_sequential=args.no_moe_grouped_gemm,
+    student_has_modelopt_state = args.student_megatron_path is not None and has_modelopt_state(
+        args.student_megatron_path
+    )
+    # A quantized student pins the layout: it must match what quantize.py wrote, so reuse the same
+    # data-driven choice. An unquantized (e.g. pruned) student exports via Megatron-Bridge, which
+    # reads either layout, so it keeps the faster grouped GEMM.
+    moe_grouped_gemm = (
+        use_moe_grouped_gemm(
+            args.student_hf_path,
+            trust_remote_code=args.trust_remote_code,
+            force_sequential=args.no_moe_grouped_gemm,
+        )
+        if student_has_modelopt_state
+        else not args.no_moe_grouped_gemm
     )
     checkpoint_dir = os.path.join(args.output_dir, "checkpoints")
     tensorboard_dir = os.path.join(args.output_dir, "tb_logs")
@@ -381,9 +390,6 @@ def _build_model_provider(hf_path, load_weights=True):
     # The student structure is always built from --student_hf_path. When --student_megatron_path is
     # given, the HF weights are skipped (they are overwritten by the Megatron checkpoint, loaded into
     # the built student inside the patched provide() below).
-    student_has_modelopt_state = args.student_megatron_path is not None and has_modelopt_state(
-        args.student_megatron_path
-    )
     student_provider = _build_model_provider(
         args.student_hf_path, load_weights=args.student_megatron_path is None
     )
diff --git a/examples/megatron_bridge/export_distilled_megatron_to_hf.py b/examples/megatron_bridge/export_distilled_megatron_to_hf.py
index 414efb2eda1..fbb2f0c8760 100644
--- a/examples/megatron_bridge/export_distilled_megatron_to_hf.py
+++ b/examples/megatron_bridge/export_distilled_megatron_to_hf.py
@@ -59,6 +59,7 @@
 
 import torch
 from megatron.bridge import AutoBridge
+from megatron.bridge.training.post_training.checkpointing import has_modelopt_state
 from transformers import AutoConfig
 
 import modelopt.torch.utils.distributed as dist
@@ -237,6 +238,14 @@ def get_args() -> argparse.Namespace:
 
 def main(args: argparse.Namespace):
     checkpoint_export_paths: list[tuple[Path, Path]] = _get_checkpoint_export_paths(args)
+    # This path drops quantization, so a QAD checkpoint would export silently unquantized.
+    # ``has_modelopt_state`` ignores ``kd_loss``, so plain distillation still passes.
+    quantized = [str(p) for p, _ in checkpoint_export_paths if has_modelopt_state(str(p))]
+    if quantized:
+        raise ValueError(
+            f"{quantized[0]} is quantized; this script exports full precision only and would drop "
+            "the quantizers. Use export_quantized_megatron_to_hf.py instead."
+        )
     is_vlm = is_vlm_config(args.student_hf_path, trust_remote_code=args.trust_remote_code)
 
     if is_vlm:
diff --git a/modelopt/torch/utils/plugins/mbridge.py b/modelopt/torch/utils/plugins/mbridge.py
index 7fd36c906ee..131f8f5220b 100644
--- a/modelopt/torch/utils/plugins/mbridge.py
+++ b/modelopt/torch/utils/plugins/mbridge.py
@@ -144,7 +144,6 @@ def load_mbridge_model_from_hf(
         provider_overrides: Overrides for the provider.
         init_model_parallel: Whether to initialize model parallel.
         moe_grouped_gemm: Whether to use grouped GEMM for MoE.
-            Pruning does not support grouped GEMM yet.
         load_weights: Whether to load the HF weights into the model. Set to ``False`` when the
             weights will be loaded from a Megatron checkpoint instead (e.g. for export), in which
             case only the model structure (with the correct layer spec) is built.

From 5e1a141aafc4b70c91fdd5fadeb5dd871fd5e435 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Sat, 29 Aug 2026 10:56:13 -0700
Subject: [PATCH 11/12] Cover packed-expert and quantized hybrid export in the
 fast test suite

Three of the bugs in this branch reached real checkpoints because the fast export
suite could not see them.

The Qwen3.5 fixture was unfaithful on disk in the direction that hides the bug:
transformers unpacks routed experts on `save_pretrained`, writing
`experts.0.gate_proj.weight`, while every released Qwen3.5 checkpoint stores them
packed. The saved reference therefore pushed the exporter toward per-expert names.
`_pack_qwen3_5_moe_experts` repacks after save, the same class of fix as
`_match_released_nemotron_h_embedding_name`.

Qwen3.5 had no case in this suite at all -- it appeared only in the slow example
suites, which never diff against a reference. Adds grouped GEMM x {NVFP4, FP8} and
SequentialMLP x NVFP4, so both expert layouts must reach the same packed HF tensors
via `GroupedMLPPacking` and `PackNameRemapping`. Reverting only the NVFP4 scale fix
makes the grouped case fail, so these bite.

Replaces the quantized Nemotron cases with NemotronH, the arch that superseded it and
the only one exporting fused grouped-GEMM experts; the plain/eagle/medusa cases still
smoke-test the old arch. The existing NemotronH test asserts key presence only, so it
never noticed that the hybrid builder defaults to LayerNorm and exported `norm.bias`
tensors NemotronH's RMSNorm has no counterpart for.

`assert_exported_checkpoint_matches` only understood the dotted `...proj.weight_scale`
form, so it silently mis-resolved packed scale keys; it now handles both layouts, which
is what makes `check_values=True` meaningful for packed experts.

Net suite cost is +2.4s; MTP remains uncovered here because `get_mcore_gpt_model`
cannot build it.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .../torch/export/unified_checkpoint.py        |  16 ++-
 .../_test_utils/torch/transformers_models.py  |  47 +++++++-
 .../export/plugins/test_moe_layout_choice.py  |   4 +-
 .../export/test_unified_export_megatron.py    | 108 ++++++++++++++----
 4 files changed, 148 insertions(+), 27 deletions(-)

diff --git a/tests/_test_utils/torch/export/unified_checkpoint.py b/tests/_test_utils/torch/export/unified_checkpoint.py
index e087f5f738d..e91e400f1c1 100644
--- a/tests/_test_utils/torch/export/unified_checkpoint.py
+++ b/tests/_test_utils/torch/export/unified_checkpoint.py
@@ -92,21 +92,26 @@ def _unpack_nvfp4(packed: torch.Tensor) -> torch.Tensor:
     return NVFP4QTensor.get_e2m1_values(packed.device).to(torch.float32)[codes]
 
 
+def _scale_key(key: str, suffix: str) -> str:
+    """Scale key for the dotted (``...proj.weight``) and packed (``...gate_up_proj``) layouts."""
+    return key.replace(".weight", f".{suffix}") if key.endswith(".weight") else f"{key}_{suffix}"
+
+
 def _dequantize(key: str, exported: dict[str, torch.Tensor]) -> torch.Tensor:
     """Undo the exporter's weight scaling so the value can be compared to the source."""
-    scale = exported.get(key.replace(".weight", ".weight_scale"))
+    scale = exported.get(_scale_key(key, "weight_scale"))
     if _is_packed(exported[key]):
         # NVFP4: E2M1 codes, an E4M3 scale per block along the last dim, and a global FP32 scale.
         weight = _unpack_nvfp4(exported[key])
         block_size = weight.shape[-1] // scale.shape[-1]
         weight = weight * scale.to(torch.float32).repeat_interleave(block_size, dim=-1)
-        scale_2 = exported.get(key.replace(".weight", ".weight_scale_2"))
+        scale_2 = exported.get(_scale_key(key, "weight_scale_2"))
         return weight if scale_2 is None else weight * scale_2.to(torch.float32)
     weight = exported[key].to(torch.float32)
     if scale is None:
         return weight
     scale = scale.to(torch.float32)
-    return weight * (scale if scale.ndim == 0 else scale.reshape(-1, *([1] * (weight.ndim - 1))))
+    return weight * scale.reshape(*scale.shape, *([1] * (weight.ndim - scale.ndim)))
 
 
 def assert_exported_checkpoint_matches(
@@ -175,7 +180,10 @@ def assert_exported_checkpoint_matches(
     wrong = []
     for key in shared:
         got, want = _dequantize(key, exported), reference[key].to(torch.float32)
-        if exported[key].dtype == reference[key].dtype and key + "_scale" not in exported:
+        if (
+            exported[key].dtype == reference[key].dtype
+            and _scale_key(key, "weight_scale") not in exported
+        ):
             # Copied through untouched (norms, router, vision tower): must be bit-exact.
             if not torch.equal(exported[key], reference[key]):
                 wrong.append((key, "not bit-exact"))
diff --git a/tests/_test_utils/torch/transformers_models.py b/tests/_test_utils/torch/transformers_models.py
index 5f8ee749aab..712cece2055 100644
--- a/tests/_test_utils/torch/transformers_models.py
+++ b/tests/_test_utils/torch/transformers_models.py
@@ -14,12 +14,15 @@
 # limitations under the License.
 
 import contextlib
+import re
+from collections import defaultdict
 from functools import partial
 from pathlib import Path
 
 import pytest
 import torch
 from _test_utils.torch.misc import set_seed
+from safetensors.torch import load_file, save_file
 
 transformers = pytest.importorskip("transformers")
 from transformers import (
@@ -390,6 +393,45 @@ def _get_tiny_qwen3_5_vl(moe: bool = False, **config_kwargs) -> PreTrainedModel:
     return AutoModelForImageTextToText.from_config(cfg)
 
 
+def _pack_qwen3_5_moe_experts(dir_path: Path | str) -> None:
+    """Repack a saved Qwen3.5-MoE fixture's experts the way released checkpoints store them.
+
+    transformers unpacks routed experts into per-expert ``experts.{i}.*`` on save, but every real
+    Qwen3.5 checkpoint stores them packed as ``experts.gate_up_proj`` / ``experts.down_proj``.
+    Without this the fixture would push the exporter toward a layout no real checkpoint uses.
+    """
+    dir_path = Path(dir_path)
+    shards = sorted(dir_path.glob("*.safetensors"))
+    state_dict: dict[str, torch.Tensor] = {}
+    for shard in shards:
+        state_dict.update(load_file(str(shard)))
+
+    per_expert = re.compile(r"^(.*\.mlp\.experts)\.(\d+)\.(gate_proj|up_proj|down_proj)\.weight$")
+    grouped: dict[str, dict[int, dict[str, torch.Tensor]]] = defaultdict(lambda: defaultdict(dict))
+    for key in list(state_dict):
+        match = per_expert.match(key)
+        if match:
+            base, expert_id, proj = match.groups()
+            grouped[base][int(expert_id)][proj] = state_dict.pop(key)
+    if not grouped:
+        return
+
+    for base, experts in grouped.items():
+        ids = sorted(experts)
+        # gate first, then up, along the output dim -- the order the packed tensor is split on.
+        state_dict[f"{base}.gate_up_proj"] = torch.stack(
+            [torch.cat([experts[i]["gate_proj"], experts[i]["up_proj"]], dim=0) for i in ids]
+        )
+        state_dict[f"{base}.down_proj"] = torch.stack([experts[i]["down_proj"] for i in ids])
+
+    for shard in shards:
+        shard.unlink()
+    save_file(state_dict, str(dir_path / "model.safetensors"), metadata={"format": "pt"})
+    index = dir_path / "model.safetensors.index.json"
+    if index.exists():
+        index.unlink()
+
+
 def _create_tiny_qwen3_5_vl_dir(
     tmp_path: Path | str,
     with_processor: bool = False,
@@ -398,7 +440,7 @@ def _create_tiny_qwen3_5_vl_dir(
     moe: bool = False,
     **config_kwargs,
 ) -> Path | tuple[Path, PreTrainedModel]:
-    return _create_tiny_vlm_dir(
+    result = _create_tiny_vlm_dir(
         Path(tmp_path) / ("tiny_qwen3_5_moe_vl" if moe else "tiny_qwen3_5_vl"),
         QWEN3_5_VL_REF,
         _get_tiny_qwen3_5_vl,
@@ -407,6 +449,9 @@ def _create_tiny_qwen3_5_vl_dir(
         moe=moe,
         **config_kwargs,
     )
+    if moe:
+        _pack_qwen3_5_moe_experts(result[0] if isinstance(result, tuple) else result)
+    return result
 
 
 get_tiny_qwen3_5_vl = partial(_get_tiny_qwen3_5_vl, moe=False)
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 11421b53e80..197f5923dac 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
@@ -22,13 +22,13 @@
 @pytest.mark.parametrize(
     ("arch", "grouped_is_exportable"),
     [
-        # Only Nemotron-H maps fused grouped-GEMM experts, so it keeps the faster layout.
+        # These map fused grouped-GEMM experts, so they keep the faster layout.
         ("NemotronHForCausalLM", True),
+        ("Qwen3_5MoeForConditionalGeneration", True),
         ("Qwen3MoeForCausalLM", False),
         ("DeepseekV3ForCausalLM", False),
         ("GptOssForCausalLM", False),
         ("Llama4ForConditionalGeneration", False),
-        ("Qwen3_5MoeForConditionalGeneration", False),
     ],
 )
 def test_grouped_gemm_exportability(arch, grouped_is_exportable):
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 8b0b77a1f36..6b5620206ff 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -22,11 +22,13 @@
 import torch
 import transformers
 from _test_utils.torch.export.unified_checkpoint import assert_exported_checkpoint_matches
-from _test_utils.torch.megatron.models import get_mcore_gpt_model
+from _test_utils.torch.megatron.models import get_mcore_gpt_model, get_mcore_hybrid_model
 from _test_utils.torch.megatron.utils import get_forward
 from _test_utils.torch.transformers_models import (
     create_tiny_llama_dir,
     create_tiny_nemotron_dir,
+    create_tiny_nemotron_h_dir,
+    create_tiny_qwen3_5_moe_vl_dir,
     create_tiny_qwen3_moe_dir,
     create_tiny_qwen3vl_dir,
 )
@@ -89,7 +91,34 @@ def _test_unified_export_megatron(
     size,
     model_dir=None,
 ):
-    if model_type == "qwen3vl":
+    if model_type == "nemotron_h":
+        config = transformers.AutoConfig.from_pretrained(model_dir)
+        model = get_mcore_hybrid_model(
+            tensor_model_parallel_size=size,
+            pipeline_model_parallel_size=1,
+            initialize_megatron=True,
+            num_layers=config.num_hidden_layers,
+            hybrid_layer_pattern=config.hybrid_override_pattern,
+            hidden_size=config.hidden_size,
+            num_attention_heads=config.num_attention_heads,
+            num_query_groups=config.num_key_value_heads,
+            ffn_hidden_size=config.intermediate_size,
+            max_sequence_length=config.max_position_embeddings,
+            vocab_size=config.vocab_size,
+            mamba_state_dim=config.ssm_state_size,
+            mamba_num_heads=config.mamba_num_heads,
+            mamba_head_dim=config.mamba_head_dim,
+            mamba_num_groups=config.n_groups,
+            num_moe_experts=config.n_routed_experts,
+            moe_ffn_hidden_size=config.moe_intermediate_size,
+            moe_shared_expert_intermediate_size=config.moe_shared_expert_intermediate_size,
+            # NemotronH is the only arch that exports fused grouped-GEMM experts.
+            moe_grouped_gemm=True,
+            # NemotronH norms are RMSNorm; the builder defaults to LayerNorm, whose biases have no
+            # counterpart in the HF checkpoint.
+            normalization="RMSNorm",
+        ).cuda()
+    elif model_type == "qwen3vl":
         config = transformers.AutoConfig.from_pretrained(model_dir)
         text_cfg = config.text_config
         num_layers = text_cfg.num_hidden_layers
@@ -100,6 +129,30 @@ def _test_unified_export_megatron(
         max_sequence_length = text_cfg.max_position_embeddings
         vocab_size = text_cfg.vocab_size
         extra_kwargs = {"kv_channels": text_cfg.head_dim, "qk_layernorm": True}
+    elif model_type in {"qwen3_5_moe_vl_grouped", "qwen3_5_moe_vl_sequential"}:
+        text_cfg = transformers.AutoConfig.from_pretrained(model_dir).text_config
+        num_layers = text_cfg.num_hidden_layers
+        hidden_size = text_cfg.hidden_size
+        num_attention_heads = text_cfg.num_attention_heads
+        num_query_groups = text_cfg.num_key_value_heads
+        ffn_hidden_size = text_cfg.intermediate_size
+        max_sequence_length = text_cfg.max_position_embeddings
+        vocab_size = text_cfg.vocab_size
+        # Hybrid GatedDeltaNet + gated attention, with routed experts stored packed.
+        extra_kwargs = {
+            "kv_channels": text_cfg.head_dim,
+            "qk_layernorm": True,
+            "experimental_attention_variant": "gated_delta_net",
+            "num_moe_experts": text_cfg.num_experts,
+            "moe_ffn_hidden_size": text_cfg.moe_intermediate_size,
+            "moe_shared_expert_intermediate_size": text_cfg.shared_expert_intermediate_size,
+            "moe_shared_expert_gate": True,
+            # Match the HF layer_types pattern (every Nth layer is full attention, rest GDN).
+            "linear_attention_freq": len(text_cfg.layer_types),
+            # Both layouts must reach the same packed HF tensors, via GroupedMLPPacking
+            # (TEGroupedMLP) and PackNameRemapping (SequentialMLP) respectively.
+            "moe_grouped_gemm": model_type.endswith("grouped"),
+        }
     elif model_type == "qwen3_moe":
         config = transformers.AutoConfig.from_pretrained(model_dir)
         num_layers = config.num_hidden_layers
@@ -133,22 +186,26 @@ def _test_unified_export_megatron(
     activation_func = "squared_relu" if model_type == "nemotron" else "swiglu"
     normalization = "LayerNorm" if model_type == "nemotron" else "RMSNorm"
 
-    model = get_mcore_gpt_model(
-        tensor_model_parallel_size=size,
-        pipeline_model_parallel_size=1,
-        initialize_megatron=True,
-        num_layers=num_layers,
-        hidden_size=hidden_size,
-        num_attention_heads=num_attention_heads,
-        num_query_groups=num_query_groups,
-        ffn_hidden_size=ffn_hidden_size,
-        max_sequence_length=max_sequence_length,
-        vocab_size=vocab_size,
-        activation_func=activation_func,
-        normalization=normalization,
-        transformer_impl="modelopt",
-        **extra_kwargs,
-    ).cuda()
+    model = (
+        model
+        if model_type == "nemotron_h"
+        else get_mcore_gpt_model(
+            tensor_model_parallel_size=size,
+            pipeline_model_parallel_size=1,
+            initialize_megatron=True,
+            num_layers=num_layers,
+            hidden_size=hidden_size,
+            num_attention_heads=num_attention_heads,
+            num_query_groups=num_query_groups,
+            ffn_hidden_size=ffn_hidden_size,
+            max_sequence_length=max_sequence_length,
+            vocab_size=vocab_size,
+            activation_func=activation_func,
+            normalization=normalization,
+            transformer_impl="modelopt",
+            **extra_kwargs,
+        ).cuda()
+    )
 
     if quant_config:
         quant_config_dict = getattr(mtq, quant_config)
@@ -210,8 +267,10 @@ def _test_unified_export_megatron(
     ("model_type", "extra_module", "quant_config", "kv_cache_quant_cfg"),
     [
         ("nemotron", None, None, None),
-        ("nemotron", None, "NVFP4_DEFAULT_CFG", None),
-        ("nemotron", None, "NVFP4_DEFAULT_CFG", "FP8_KV_CFG"),
+        # NemotronH (Mamba + attention + grouped-GEMM MoE) supersedes the older Nemotron for
+        # quantized coverage; the plain/eagle/medusa cases below still smoke-test the old arch.
+        ("nemotron_h", None, "NVFP4_DEFAULT_CFG", None),
+        ("nemotron_h", None, "NVFP4_DEFAULT_CFG", "FP8_KV_CFG"),
         ("nemotron", "eagle", None, None),
         ("nemotron", "medusa", None, None),
         ("llama", None, None, None),
@@ -224,6 +283,11 @@ def _test_unified_export_megatron(
         # Regression guard: routed experts used to be dropped silently from the export.
         ("qwen3_moe", None, None, None),
         ("qwen3_moe", None, "FP8_DEFAULT_CFG", None),
+        # Packed routed experts (Qwen3.5). NVFP4 keeps per-expert block scales while FP8 merges a
+        # single scale, so the packing rules only get full coverage across both formats.
+        ("qwen3_5_moe_vl_grouped", None, "NVFP4_DEFAULT_CFG", None),
+        ("qwen3_5_moe_vl_grouped", None, "FP8_DEFAULT_CFG", None),
+        ("qwen3_5_moe_vl_sequential", None, "NVFP4_DEFAULT_CFG", None),
     ],
 )
 def test_unified_export_megatron(
@@ -235,8 +299,12 @@ def test_unified_export_megatron(
         model_dir = create_tiny_qwen3vl_dir(tmp_path)
     elif model_type == "nemotron":
         model_dir = create_tiny_nemotron_dir(tmp_path)
+    elif model_type == "nemotron_h":
+        model_dir = create_tiny_nemotron_h_dir(tmp_path)
     elif model_type == "qwen3_moe":
         model_dir = create_tiny_qwen3_moe_dir(tmp_path)
+    elif model_type.startswith("qwen3_5_moe_vl"):
+        model_dir = create_tiny_qwen3_5_moe_vl_dir(tmp_path)
     else:
         raise ValueError(f"Unsupported model_type: {model_type}")
     # TODO: Fix TP>1 failures

From 215d4aaef8cb953eee0003d56c85be8b5ba13da3 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Sat, 29 Aug 2026 12:05:46 -0700
Subject: [PATCH 12/12] Address review: quant-config recording, GDN guards, VLM
 probe, teacher layout

`_grouped_mlp_packing` reused `_grouped_mlp_slicing` through an internal `\x00pack\x00`
marker prefix, and the slicer records per-layer quant metadata itself. That marker
contains no `{`, so `_record_layer_quant_config`'s placeholder guard did not skip it and
wrote one NUL-byte key per expert per layer, while the real packed prefix was never
recorded at all. Homogeneous exports pop `quantized_layers`, which is why the new tests
stayed green; a mixed-precision export would emit the junk keys and report the packed
routed experts as unquantized. The slicer no longer records when packing, and packing
records against the real prefix as `_pack_name_remapping` does.

The zero-centered-gamma guard read `module.config`, which TE norms do not have, so it
resolved to the `True` default and could never fire. It now reads the module's own
`zero_centered_gamma`, which those norms do carry. Only the GatedDeltaNet output norm is
zero-centered in practice -- exported `input_layernorm` and final `norm` are bit-exact
against the reference -- so this guard fires exactly when adding 1.0 would be wrong.

`_gated_delta_net_slicing` looked section sizes up by name but split on physical row
order; that assumption is now asserted, turning an upstream reorder into a message
instead of silently mis-sliced weights.

The VLM checkpoint probe keyed off `vision_model.` while `get_language_model` keys off
`.language_model`, so a renamed vision tower would misclassify a full-VLM checkpoint as
language-model-only and load it into `.language_model`. Both now use the same signal.

distill.py forced the student's expert layout onto the teacher. The teacher only runs
forward, is loaded from HF, and is hidden from the checkpoint by
`expose_minimal_state_dict`, so it keeps grouped GEMM independently.

Also documents that `--no_moe_grouped_gemm` is a no-op on the LLM path of
export_distilled_megatron_to_hf.py, and raises the megatron_bridge example-test CI
timeout to 75 minutes.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/workflows/example_tests.yml           |  2 +-
 examples/megatron_bridge/distill.py           | 14 +++++--
 .../export_distilled_megatron_to_hf.py        |  3 +-
 .../torch/export/unified_export_megatron.py   | 39 ++++++++++++++++---
 modelopt/torch/utils/plugins/mbridge.py       |  8 ++--
 5 files changed, 51 insertions(+), 15 deletions(-)

diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml
index 9a23e8181b0..14cf7247e04 100644
--- a/.github/workflows/example_tests.yml
+++ b/.github/workflows/example_tests.yml
@@ -155,7 +155,7 @@ jobs:
     with:
       docker_image: "nvcr.io/nvidia/nemo:26.08"
       example: megatron_bridge
-      timeout_minutes: 60
+      timeout_minutes: 75
       pip_install_extras: "[hf,puzzletron,dev-test]"
       runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
       allow_failure: ${{ contains(format(',{0},', vars.ALLOW_FAILURE_EXAMPLE_TESTS), ',megatron_bridge,') }}
diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py
index 7df5489159a..25bd9242818 100644
--- a/examples/megatron_bridge/distill.py
+++ b/examples/megatron_bridge/distill.py
@@ -360,7 +360,7 @@ def main(args: argparse.Namespace):
     tensorboard_dir = os.path.join(args.output_dir, "tb_logs")
 
     # Build student and teacher model providers
-    def _build_model_provider(hf_path, load_weights=True):
+    def _build_model_provider(hf_path, load_weights=True, moe_grouped_gemm=True):
         bridge = AutoBridge.from_hf_pretrained(hf_path, trust_remote_code=args.trust_remote_code)
         provider = bridge.to_megatron_provider(load_weights=load_weights)
 
@@ -373,7 +373,6 @@ def _build_model_provider(hf_path, load_weights=True):
         provider.expert_model_parallel_size = args.ep_size
         provider.expert_tensor_parallel_size = 1  # Expert tensor parallelism is not supported
         provider.seq_length = args.seq_length
-        # Must match the expert layout of --student_megatron_path (see quantize.py).
         set_moe_expert_layout(provider, moe_grouped_gemm)
         if args.sft:
             # A response-only loss mask needs per-token reduction to combine across CP ranks.
@@ -390,14 +389,21 @@ def _build_model_provider(hf_path, load_weights=True):
     # The student structure is always built from --student_hf_path. When --student_megatron_path is
     # given, the HF weights are skipped (they are overwritten by the Megatron checkpoint, loaded into
     # the built student inside the patched provide() below).
+    # Only the student's layout is pinned -- it must match --student_megatron_path (see quantize.py).
     student_provider = _build_model_provider(
-        args.student_hf_path, load_weights=args.student_megatron_path is None
+        args.student_hf_path,
+        load_weights=args.student_megatron_path is None,
+        moe_grouped_gemm=moe_grouped_gemm,
     )
     if student_has_modelopt_state:
         # Gradient accumulation fusion is not supported with ModelOpt quantized models. Disable it
         # before the model is built so the student's linear layers are constructed accordingly.
         student_provider.gradient_accumulation_fusion = False
-    teacher_provider = _build_model_provider(args.teacher_hf_path)
+    # The teacher only runs forward, is loaded from HF, and is hidden from the checkpoint
+    # (``expose_minimal_state_dict``), so it keeps the faster grouped GEMM regardless.
+    teacher_provider = _build_model_provider(
+        args.teacher_hf_path, moe_grouped_gemm=not args.no_moe_grouped_gemm
+    )
 
     # The KD losses compare logits elementwise over the vocab dim, so both output layers must have
     # the same padded width. A shared tokenizer does not imply it: the HF configs can disagree.
diff --git a/examples/megatron_bridge/export_distilled_megatron_to_hf.py b/examples/megatron_bridge/export_distilled_megatron_to_hf.py
index fbb2f0c8760..5957f49b389 100644
--- a/examples/megatron_bridge/export_distilled_megatron_to_hf.py
+++ b/examples/megatron_bridge/export_distilled_megatron_to_hf.py
@@ -222,7 +222,8 @@ def get_args() -> argparse.Namespace:
         help=(
             "Force SequentialMLP for MoE experts instead of the fused TEGroupedMLP (grouped GEMM). "
             "By default grouped GEMM is used unless the architecture cannot export it to "
-            "HuggingFace, in which case SequentialMLP is selected automatically."
+            "HuggingFace, in which case SequentialMLP is selected automatically. VLMs only: the "
+            "LLM path reads the expert layout from the checkpoint and ignores this flag."
         ),
     )
     parser.add_argument("--tp_size", type=int, default=1, help="Tensor parallel size")
diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py
index 281271a210e..070ad8d34e4 100644
--- a/modelopt/torch/export/unified_export_megatron.py
+++ b/modelopt/torch/export/unified_export_megatron.py
@@ -1137,9 +1137,8 @@ def _name_remapping(
         if zero_centered_gamma:
             # Megatron centres this gamma on 0, HF on 1. Assert, don't derive: the config flag
             # is model-wide while this one is per-norm.
-            module_config = getattr(module, "config", None)
-            assert getattr(module_config, "layernorm_zero_centered_gamma", True), (
-                f"{prefix} is mapped as zero-centered gamma but its config disables it; "
+            assert getattr(module, "zero_centered_gamma", False), (
+                f"{prefix} is mapped as zero-centered gamma but the module reports otherwise; "
                 "exporting would shift the weights by 1.0"
             )
             weight = weight + 1.0
@@ -1246,7 +1245,12 @@ def _grouped_mlp_packing(self, module, prefix, parallel_config=None, is_mtp=Fals
         self._state_dict = OrderedDict()
         try:
             qformat, block_size = self._grouped_mlp_slicing(
-                module, marker + "{}", parallel_config=parallel_config, is_mtp=False, quantize=False
+                module,
+                marker + "{}",
+                parallel_config=parallel_config,
+                is_mtp=False,
+                quantize=False,
+                record_quant_config=False,
             )
             per_expert = self._state_dict
         finally:
@@ -1263,6 +1267,12 @@ def collect(suffix):
         weights = collect(".weight")
         if not weights:
             return
+        # Record against the packed prefix, as _pack_name_remapping does for the other packed path.
+        if qformat in (None, QUANTIZATION_NONE):
+            self._record_excluded_module(prefix)
+        else:
+            assert block_size is not None
+            self._record_layer_quant_config(prefix, qformat, block_size)
         scales, scales_2 = collect(".weight_scale"), collect(".weight_scale_2")
         input_scales = collect(".input_scale")
 
@@ -1290,7 +1300,13 @@ def collect(suffix):
             )[0]
 
     def _grouped_mlp_slicing(
-        self, module, prefix, parallel_config=None, is_mtp=False, quantize=True
+        self,
+        module,
+        prefix,
+        parallel_config=None,
+        is_mtp=False,
+        quantize=True,
+        record_quant_config=True,
     ):
         """Export TEGroupedMLP weight0..weight{N-1} as one HF-style entry per expert.
 
@@ -1447,7 +1463,7 @@ def _grouped_mlp_slicing(
         # Record quant config for ALL global experts on every rank; otherwise the writer's
         # hf_quant_config.json would miss (EP-1)/EP of the routed experts. All experts in
         # a TEGroupedMLP layer share qformat/block_size, so local values apply globally.
-        if seen_qformat is not None:
+        if seen_qformat is not None and record_quant_config:
             assert seen_block_size is not None
             num_total_experts = num_experts * ep_size
             for global_id in range(num_total_experts):
@@ -1651,6 +1667,17 @@ def _gated_delta_net_slicing(self, module, prefix, is_mtp=False):
             in_proj, self.dtype, prefix=prefix
         )
 
+        assert tuple(module.in_proj_split_names) == (
+            "query",
+            "key",
+            "value",
+            "z",
+            "beta",
+            "alpha",
+        ), (
+            f"Unexpected GatedDeltaNet in_proj layout {tuple(module.in_proj_split_names)}; the "
+            "split below assumes [query, key, value, z, beta, alpha]"
+        )
         sections = dict(zip(module.in_proj_split_names, module.in_proj_split_sections))
         split_sizes = [
             sections["query"] + sections["key"] + sections["value"],
diff --git a/modelopt/torch/utils/plugins/mbridge.py b/modelopt/torch/utils/plugins/mbridge.py
index 131f8f5220b..23fdaa8b324 100644
--- a/modelopt/torch/utils/plugins/mbridge.py
+++ b/modelopt/torch/utils/plugins/mbridge.py
@@ -218,12 +218,14 @@ def load_modelopt_megatron_checkpoint(
     # _load_model_weights_from_checkpoint does not resolve the latest iter_* directory, so resolve it explicitly
     checkpoint_path = _get_modelopt_checkpoint_path(megatron_path)
 
-    # ``distill.py`` distills a VLM's language model only, so its checkpoint holds no ``vision_model.*``
-    # weights and must be loaded into ``.language_model`` rather than the full VLM wrapper.
+    # ``distill.py`` distills a VLM's language model only, so its checkpoint holds the language
+    # model at the root and must be loaded into ``.language_model``, not the full VLM wrapper.
+    # Keyed off the same ``language_model.`` prefix ``get_language_model`` navigates, so a renamed
+    # vision tower cannot make a full-VLM checkpoint look language-model-only.
     checkpoint_keys = _checkpoint_keys(checkpoint_path)
     unwrapped_model = unwrap_model(model)
     if any(get_language_model(m)[1] for m in unwrapped_model) and not any(
-        key.startswith("vision_model.") for key in checkpoint_keys
+        key.startswith("language_model.") for key in checkpoint_keys
     ):
         print_rank_0("Language-model-only checkpoint: loading into the VLM's `.language_model`.")
         model = [get_language_model(m)[0] for m in unwrapped_model]