Support quantized Qwen3-VL / Qwen3.5-VL (dense + MoE) export from Megatron-Bridge and verify exported checkpoints (#2276)

### What does this PR do?

Type of change: Bug fix + new feature

Enables quantized **Qwen3-VL** and **Qwen3.5-VL** (dense and MoE) →
unified HuggingFace export from Megatron-Bridge, and fixes the bugs
found along the way (ten from testing, plus a further round from
review). Most of them produced a valid-looking checkpoint and a green
test run, so the PR also makes the export path verify its own output.

Review is easiest commit-by-commit — each of the eleven commits is
self-contained and independently green.

#### Two blockers

1. **The exporter rejected the Megatron-Bridge VLM wrapper.**
`GPTModelExporter` only unwrapped MCore's `LLaVAModel`, so
`Qwen3VLModel` raised `ValueError: Input to GPTModelExport must be a
megatron.core.models.GPTModel!`. It now unwraps any wrapper exposing
`.language_model`.
2. **A VLM QAD checkpoint couldn't be loaded back.** `distill.py` passes
`distill_submodule="language_model"`, so the checkpoint holds only the
language model and the load died on `KeyError:
vision_model.patch_embed.proj.weight`. The loader now reads the
checkpoint metadata and targets `.language_model` when there are no
vision weights.

#### Four silent-corruption bugs

3. **VLM QAD discarded all ModelOpt state** (shipped in 0.46).
`ModeloptStateManager` requires state on the **root** of whatever gets
checkpointed. `quantize.py` quantizes the VLM root, so PTQ anchors it
there — but QAD checkpoints only `language_model`, orphaning it. The
saved `modelopt_state_dict` was literally `[]`; the `*_quantizer._amax`
tensors were still present but got dropped on load
(`dist_ckpt_strictness="assume_ok_unexpected"`), and the export came out
plain BF16 with no `hf_quant_config.json`.
4. **Fused grouped-GEMM MoE experts were omitted entirely.** The MoE
dispatch had no `else`, so an architecture without an
`experts.linear_fc1` rule exported *zero routed experts*. This hit
**`Qwen3MoeForCausalLM`** — a registered, supported architecture with no
export test — not just VLMs. A tiny Qwen3-MoE exported 37 of 45 tensors,
exit 0, no warning.
5. **Qwen3.5's GatedDeltaNet output norm was off by exactly 1.0.**
Megatron stores that gamma zero-centered, HF centers it on 1. Correct
names, correct shapes, wrong values — invisible to any structural check.
Megatron-Bridge's importer confirms the convention
(`RMSNorm2ZeroCenteredRMSNormMapping`).
6. **The disabled-quantizer patterns silently no-op on Megatron paths.**
They are written against HuggingFace module names. `*mixer.conv1d*`
matches only because MCore and HF happen to agree on "mixer" for Mamba;
`*linear_attn.conv1d*` never matched (Megatron calls it
`self_attention.conv1d`), so the conv1d was calibrated.
`*linear_attn.in_proj_a/b*` **cannot** match at all — Megatron fuses all
six GDN sections behind one quantizer — so the alpha/beta gates the
recipe wants in BF16 were exported in FP8.

#### Four more bugs, found only by running real checkpoints

The tiny fixtures could not reach these; each came from a real model or
a real quant format.

7. **Routed experts were written in a layout no real Qwen3.5 checkpoint
uses.** Real Qwen3.5 stores experts packed as `[num_experts, out, in]`;
the mapping emitted per-expert names, so every routed expert was
dropped. The fixture actively hid this: transformers *unpacks* experts
on `save_pretrained`, so the saved reference agreed with the wrong
output. Fixed with a `transpose` kwarg on `_pack_name_remapping` plus a
`GroupedMLPPacking` rule, so fused `TEGroupedMLP` reaches the same
packed tensors — which is also what lets Qwen3.5 keep grouped GEMM
(**22.1 GB/GPU vs 38.9 GB/GPU** on a 20-layer, 256-expert model).
8. **`_grouped_mlp_packing` was broken for NVFP4.** It 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`.
`weight_scale_2` is never populated in an FP8 run, so the whole branch
was dead code under FP8-only testing. `_grouped_mlp_slicing` gained
`quantize=False` so packing can quantize once over the stack, matching
`_pack_name_remapping`.
9. **`_mtp_prefix` corrupted every VLM's MTP tensor names.** It did
`prefix.replace("model", "mtp")` uncounted, so
`model.language_model.layers.{}` became `mtp.language_mtp.layers.0.*` —
tensors present and correctly valued, under names nothing loads.
LLM-only prefixes contain one occurrence, so this was invisible until a
VLM with MTP was exported.
10. **`load_multimodal_components` rejected HF repo ids.** `quantize.py
--hf_model_name_or_path Qwen/Qwen3.5-0.8B` worked, but the documented
export step failed with *"It should be a directory"*. Its sibling in the
same file already resolved repo ids via `snapshot_download`; now it does
too. This affected **every** VLM export.

`Qwen3_5ForConditionalGeneration` (dense Qwen3.5-VL) is now registered
for export and vision passthrough, which bugs 9 and 10 were blocking.

#### New: Qwen3.5-VL

`GatedDeltaNetSlicing` splits Megatron's fused `in_proj` (`[query, key,
value, z, beta, alpha]`) into HF's `in_proj_qkv` / `_z` / `_b` / `_a`,
taking sizes from the module's own `in_proj_split_sections` so TP
sharding falls out. Widening coverage to Qwen3.5's *gated
full-attention* layers then exposed a further split bug: gated attention
packs a per-head output gate beside each query head, so `_qkv_slicing`
split 192 rows as 96/48/48 instead of 128/32/32. It now derives the
group stride from `config.attention_output_gate`, matching
Megatron-Bridge's `split_qkv_weights`. The non-gated path is unchanged.

#### New: the export path verifies itself

- `assert_exported_checkpoint_matches` compares an exported checkpoint
against the model it came from — key set, shapes (accounting for NVFP4
`uint8` packing), safetensors index consistency, and values — replacing
existence-only assertions in all three export tests.
- `GPTModelExporter.save_pretrained` now raises if the export dropped
tensors the source checkpoint has, so *user* runs on architectures CI
never sees are protected too, not just tiny models.
- Loading a checkpoint whose quantizer tensors have no restorable state
now raises instead of silently loading unquantized.
- `assert_has_modelopt_state` replaces `rglob("modelopt_state")`, which
passes on an empty state; `assert_no_quantizers_matching` fails on
future HF↔Megatron name drift.

The mapping is also table-driven now: vision-tower prefixes live in
`all_mcore_hf_vision_passthrough_mapping` and
`with_language_model_prefix` is shared, so adding a VLM no longer means
editing `unified_export_megatron.py`. Five call sites that answered "is
this a VLM" three different ways now share `get_language_model` /
`is_vlm_config`.

### Usage

```bash
# Dense VLM (Qwen3-VL) -- no extra flags
torchrun --nproc_per_node 2 quantize.py \
    --hf_model_name_or_path Qwen/Qwen3-VL-8B-Instruct \
    --quant_cfg nvfp4 --tp_size 2 \
    --export_megatron_path /tmp/Qwen3-VL-8B-NVFP4-megatron

torchrun --nproc_per_node 2 export_quantized_megatron_to_hf.py \
    --hf_model_name_or_path Qwen/Qwen3-VL-8B-Instruct \
    --megatron_path /tmp/Qwen3-VL-8B-NVFP4-megatron \
    --pp_size 2 --export_unified_hf_path /tmp/Qwen3-VL-8B-NVFP4-hf

# Gated MoE (Qwen3.5-VL, Qwen3-MoE) -- no extra flags either. The scripts derive the
# expert layout from the model config, so quantize / distill / export all agree.
# --no_moe_grouped_gemm forces SequentialMLP if you want it explicitly.
```

### Testing

All in `nvcr.io/nvidia/nemo:26.08` on 2x RTX 6000 Ada.

| Suite | Result | Time |
|---|---|---|
| `tests/examples/megatron_bridge/` (full) | 18 passed | 27m58 |
| `tests/gpu_megatron/torch/export/` | 38 passed | 2m13 |
| `tests/unit/torch/export/` | 186 passed | 1.5s |
| pre-commit (ruff, ruff format, mypy, bandit) | clean | — |
| `tests/examples/megatron_bridge/test_quantize_export.py` on **2 GPUs**
(`pp_size=2`) | 3 passed | 5m |

The export leg of `test_quantize_and_export` now scales with `num_gpus`
like its quantize leg
already did. Previously it was hardcoded to one process, so the
collective checkpoint load ran at
PP=1 on both the 1-GPU PR runner and the 2-GPU nightly — which is how a
guard that raised on only
some pipeline stages (and therefore hung the job) reached review. The
dense `qwen3` case was dropped
in exchange: `qwen3_moe` already covers the non-VLM script path,
`qwen3vl` covers a dense decoder,
and that case was the one exceeding the 300s cap in CI.

#### Model coverage

`tests/gpu_megatron` runs in-process and is cheap, so it owns
per-architecture **mapping**
correctness. The example tests spawn `torchrun` per step and are ~50x
slower per case, so they
cover **script wiring** only — CLI flags, recipe resolution, and
checkpoint hand-off between steps.

| Suite | Models |
|---|---|
| `test_unified_export_megatron` | llama, nemotron, nemotron_h, qwen3vl,
qwen3_moe, qwen3_5_moe_vl x {none, FP8, NVFP4, +/-KV} x {grouped GEMM,
SequentialMLP} + eagle / medusa / MTP (29 params) |
| `test_megatron_importer` | nemotron_h, llama export->import round-trip
|
| `test_moe_layout_choice` | per-architecture grouped-GEMM exportability
(6 architectures) |
| `test_distill_megatron` | KD loss mechanics |

| Model | prune | quantize+export | QAD | distill+export |
|---|:--:|:--:|:--:|:--:|
| qwen3 | Y | Y | Y | Y |
| qwen3_moe | - | **Y (new)** | - | - |
| qwen3vl | - | **Y (moved from QAD)** | - | - |
| nemotron_h | Y | **Y (new)** | - | - |
| qwen3_5_vl | - | - | - | Y |
| qwen3_5_moe_vl | Y | **Y (new, both expert layouts)** | Y | - |
| deepseek_v3 | Y | - | - | - |
| gemma3vl | Y | - | ~~manual~~ removed | - |

QAD's unique property is that ModelOpt state survives distillation,
which needs one LLM and one
VLM rather than one case per architecture. Moving the rest to
quantize+export drops a `torchrun`
launch each: QAD went from 3 CI cases to 2 while quantize+export went
from 1 to 4, adding two
architectures for about a minute.

#### Real-model validation

Tiny fixtures cannot catch layout or scale bugs that only appear at real
dimensions, so the export
path was run end-to-end on released checkpoints. This is where bugs 7-10
came from.

| Model | Run | Result |
|---|---|---|
| Nemotron-3.5-Lightning-30B-A3B | NVFP4 4o6 PTQ → export → MMLU |
**0.7825 ± 0.0105** (gate 0.75) |
| Nemotron-3.5-Lightning-30B-A3B | Minitron pruning | 22.28B/3.00B
active, **0.5944** (gate 0.58) |
| Qwen3.5-0.8B (dense VLM) | FP8 PTQ → export → MMLU | BF16 0.4895 →
**0.4832** (±0.0127) |
| Qwen3.5-35B-A3B, half-depth (20 layers, 256 experts) | FP8 + NVFP4 PTQ
→ export | keys + shapes + **values** match reference |
| Qwen3.5-35B-A3B, full | FP8 PTQ | OOM on 2x48GB (see below) |

The half-depth model keeps real weights, real dims and all 256 experts.
Both expert layouts produce
identical key sets, and all exports pass
`assert_exported_checkpoint_matches(..., check_values=True)`
— every tensor, including all 20 x 256 experts, dequantizes to within
tolerance of the BF16
reference, so a transposed or mis-ordered expert stack would fail. NVFP4
lands in the correct packed
layout (`gate_up_proj [256, 1024, 1024]` U8, `weight_scale [256, 1024,
128]` E4M3,
`weight_scale_2 []` F32). Its *accuracy* is not meaningful — truncating
to 20 of 40 layers leaves a
chance-level model (BF16 0.2322, FP8 0.2538) — so it validates
correctness, not quality.

**Re-validated on the final code.** The numbers above were first taken
mid-review; since then the
NVFP4 block-scale merge changed on both packed paths, the vision-tower
download became two-stage,
and an expert-layout load guard was added. Both gating runs were
therefore repeated end to end:
Nemotron went 0.7748 → **0.7825 ± 0.0105** and Qwen3.5-0.8B went 0.4678
→ **0.4832 ± 0.0127**, with
the rest of the Nemotron pipeline reproducing exactly (3519 quantizers,
69GB checkpoint, 21GB
export). Both deltas are inside their own stderr, so the claim is that
the rework costs no accuracy
— not that it improved it. The Nemotron export also runs at `--pp_size
2`, exercising the new
collective layout guard on a real 30B MoE across pipeline stages.

Two limitations worth stating plainly:

- **No quantized accuracy number for a full-size MoE.** The full 35B
OOMs at 47.37 GiB while
*constructing* the model on 2x48GB, with grouped GEMM already enabled,
so no calibration knob
  helps. Needs more GPUs than this setup has.
- **vLLM cannot yet serve packed FP8 Qwen3.5 experts.** `vllm
0.24.1.dev0` builds its fused expert
mapping weight-only, rewriting `experts.down_proj_input_scale` to
`w2_weight_input_scale` while the
parameter it registers is `w2_input_scale`. This is upstream and
independent of how the checkpoint
is produced — both of our export paths fail it identically. The 0.8B
numbers above are unaffected
(dense), and the packed exports are verified against the reference
checkpoint instead.

#### Guard verification

Each new guard was made to fire, not just to compile:

| Guard | Verification |
|---|---|
| Export self-check | Disabled the MoE guard, re-exported Qwen3-MoE -
independently reported all 24 dropped tensors. No false positives across
llama, nemotron, qwen3, qwen3-moe, qwen3vl, qwen3.5-vl, deepseek_v3
incl. eagle / medusa / MTP |
| Dropped-state raise | Deleted `modelopt_state` from a checkpoint with
50 quantizer tensors - raised instead of loading unquantized |
| NVFP4 value check | Flipped a `q_proj` - failed at `max_rel_err=1.74`
against a 0.3 threshold |
| Zero-centered gamma | Reproduced the off-by-1.0 on a good export -
caught as "not bit-exact" |
| Exclusion guard | Asserts no calibrated quantizer matches `conv1d` /
`mlp.router` / `output_layer` |

Exported artifacts are validated, not just their existence: 0 missing
keys vs reference, vision
tower bitwise-identical, dequantized weights within FP8 E4M3 error
(<=4.6%). The
`in_proj_a`/`in_proj_b` check is load-bearing - swapped alpha/beta would
still match on shape but
show ~100% error.

Also ran a tiny-Qwen3 **LLM** control through both steps to confirm the
exporter changes are a
no-op off the VLM path.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅ — the scripts now derive the
MoE expert layout from the model config, building SequentialMLP only for
architectures with no `experts.linear_fc1` rule, and the exporter raises
rather than dropping experts it has no rule for. Those runs previously
"succeeded" while writing a checkpoint containing no expert weights, so
no working behaviour is removed. `--no_moe_grouped_gemm` forces
SequentialMLP explicitly.
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅
- Did you get Claude approval on this PR?: ✅ — approved (round 10: 0
CRITICAL, 0 IMPORTANT, 0 new suggestions); CodeRabbit approved earlier

### Additional Information

**MoE expert layout is now chosen automatically.** Only Nemotron-H can
export fused grouped-GEMM experts, so every other MoE architecture would
otherwise need `--no_moe_grouped_gemm` on all four scripts or hit a wall
at export. The scripts derive the layout from the model config — grouped
GEMM unless it would not be exportable — so they agree without threading
a flag. This changes MoE activation scales from one shared scale to
per-expert for the affected architectures.

Known gaps, unchanged by this PR:

- **Gated MoE still cannot use fused grouped GEMM.**
`_grouped_mlp_slicing` emits one weight per expert with no gate/up split
— its only prior caller, Nemotron-H, is non-gated, so every other MoE
architecture is built as `SequentialMLP` (see below). Adding that split
would restore the faster layout, but it needs a deliberate call on
activation-scale semantics: grouped GEMM keeps **one shared** activation
scale across experts while `SequentialMLP` has **per-expert** scales, so
the two are not numerically equivalent. It also needs EP>1 coverage.
- **Qwen3.5's alpha/beta gates share Megatron's fused `in_proj`
quantizer,** so they can only be kept in BF16 at export, not excluded by
name. Full fidelity needs per-section quantizers on the fused
projection.
- **Anchoring ModelOpt state on `.language_model`** (which would let
`quantize.py` quantize the language model directly and drop its
name-based non-LM disabling) needs a coordinated Megatron-Bridge change:
`save_sharded_modelopt_state` is ModelOpt code, but the restore the
Bridge path uses is Bridge's own and unconditionally restores onto the
root.
- **Gemma3-VL** remains Megatron-checkpoint only (`OMNIML-5366`).

🤖 Generated with [Claude Code](https://claude.com/claude-code)


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
  * Added Muse Glimmer AutoQuantize and Alpamayo QAD workflows.
* Added streaming Kimi-K3 conversion and NVFP4 activation headroom
calibration.
  * Added SFT-masked distillation for Megatron-Bridge.
* Added unified Hugging Face export for quantized Qwen3-VL and
Qwen3.5-VL checkpoints.
* MoE expert layouts are selected automatically, with an option to force
sequential experts.

* **Bug Fixes**
* Improved export validation for tensor coverage, MoE mappings,
quantizer state, and NVFP4 scales.
  * Fixed Qwen3.5-VL GatedDeltaNet export handling.
  * Preserved visual-model weights exactly during export.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

.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,') }}

CHANGELOG.rst

@@ -19,8 +19,10 @@ Changelog
 *Megatron Framework (M-LM / M-Bridge)*
 
 - 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 per-expert weight quantization for Transformer Engine ``TEGroupedLinear`` (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``, 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*
 
@@ -40,7 +42,7 @@ Changelog
 - Remove the legacy ``layerwise`` bool form, its ``use_sequential`` alias, and the top-level ``layerwise_checkpoint_dir`` key from calibration algorithm configs (deprecated in 0.45). Use the nested form, e.g. ``layerwise: {enable: true, checkpoint_dir: /path}``. A pre-0.45 ``modelopt_state`` carrying either legacy key now fails validation on restore instead of being migrated; re-save it with a 0.45/0.46 release first.
 - Remove in-trainer quantization via ``QuantizationArguments.quant_cfg`` / ``--quant_cfg`` (deprecated in 0.45); use ``--recipe``. New recipes ``general/ptq/mxfp4_mlp_weight_only`` and ``general/ptq/nvfp4_mlp_weight_only`` replace ``MXFP4_MLP_WEIGHT_ONLY_CFG`` / ``NVFP4_MLP_WEIGHT_ONLY_CFG`` in the ``examples/gpt-oss`` QAT flow.
 - Remove the ``QuantizationArgumentsWithConfig`` alias in ``modelopt.torch.quantization.plugins.transformers_trainer`` (deprecated in 0.45). Use ``QuantizationArguments``.
-- Transformer Engine ``TEGroupedMLP`` (fused MoE experts) now uses **per-expert** weight quantization (one ``amax`` per expert) instead of a single shared ``amax``, so ModelOpt checkpoints containing quantized ``TEGroupedMLP`` modules saved before 0.47 are **not compatible** with 0.47. Re-run PTQ to regenerate compatible checkpoints.
+- Transformer Engine ``TEGroupedLinear`` (fused MoE experts) now uses **per-expert** weight quantization (one ``amax`` per expert) instead of a single shared ``amax``, so ModelOpt checkpoints containing quantized ``TEGroupedLinear`` modules saved before 0.47 are **not compatible** with 0.47. Re-run PTQ to regenerate compatible checkpoints.
 
 **Deprecations**
 
@@ -51,6 +53,11 @@ 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. 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.
 - Fix image-text calibration for non-Nemotron VLMs by forwarding multimodal batches through the complete VLM.
 - 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``.
@@ -90,7 +97,7 @@ Changelog
   - **GatedDeltaNet** (linear attention) and **gated attention** (``attention_output_gate``), such as Qwen3.5 (hybrid GatedDeltaNet + gated-attention) language models, including MoE variants — attention / linear-attention heads are not pruned.
   - **Multi-Latent Attention (MLA)**, such as DeepSeek — MLA latent ranks are not pruned.
   - **Latent MoE**, such as Nemotron-3-Super — ``hidden_size`` pruning resizes the latent projections while the MoE latent dim itself is not pruned.
-- Optimize Minitron pruning support for MoE models using the fused **grouped GEMM** experts (``TEGroupedMLP``) in addition to the existing ``SequentialMLP`` path. ``examples/megatron_bridge/prune_minitron.py`` now uses grouped GEMM by default (pass ``--no_moe_grouped_gemm`` to fall back to ``TESequentialMLP``).
+- Optimize Minitron pruning support for MoE models using the fused **grouped GEMM** experts (``TEGroupedMLP``) in addition to the existing ``SequentialMLP`` path. ``examples/megatron_bridge/prune_minitron.py`` now uses grouped GEMM by default (pass ``--no_moe_grouped_gemm`` to fall back to ``SequentialMLP``).
 - Add Minitron pruning support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/prune_minitron.py``. The language model is pruned while the vision tower is left intact and the full VLM is saved back; ``hidden_size`` is not pruned if it is shared with the vision projector. Pruning importance is estimated from image-text calibration (the full VLM forward over vision-conditioned activations) by default, or from a text dataset for text-only ablations.
 - Add PTQ support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/quantize.py``. Only the language model is quantized (vision tower + projector left in full precision) and the full VLM is saved as a Megatron checkpoint. The calibration modality is inferred from ``--calib_dataset_name``: an image-text dataset drives the full VLM forward (vision-conditioned activations), while a text dataset runs text-only calibration of the language model. Image-text calibration shards across data-parallel ranks (context parallelism is supported only for text-only calibration). HuggingFace unified export of a quantized VLM is not yet supported.
 - Add Megatron-Bridge distillation and Quantization-Aware Distillation (QAD) support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/distill.py``.

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** and **Qwen3.5-VL**. Other VLMs such as Gemma3-VL are saved in Megatron checkpoint format only.
 
 ## Distillation
 

examples/megatron_bridge/distill.py

@@ -50,12 +50,18 @@ from megatron.bridge.utils.vocab_utils import calculate_padded_vocab_size
 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,
+    set_moe_expert_layout,
+    use_moe_grouped_gemm,
+)
 
 with contextlib.suppress(ModuleNotFoundError):
     import modelopt.torch.puzzletron.plugins.mbridge  # noqa: F401
@@ -92,6 +98,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=(
+            "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(
         "--student_megatron_path",
         type=str,
@@ -326,11 +341,26 @@ def _tokenizer_prepends_bos(args) -> bool:
 
 
 def main(args: argparse.Namespace):
+    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")
 
     # 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)
 
@@ -343,6 +373,7 @@ def main(args: argparse.Namespace):
         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
+        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.
@@ -358,17 +389,21 @@ def main(args: argparse.Namespace):
     # 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
-    )
+    # 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.
@@ -390,10 +425,7 @@ def main(args: argparse.Namespace):
 
     # 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(
@@ -420,9 +452,14 @@ def main(args: argparse.Namespace):
             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])
+            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. Resuming from a
+                # language-model-only checkpoint already restores it there.
+                ModeloptStateManager.transfer_state_dict(student, student.language_model)
             return model_chunks
 
         distill_provider.register_pre_wrap_hook(_restore_student_hook, prepend=True)

examples/megatron_bridge/export_distilled_megatron_to_hf.py

@@ -59,12 +59,14 @@ from pathlib import Path
 
 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
 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,
 )
@@ -213,6 +215,16 @@ 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. 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")
     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")
@@ -226,10 +238,15 @@ 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",
-    )
+    # 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:
         # Build the full VLM (vision tower / projector + original LM from HF), then overwrite the LM
@@ -238,6 +255,8 @@ 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,
+            # Mirrors distill.py's unquantized branch
+            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,

examples/megatron_bridge/export_quantized_megatron_to_hf.py

@@ -47,6 +47,7 @@ from modelopt.torch.utils import print_args, print_rank_0
 from modelopt.torch.utils.plugins.mbridge import (
     load_mbridge_model_from_hf,
     load_modelopt_megatron_checkpoint,
+    use_moe_grouped_gemm,
 )
 
 
@@ -71,6 +72,15 @@ 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=(
+            "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(
         "--export_extra_modules",
         action="store_true",
@@ -108,6 +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=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,
@@ -144,9 +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): 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: 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,

examples/megatron_bridge/prune_minitron.py

@@ -66,7 +66,7 @@ from modelopt.torch.utils import (
     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 "

examples/megatron_bridge/quantize.py

@@ -62,6 +62,7 @@ import copy
 import gc
 
 import torch
+from megatron.bridge.models.hf_pretrained.utils import is_safe_repo
 from transformers import AutoProcessor
 
 import modelopt.torch.quantization as mtq
@@ -70,7 +71,11 @@ 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.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,
+    use_moe_grouped_gemm,
+)
 from modelopt.torch.utils.plugins.megatron_calibration import (
     get_megatron_calibration_forward_loop,
     get_megatron_vlm_calibration_forward_loop,
@@ -96,6 +101,15 @@ 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=(
+            "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(
         "--export_megatron_path",
         type=str,
@@ -276,9 +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,
+        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,
@@ -293,8 +318,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)."
@@ -377,7 +401,7 @@ def main(args: argparse.Namespace):
         # 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)
@@ -412,12 +436,12 @@ def main(args: argparse.Namespace):
         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 "
-            "(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."
         )
     else:
         print_rank_0(

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,8 +215,44 @@ 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.
+        # Fetched in two stages: the vision tower is a small fraction of a VLM checkpoint, so
+        # pulling every shard to keep a few would waste tens of GB.
+        local_files_only = _is_hf_hub_offline()
+        repo_id = str(pretrained_model_path)
+        try:
+            index_dir = Path(
+                snapshot_download(
+                    repo_id=repo_id,
+                    allow_patterns=["model.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
+
+        index_file = index_dir / "model.safetensors.index.json"
+        if index_file.is_file():
+            try:
+                weight_map = json.loads(index_file.read_text())["weight_map"]
+            except (json.JSONDecodeError, KeyError) as exc:
+                raise ValueError(f"Malformed safetensors index in {repo_id}.") from exc
+            wanted = sorted(
+                {shard for key, shard in weight_map.items() if key.startswith(prefixes)}
+            )
+        else:
+            wanted = ["model.safetensors"]  # unsharded checkpoint
+        # Kept separate from the resolution failure above: a hub outage or a full disk here is
+        # retryable, not a bad path.
+        hf_checkpoint_path = Path(
+            snapshot_download(
+                repo_id=repo_id,
+                allow_patterns=["model.safetensors.index.json", *wanted],
+                local_files_only=local_files_only,
+            )
         )
 
     safetensors_file = Path(hf_checkpoint_path) / "model.safetensors"
@@ -240,7 +276,13 @@ def load_multimodal_components(
         with open(safetensors_index_file) as f:
             safetensors_index = json.load(f)
 
-        all_shard_files = sorted(set(safetensors_index["weight_map"].values()))
+        all_shard_files = sorted(
+            {
+                shard
+                for key, shard in safetensors_index["weight_map"].items()
+                if key.startswith(prefixes)
+            }
+        )
         for shard_file in all_shard_files:
             safetensors_filepath = Path(hf_checkpoint_path) / shard_file
             with safe_open(safetensors_filepath, framework="pt") as f:
@@ -251,6 +293,12 @@ def load_multimodal_components(
     else:
         print(f"Warning: No safetensors files found in {hf_checkpoint_path}")
 
+    if not multimodal_state_dict:
+        raise ValueError(
+            f"No tensors under {prefixes} in {pretrained_model_path}; the vision tower would be "
+            "missing from the export. The checkpoint's prefixes have likely changed."
+        )
+
     print(f"Successfully loaded {len(multimodal_state_dict)} multimodal tensors")
     return multimodal_state_dict
 

modelopt/torch/export/plugins/mcore_common.py

@@ -39,7 +39,12 @@ from .mcore_qwen import (
     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,17 @@ all_mcore_hf_export_mapping: dict[str, Any] = {
     "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,
+}
+
+# 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_5ForConditionalGeneration": QWEN3_5_VL_VISION_PREFIXES,
+    "Qwen3_5MoeForConditionalGeneration": QWEN3_5_VL_VISION_PREFIXES,
 }
 
 all_mcore_hf_import_mapping: dict[str, Any] = {

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
@@ -127,6 +128,18 @@ class GroupedMLPSlicing(CustomModuleMapping):
         )
 
 
+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."""
 
@@ -175,6 +188,18 @@ class GatedMLPSlicing(CustomModuleMapping):
         )
 
 
+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 +248,29 @@ class UnpackNameRemappingGPT(CustomModuleMapping):
         )
 
 
+# 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]:
+    """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.") :]
+        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()

modelopt/torch/export/plugins/mcore_qwen35vl.py

@@ -0,0 +1,71 @@
+# 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 (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.
+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
+
+# 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."),
+    # 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."),
+    "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}
+)

modelopt/torch/export/plugins/mcore_qwen3vl.py

@@ -15,48 +15,17 @@
 
 """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.  The mappings below are derived automatically from
-the Qwen3 mappings by inserting ``language_model.`` after ``model.`` for every
-prefix that starts with ``model.``.
-
-Note: the visual encoder (``model.visual.*``) is intentionally excluded — this
-mapping covers only the language-model decoder used for quantization and export.
-
-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.
 """
 
-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)

modelopt/torch/export/plugins/megatron_importer.py

@@ -331,10 +331,10 @@ class GPTModelImporter:
         state_dict = module.state_dict()
 
         assert module.num_gemms == num_local_experts, (
-            "num_gemms must be equal to num_local_experts in TEGroupedMLP"
+            "num_gemms must be equal to num_local_experts in TEGroupedLinear"
         )
         # init_expert_id is the global index of this rank's first local expert.
-        # TEGroupedMLP stores weights as weight0..weight{num_local-1} locally, so we
+        # TEGroupedLinear stores weights as weight0..weight{num_local-1} locally, so we
         # map global expert_id -> local slot (expert_id - init_expert_id).
         for local_id in range(num_local_experts):
             global_expert_id = init_expert_id + local_id

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
@@ -28,7 +29,7 @@ from typing import Any
 
 import torch
 import torch.distributed
-from huggingface_hub import hf_hub_download
+from huggingface_hub import get_safetensors_metadata, hf_hub_download
 from huggingface_hub.errors import EntryNotFoundError
 from safetensors import safe_open
 from safetensors.torch import save_file
@@ -53,8 +54,12 @@ from .plugins.hf_checkpoint_utils import (
     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,
@@ -133,7 +138,14 @@ class GPTModelExporter:
         moe_router_dtype: str | None = None,
     ):
         """Create a GPTModel exporter instance."""
-        if not isinstance(model, (GPTModel, HybridModel, LLaVAModel)):
+        # 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))
+            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,22 +165,25 @@ class GPTModelExporter:
         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]
+        # ``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
+        )
         # 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"):
@@ -393,14 +408,13 @@ class GPTModelExporter:
             # 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).
@@ -424,6 +438,84 @@ class GPTModelExporter:
             name_template="model-{:05d}-of-{:05d}",
         )
 
+        # Every rank has written its shards; one rank now checks nothing was dropped. The result is
+        # shared so every rank raises together -- this is public API, and a lone raise would leave
+        # peers hanging in the next collective instead of surfacing the error.
+        torch.distributed.barrier()
+        failure = ""
+        if is_writer_rank:
+            try:
+                self._verify_exported_keys(save_directory, pretrained_model_name_or_path)
+            except Exception as e:
+                # Any escape would strand peers in the all_gather below, which is what this
+                # block exists to prevent.
+                failure = f"{type(e).__name__}: {e}"
+        if torch.distributed.is_initialized():
+            # all_gather rather than broadcast: the writer is not necessarily rank 0, and ``src``
+            # must be identical on every rank.
+            gathered: list = [None] * torch.distributed.get_world_size()
+            torch.distributed.all_gather_object(gathered, failure)
+            failure = next((f for f in gathered if f), "")
+        if failure:
+            raise RuntimeError(failure)
+
+    def _verify_exported_keys(self, save_directory, pretrained_model_name_or_path) -> None:
+        """Raise if the export dropped tensors the source has: a missing rule emits nothing."""
+        if pretrained_model_name_or_path is None:
+            return
+        source_dir = str(pretrained_model_name_or_path)
+        if os.path.isdir(source_dir):
+            source = _read_checkpoint_keys(source_dir)
+        else:
+            # A repo id is the documented invocation. Read the safetensors headers rather than
+            # downloading weights: the export deliberately never fetches them.
+            try:
+                source = set(get_safetensors_metadata(source_dir).weight_map)
+            except Exception:
+                warn_rank_0(
+                    f"Export self-check skipped: cannot read {pretrained_model_name_or_path}."
+                )
+                return
+        index_file = Path(save_directory) / "model.safetensors.index.json"
+        if not index_file.exists():
+            single = Path(save_directory) / "model.safetensors"
+            if not single.exists():
+                warn_rank_0("Export self-check skipped: no safetensors written.")
+                return
+            with safe_open(str(single), framework="pt", device="cpu") as f:
+                exported = set(f.keys())
+        else:
+            with open(index_file) as f:
+                exported = set(json.load(f)["weight_map"])
+        if not source:
+            warn_rank_0(f"Export self-check skipped: no tensor index found in {source_dir}.")
+            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 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.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]}. The checkpoint "
+                f"written to {save_directory} is incomplete -- the architecture has no export "
+                "rule for one of its decoder modules."
+            )
+
     @property
     def state_dict(self):
         """Return the real quantized state_dict of the base model."""
@@ -497,8 +589,12 @@ class GPTModelExporter:
         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:
@@ -531,6 +627,9 @@ class GPTModelExporter:
                     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)
@@ -585,6 +684,13 @@ class GPTModelExporter:
                     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):
@@ -613,6 +719,14 @@ class GPTModelExporter:
                     self.rules["experts.linear_fc2"](
                         layer.mlp.experts.linear_fc2, layer_id, is_mtp=is_mtp
                     )
+                else:
+                    # 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' "
+                        "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)
@@ -629,7 +743,9 @@ class GPTModelExporter:
         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
 
@@ -728,6 +844,16 @@ class GPTModelExporter:
             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)
@@ -863,7 +989,9 @@ class GPTModelExporter:
                 "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,
+                "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,
             }
@@ -1010,8 +1138,11 @@ class GPTModelExporter:
     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,
@@ -1021,6 +1152,7 @@ class GPTModelExporter:
         mapping={},
         dtype: torch.dtype | None = None,
         is_mtp: bool = False,
+        zero_centered_gamma: bool = False,
     ):
         if is_mtp:
             prefix = self._mtp_prefix(prefix)
@@ -1028,13 +1160,21 @@ class GPTModelExporter:
             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 centres this gamma on 0, HF on 1. Assert, don't derive: the config flag
+            # is model-wide while this one is per-norm.
+            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
         weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat)
 
         if weight_scale is None:
@@ -1125,8 +1265,90 @@ class GPTModelExporter:
                 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):
-        """Export TEGroupedMLP weight0..weight{N-1} as one HF-style entry per expert.
+    def _grouped_mlp_packing(self, module, prefix, parallel_config=None, is_mtp=False):
+        """Pack TEGroupedLinear 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,
+                record_quant_config=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
+        handled = (".weight", ".weight_scale", ".weight_scale_2", ".input_scale")
+        unhandled = {k.split(".", 1)[1] for k in per_expert if not k.endswith(handled)}
+        assert not unhandled, (
+            f"{prefix}: grouped-expert packing has no rule for {sorted(unhandled)}"
+        )
+        # 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")
+
+        # 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, rescaled onto the merged global scale.
+                merged_scale, merged_scale_2 = self._merge_nvfp4_expert_scales(scales, scales_2)
+            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,
+        record_quant_config=True,
+    ):
+        """Export TEGroupedLinear 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
@@ -1148,7 +1370,7 @@ class GPTModelExporter:
         )
         if grouped_wq is not None and num_experts > len(grouped_wq):
             warn_rank_0(
-                f"TEGroupedMLP has {num_experts} local experts but only {len(grouped_wq)} "
+                f"TEGroupedLinear has {num_experts} local experts but only {len(grouped_wq)} "
                 f"per-expert weight quantizers; experts >= {len(grouped_wq)} reuse expert "
                 f"{len(grouped_wq) - 1}'s scales (TP/EP-mismatch fallback)."
             )
@@ -1194,11 +1416,11 @@ class GPTModelExporter:
             )
             if missing_flag.item() != 0:
                 raise ValueError(
-                    f"TEGroupedMLP missing expert weights on at least one EP rank "
+                    f"TEGroupedLinear missing expert weights on at least one EP rank "
                     f"(local missing on rank {ep_rank}: {local_missing})"
                 )
         elif local_missing:
-            raise ValueError(f"TEGroupedMLP missing expert weights: {local_missing}")
+            raise ValueError(f"TEGroupedLinear missing expert weights: {local_missing}")
 
         # Per expert, temporarily assign weight = weight{i} and, for the per-expert
         # quantizer layout (GroupedQuantizer), swap in that expert's own TensorQuantizer,
@@ -1245,12 +1467,16 @@ class GPTModelExporter:
                 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()
 
@@ -1273,8 +1499,8 @@ class GPTModelExporter:
 
         # 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:
+        # a TEGroupedLinear layer share qformat/block_size, so local values apply globally.
+        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):
@@ -1302,6 +1528,7 @@ class GPTModelExporter:
             del gathered_bytes
         else:
             self._state_dict.update(local_expert_state)
+        return seen_qformat, seen_block_size
 
     def _qkv_slicing(
         self,
@@ -1338,7 +1565,11 @@ class GPTModelExporter:
         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 (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
 
         weight = name_to_value.pop("weight")
 
@@ -1360,12 +1591,24 @@ class GPTModelExporter:
 
         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]
@@ -1378,7 +1621,16 @@ class GPTModelExporter:
         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, with_gate=False):
+            """Gather ``index`` heads, appending the gate heads for q under gated attention."""
+            taken = tensor[index]
+            if with_gate:
+                taken = torch.cat([taken, tensor[gate_slice]], dim=1)
+            return taken.reshape(-1, last_dim)
+
+        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:
@@ -1393,10 +1645,8 @@ class GPTModelExporter:
                     [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)
-                    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
@@ -1432,7 +1682,7 @@ class GPTModelExporter:
                 # 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, 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
@@ -1441,6 +1691,104 @@ class GPTModelExporter:
                 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 HF's qkv / z / b / a projections.
+
+        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)
+        in_proj = module.in_proj
+        name_to_value, qformat, block_size = self._get_quantized_state(
+            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"],
+            sections["z"],
+            sections["beta"],
+            sections["alpha"],
+        ]
+        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:
+            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 = [
+                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_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
+                )
+                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 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.
+        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:
+                    if proj_prefix in keep_bf16:
+                        continue
+                    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
     ):
@@ -1460,8 +1808,28 @@ class GPTModelExporter:
                 # 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."""
+    @staticmethod
+    def _merge_nvfp4_expert_scales(scales: list, scales_2: list):
+        """Merge per-expert NVFP4 scales onto one global scale, preserving each expert's FP4 range.
+
+        Each expert's block scales were derived against its own ``scale_2``; rescaling them by
+        ``scale_2_i / scale_2_max`` keeps the quieter experts from losing a mantissa bit.
+        """
+        merged_scale_2 = torch.max(torch.stack(scales_2, dim=0), dim=0)[0].clamp_min(
+            torch.finfo(torch.float32).tiny
+        )
+        stacked_2 = torch.stack(scales_2, dim=0).reshape(-1, *([1] * scales[0].dim()))
+        # Rescaling only ever shrinks a block scale, so clamp before the cast: a block already
+        # near the E4M3 floor would otherwise flush to zero and take its weights with it.
+        rescaled = torch.stack(scales, dim=0).to(torch.float32) * (stacked_2 / merged_scale_2)
+        smallest = torch.finfo(scales[0].dtype).smallest_normal
+        merged_scale = torch.where(rescaled > 0, rescaled.clamp_min(smallest), rescaled).to(
+            scales[0].dtype
+        )
+        return merged_scale, merged_scale_2
+
+    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 = []
@@ -1488,10 +1856,10 @@ class GPTModelExporter:
 
         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
@@ -1501,10 +1869,11 @@ class GPTModelExporter:
                 merged_weight_scale = None
         else:
             # 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()
+            merged_weight_scale, merged_weight_scale_2 = self._merge_nvfp4_expert_scales(
+                weight_scale_list, weight_scale_2_list
+            )
+            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]
@@ -1684,6 +2053,20 @@ class GPTModelExporter:
         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,

modelopt/torch/quantization/config.py

@@ -922,7 +922,7 @@ class MaxCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig):
         description=(
             "If True, max-calibration synchronizes the weight quantizer amax across local "
             "experts within each SequentialMLP layer, so all experts in that layer share "
-            "one effective weight amax. TEGroupedMLP keeps a per-expert weight quantizer "
+            "one effective weight amax. TEGroupedLinear keeps a per-expert weight quantizer "
             "(GroupedQuantizer) whose amax follows the same expert-parallel sync rule."
         ),
     )

modelopt/torch/quantization/plugins/megatron.py

@@ -697,8 +697,9 @@ class _MegatronSequentialMLP(DynamicModule):
         """Sync quantizer amax across local experts in a SequentialMLP.
 
         Always syncs input quantizer amax across experts. Optionally syncs weight
-        quantizer amax as well, which matches TEGroupedMLP behavior where all
-        experts are fused into a single GEMM with one quantizer per linear layer.
+        quantizer amax as well, so all experts share one effective weight scale --
+        opt-in, since experts otherwise keep an independent amax each, as
+        ``TEGroupedLinear``'s per-expert ``GroupedQuantizer`` does.
 
         Args:
             sync_weight_amax: If True, also sync weight quantizer amax across experts.
@@ -754,7 +755,7 @@ if HAS_TE:
     ):
         pass
 
-    # Quantized subclasses to support TEGroupedMLP quantization
+    # Quantized subclasses to support TEGroupedLinear quantization
     class _QuantMegatronTEGroupedLinear(_QuantTEGroupedLinear, _MegatronParallelLinear):
         def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
             # _sharded_state_dict_grouped adds _extra_state{gemm_idx} for gemm_idx:[1, num_gemms] in

modelopt/torch/quantization/utils/core_utils.py

@@ -713,8 +713,9 @@ def sync_moe_expert_amax(experts, sync_weight_amax=False):
     1. Takes the element-wise max of each ``input_quantizer`` amax across all experts
        and writes it back, so every expert shares the same input amax.
     2. If ``sync_weight_amax`` is True, also syncs ``weight_quantizer`` amax across
-       experts (max across experts). This matches TEGroupedMLP behavior where all
-       experts share a single weight quantizer.
+       experts (max across experts), so the layer ends up with one effective weight
+       scale. Off by default: experts otherwise keep an independent amax each, which
+       is also what ``TEGroupedLinear``'s per-expert ``GroupedQuantizer`` does.
     3. For any ``weight_quantizer`` that is enabled but has ``amax is None`` (expert
        received no tokens during calibration), runs a weight-only ``max_calibrate``
        to populate the missing amax.

modelopt/torch/utils/plugins/mbridge.py

@@ -14,8 +14,11 @@
 # limitations under the License.
 """Megatron-Bridge plugins for using with Model-Optimizer."""
 
+import re
+from functools import cache
 from typing import Any
 
+import torch
 from megatron.bridge import AutoBridge
 from megatron.bridge.models.gpt_provider import GPTModelProvider
 from megatron.bridge.models.hf_pretrained.utils import is_safe_repo
@@ -30,12 +33,94 @@ from megatron.core.models.gpt import GPTModel
 from megatron.core.models.hybrid.hybrid_model import HybridModel
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.utils import unwrap_model
-from transformers import AutoTokenizer
+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__ = ["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",
+    "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, 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."""
+    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(
@@ -61,7 +146,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.
@@ -84,14 +168,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)
@@ -99,9 +176,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)}"
@@ -118,9 +195,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 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:
@@ -129,10 +211,69 @@ 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 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("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]
+
+    # The expert layout is baked into the checkpoint and the two layouts use different key names,
+    # so a model built with the other one would load its experts as random init: the load is
+    # non-strict. Fires only when the layout is positively identified, so unknown naming is silent.
+    grouped_re = re.compile(r"experts\.linear_fc[12]\.weight\d")
+    ckpt_grouped = any(grouped_re.search(key) for key in checkpoint_keys)
+    ckpt_sequential = any(".local_experts." in key for key in checkpoint_keys)
+    param_names = [name for m in unwrap_model(model) for name, _ in m.named_parameters()]
+    model_has_experts = any(".experts." in name for name in param_names)
+    model_grouped = any(grouped_re.search(name) for name in param_names)
+    if torch.distributed.is_initialized():
+        # checkpoint_keys is global but named_parameters() is this PP stage only, so a stage
+        # holding no MoE layer would otherwise raise on its own and hang the ranks that do not.
+        flags = torch.tensor(
+            [int(model_has_experts), int(model_grouped)],
+            dtype=torch.int,
+            device=torch.cuda.current_device(),
+        )
+        torch.distributed.all_reduce(flags, op=torch.distributed.ReduceOp.MAX)
+        model_has_experts, model_grouped = bool(flags[0].item()), bool(flags[1].item())
+    if model_has_experts and (
+        (ckpt_grouped and not model_grouped) or (ckpt_sequential and model_grouped)
+    ):
+        raise ValueError(
+            f"{megatron_path} stores MoE experts as "
+            f"{'grouped GEMM (TEGroupedMLP)' if ckpt_grouped else 'SequentialMLP'} but the model "
+            f"was built as {'grouped GEMM' if model_grouped else 'SequentialMLP'}; the expert "
+            "weights would load as random init. Rebuild with the matching layout "
+            "(--no_moe_grouped_gemm toggles it)."
+        )
+
     # 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)
+    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):
+            # 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)
+    return model

modelopt_recipes/configs/auto_quantize/units/base_disabled_layers.yaml

@@ -20,6 +20,8 @@
 # modelopt-schema: modelopt.recipe.config.LayerPatternList
   - "*block_sparse_moe.gate*"
   - "*linear_attn.conv1d*"
+  # Megatron-Core names the GatedDeltaNet block ``self_attention``, HF ``linear_attn``.
+  - "*self_attention.conv1d*"
   - "*linear_attn.in_proj_a*"
   - "*linear_attn.in_proj_b*"
   - "*mixer.conv1d*"

modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml

@@ -20,6 +20,9 @@
     enable: false
   - quantizer_name: '*linear_attn.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: '*linear_attn.in_proj_a*'
     enable: false
   - quantizer_name: '*linear_attn.in_proj_b*'

tests/_test_utils/torch/export/unified_checkpoint.py

@@ -0,0 +1,196 @@
+# 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 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
+from pathlib import Path
+
+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",
+    "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 _unpack_nvfp4(packed: torch.Tensor) -> torch.Tensor:
+    """Expand two E2M1 values per ``uint8`` back into a float tensor (low nibble first)."""
+    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 _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(_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(_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.reshape(*scale.shape, *([1] * (weight.ndim - scale.ndim)))
+
+
+def assert_exported_checkpoint_matches(
+    export_dir: Path | str,
+    ref_hf_dir: Path | str,
+    *,
+    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: 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)
+    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) 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))
+    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 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:
+        got, want = _dequantize(key, exported), reference[key].to(torch.float32)
+        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"))
+        else:
+            denom = want.abs().max().clamp_min(torch.finfo(torch.float32).tiny)
+            rel = ((got - want).abs().max() / denom).item()
+            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]}"

tests/_test_utils/torch/megatron/modelopt_state.py

@@ -0,0 +1,52 @@
+# 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
+from torch.distributed.checkpoint import FileSystemReader
+
+__all__ = ["assert_has_modelopt_state", "assert_no_quantizers_matching"]
+
+
+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"
+    )
+
+
+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]}"

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 (
@@ -413,6 +416,45 @@ def _get_tiny_qwen3_5_vl(moe: bool = False, *, tokenizer=None, **config_kwargs)
     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,
@@ -421,7 +463,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,
@@ -430,6 +472,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)
@@ -475,6 +520,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)
 
@@ -482,6 +552,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 = {

tests/examples/megatron_bridge/test_distill.py

@@ -157,7 +157,7 @@ def test_distill_validate_only(tmp_path, num_gpus):
 
 
 # NOTE: Qwen3.5-VL-MoE covered by test_qad.py
-@pytest.mark.timeout(360)  # sometimes times out in CI env with default 300s timeout
+@pytest.mark.timeout(420)  # sometimes times out in CI env with default 300s timeout
 def test_distill_vlm(tmp_path, num_gpus):
     # Self-distillation of a tiny VLM: only the language model is distilled; the vision tower and the
     # vision->language projector must be left byte-for-byte untouched.

tests/examples/megatron_bridge/test_qad.py

@@ -18,8 +18,12 @@ 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,
+    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,
 )
@@ -27,41 +31,29 @@ from _test_utils.torch.transformers_models import (
 
 @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",
     [
-        (lambda tmp_path: create_tiny_qwen3_dir(tmp_path, with_tokenizer=True), False, 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,
-            ),
-            True,
-            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, num_hidden_layers=2
+                tmp_path,
+                with_processor=True,
+                # Cover both Qwen3.5 decoder kinds at the same layer count
+                num_hidden_layers=2,
+                layer_types=["linear_attention", "full_attention"],
             ),
-            True,
-            True,
         ),
     ],
-    ids=["qwen3", "gemma3vl", "qwen3_5_moe_vl"],
+    ids=["qwen3", "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):
     """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.
+    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()
     quantized_megatron_path = tmp_path / "quantized_megatron"
     distill_output_dir = tmp_path / "qad_output"
     train_iters = 3
@@ -81,9 +73,9 @@ def test_qad(tmp_path: Path, num_gpus, create_student, is_vlm, is_moe):
         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)
+    # 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
@@ -113,18 +105,17 @@ def test_qad(tmp_path: Path, num_gpus, create_student, is_vlm, is_moe):
     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"
-    )
-
-    if is_vlm:
-        return  # VLM quantized-HF export is unsupported; stop at the distilled Megatron checkpoint
+    assert_has_modelopt_state(distilled_megatron_path)
 
     # 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"
     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",
+        ],
         hf_model_name_or_path=hf_model_path,
         megatron_path=distilled_megatron_path,
         export_unified_hf_path=hf_export_path,
@@ -133,4 +124,11 @@ def test_qad(tmp_path: Path, num_gpus, create_student, is_vlm, is_moe):
     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 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 (),
+    )

tests/examples/megatron_bridge/test_quantize_export.py

@@ -16,26 +16,48 @@
 
 from pathlib import Path
 
+import pytest
 from _test_utils.examples.run_command import extend_cmd_parts, run_example_command
-from _test_utils.torch.transformers_models import create_tiny_qwen3_dir
+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_nemotron_h_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"),
+    [
+        # 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_moe", "qwen3vl", "nemotron_h"],
+)
+@pytest.mark.timeout(360)  # quantize + export in one test; 1-gpu CI exceeds the default 300s
+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(
@@ -51,21 +73,20 @@ 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(
-        ["torchrun", "--nproc_per_node=1", "export_quantized_megatron_to_hf.py"],
+        ["torchrun", f"--nproc_per_node={num_gpus}", "export_quantized_megatron_to_hf.py"],
         hf_model_name_or_path=hf_model_path,
         megatron_path=megatron_path,
         export_unified_hf_path=hf_export_path,
+        pp_size=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)

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"),
+    [
+        # 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),
+    ],
+)
+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"
+    )

tests/gpu_megatron/torch/export/test_unified_export_megatron.py

@@ -21,11 +21,15 @@ from pathlib import Path
 import pytest
 import torch
 import transformers
-from _test_utils.torch.megatron.models import get_mcore_gpt_model
+from _test_utils.torch.export.unified_checkpoint import assert_exported_checkpoint_matches
+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,
 )
 from safetensors import safe_open
@@ -87,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
@@ -98,6 +129,47 @@ 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
+        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
@@ -114,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)
@@ -166,20 +242,17 @@ def _test_unified_export_megatron(
     if quant_config:
         _verify_model_quant_config(tmp_export_dir, quant_config, kv_cache_quant_cfg)
 
-    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"
+    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,
+            # 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
         vl_model = Qwen3VLForConditionalGeneration.from_pretrained(
             tmp_export_dir, torch_dtype=torch.bfloat16
@@ -194,8 +267,11 @@ def _test_unified_export_megatron(
     ("model_type", "extra_module", "quant_config", "kv_cache_quant_cfg"),
     [
         ("nemotron", None, None, None),
+        # NemotronH (Mamba + attention + grouped-GEMM MoE) is the stronger quantized case, but it
+        # routes no NVFP4 weight through the dense-MLP rules, so keep one dense NVFP4 param too.
         ("nemotron", None, "NVFP4_DEFAULT_CFG", None),
-        ("nemotron", None, "NVFP4_DEFAULT_CFG", "FP8_KV_CFG"),
+        ("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),
@@ -205,6 +281,14 @@ 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),
+        # 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(
@@ -216,6 +300,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