diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md
index 3e711e3618d..a6428d79022 100644
--- a/examples/megatron_bridge/README.md
+++ b/examples/megatron_bridge/README.md
@@ -19,6 +19,8 @@ This directory contains examples of using Model Optimizer with the [NeMo Megatro
> Checkout the [Nemotron-3-Nano-30B-A3B pruning + distillation (with data blend prep) + quantization tutorial](tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md) for a complete end-to-end workflow using Megatron-Bridge!
>
> Or the [Qwen3.6-35B-A3B W4A4 NVFP4 + QAD tutorial](tutorials/Qwen3.6-35B-A3B/README.md) for an end-to-end quantization-aware distillation workflow.
+>
+> Or the [NemotronH omni VL W4A4 + QAD tutorial](tutorials/NemotronH-Omni-W4A4-QAD/README.md) for a Mamba-Transformer MoE vision-language model.
## Pre-Requisites
diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
new file mode 100644
index 00000000000..c2036264699
--- /dev/null
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -0,0 +1,175 @@
+# NemotronH Omni VL: W4A4 NVFP4 PTQ + QAD + Deployment
+
+End-to-end W4A4 quantization of a NemotronH omni vision-language checkpoint (architecture `NemotronH_Omni_Reasoning_V3`), a Mamba-Transformer MoE language model with a vision tower: mixed NVFP4/FP8 post-training quantization → Quantization-Aware Distillation (QAD) → Hugging Face export → vLLM serving. This document covers:
+
+1. **[Data Preparation](#1-data-preparation)**: the chat-format QAD blend, from `data_blend.yaml`
+2. **[Quantization](#2-quantization)**: W4A4 PTQ with `examples/megatron_bridge/quantize.py`
+3. **[QAD](#3-quantization-aware-distillation-qad)**: distillation from the BF16 teacher with `examples/megatron_bridge/distill.py`
+4. **[Export](#4-export)**: the deployable Hugging Face checkpoint
+5. **[Evaluation and Serving](#5-evaluation-and-serving)**: the vLLM settings behind the results
+
+Only the language model (including its MTP head) is quantized; the vision tower and projector stay BF16 and are copied from the source checkpoint.
+
+## Results
+
+All checkpoints were served with the same vLLM settings (see [Section 5](#5-evaluation-and-serving)); the quantized ones with an FP8 KV cache. BF16 is the mean of two runs.
+
+| Benchmark | BF16 (teacher) | W4A4 PTQ | W4A4 QAD (600 iters) |
+| --- | --- | --- | --- |
+| AA-LCR | 68.63 | 68.62 | 67.94 |
+| AA-Omniscience | 26.21 | 25.63 | 25.87 |
+| AA-Omniscience (non-hallucination) | 70.03 | 71.34 | 69.65 |
+| GPQA Diamond | 89.52 | 89.84 | 89.90 |
+| HLE (no tools) | 31.79 | 30.49 | 30.72 |
+| MMLU-ProX | 76.91 | 74.41 | 81.15 |
+| SciCode (subtask) | 47.13 | 46.48 | 46.58 |
+| SciCode (problem) | 15.10 | 13.65 | 15.38 |
+| τ³-bench banking | 18.88 | 17.40 | 20.10 |
+
+MMLU-ProX scores a response only if it ends with the answer phrase in the question's language. Most of the QAD gain comes from fewer responses missing that phrase (11.3% for PTQ, 5.4% after QAD); on the responses with a readable answer, accuracy rises from 83.9 to 85.8.
+
+The exported W4A4 checkpoint is **70 GB**.
+
+---
+
+## Steps to Reproduce
+
+**Environment:** a NeMo 26.08 container (`nvcr.io/nvidia/nemo:26.08`) on GB300 (4 GPUs per node, aarch64), with this repository mounted at `/opt/Model-Optimizer` and Megatron-Bridge `main` (and its pinned Megatron-LM) mounted at `/opt/Megatron-Bridge`. See the [Megatron-Bridge README](../../README.md) for the container setup. Serving NVFP4 checkpoints requires a Blackwell GPU.
+
+PTQ and export launch with `torchrun`; QAD runs on 8 nodes with the `sbatch` script in [Section 3](#3-quantization-aware-distillation-qad).
+
+### 1. Data Preparation
+
+QAD trains on a weighted blend of 11 text-only chat sources, defined in [`data_blend.yaml`](data_blend.yaml) and computing the loss on assistant tokens only. Set `output_dir` in that file, then write the blend as chat JSONL with the [chat JSONL mode of `prepare_megatron_data_blend`](../../../dataset/MEGATRON_DATA_PREP.md#chat-jsonl-blends-for-assistant-only-loss):
+
+```bash
+python -m modelopt.torch.utils.plugins.prepare_megatron_data_blend --config data_blend.yaml
+```
+
+This writes `train.jsonl` (35,800 conversations) and `validation.jsonl` (1,000) to `output_dir`. Run it in the container: Megatron-Bridge's text-only chat filter selects the rows.
+
+---
+
+### 2. Quantization
+
+The recipe is [`model_type/nemotron_h/ptq/nvfp4-aggressive-mse`](../../../../modelopt_recipes/model_type/nemotron_h/ptq/nvfp4-aggressive-mse.yaml), the same file `hf_ptq.py` uses: its rules list both the Hugging Face and the Megatron-Core module names where they differ.
+
+| Component | Precision |
+| --- | --- |
+| MoE routed and shared experts | **NVFP4 W4A4**, block 16, MSE-selected static weight scales |
+| Mamba `in_proj` / `out_proj` | **FP8 W8A8** per-tensor |
+| Attention `q`/`k`/`v`/`o_proj` (Megatron-Core `linear_qkv` / `linear_proj`) | **FP8 W8A8** per-tensor |
+| `lm_head` (Megatron-Core `output_layer`) | **FP8 W8A8** |
+| KV cache | **FP8** cast at a fixed range (no `k_scale` / `v_scale` exported) |
+| MTP block-1 routed and shared experts | **NVFP4 W4A4**, as above |
+| MTP block-0 KV cache | **FP8**, calibrated |
+| Vision tower, projector, the rest of the MTP head, latent MoE, routers, embeddings | BF16 |
+
+PTQ runs on **one 4-GPU GB300 node** with 512 calibration samples.
+
+
+PTQ command (click to expand)
+
+```bash
+torchrun --nproc_per_node 4 /opt/Model-Optimizer/examples/megatron_bridge/quantize.py \
+ --hf_model_name_or_path \
+ --trust_remote_code \
+ --recipe model_type/nemotron_h/ptq/nvfp4-aggressive-mse \
+ --tp_size 1 --ep_size 4 --pp_size 1 \
+ --calib_dataset_name cnn_nemotron_v2_mix \
+ --calib_num_samples 512 \
+ --calib_batch_size 1 \
+ --export_megatron_path /path/to/nemotron_h_omni_w4a4_megatron \
+ --skip_generate
+```
+
+- `--tp_size 1` is **required**: with tensor parallelism, static NVFP4 weight scales are per-shard but the Megatron checkpoint saves them as replicated, so calibration raises `NotImplementedError` for TP > 1. Expert parallelism spreads the experts instead.
+- Calibration is text-only (`cnn_nemotron_v2_mix`): the language model is calibrated on its own, and the vision tower is not run.
+- The MTP head is quantized too, so the exported checkpoint has it in the serving format: its MoE block (`language_model.mtp.layers.0.mtp_model_layer.layers.1` in Megatron-Core) gets the same NVFP4 treatment as the backbone experts, and the rest of the head stays BF16.
+
+
+
+---
+
+### 3. Quantization-Aware Distillation (QAD)
+
+QAD trains the quantized student against the frozen BF16 teacher with a logit KL loss only (no cross-entropy). See the [QAD section of the Megatron-Bridge README](../../README.md#quantization-aware-distillation-qad).
+
+We used **8 nodes × 4 GB300 (32 GPUs)** with context parallelism 4 and expert parallelism 32.
+
+
+QAD command (click to expand)
+
+`srun` starts one task per GPU (32 in all), and each task maps its Slurm rank variables to the `RANK` / `WORLD_SIZE` / `LOCAL_RANK` that `torch.distributed` reads, with a shared `MASTER_ADDR` / `MASTER_PORT`. Add the container options your cluster needs (image, mounts as in the Environment note) to `srun`.
+
+```bash
+#!/bin/bash
+#SBATCH --nodes=8 --ntasks-per-node=4 --gpus-per-node=4
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) MASTER_PORT=29500
+srun bash -c 'export RANK=$SLURM_PROCID WORLD_SIZE=$SLURM_NTASKS LOCAL_RANK=$SLURM_LOCALID
+python -u /opt/Model-Optimizer/examples/megatron_bridge/distill.py \
+ --teacher_hf_path \
+ --student_hf_path \
+ --student_megatron_path /path/to/nemotron_h_omni_w4a4_megatron \
+ --trust_remote_code \
+ --tp_size 1 --pp_size 1 --cp_size 4 --ep_size 32 \
+ --sft --sft_hf_dataset /path/to/qad_blend/train.jsonl \
+ --sft_hf_validation /path/to/qad_blend/validation.jsonl \
+ --sft_loss_mode assistant \
+ --freeze_mtp \
+ --seq_length 32768 --mbs 1 --gbs 64 \
+ --train_iters 600 \
+ --lr 2e-5 --min_lr 5e-6 --lr_warmup_iters 30 \
+ --recompute_granularity full --recompute_method uniform --recompute_num_layers 1 \
+ --no_async_save \
+ --save_interval 100 --checkpoint_keep_last -1 \
+ --eval_iters 4 --eval_interval 100 \
+ --log_interval 10 \
+ --output_dir /path/to/qad_output'
+```
+
+- `--student_megatron_path` is the PTQ checkpoint from Section 2; `--student_hf_path` still points at the BF16 model, which supplies the architecture. The PTQ checkpoint was written at EP=4 and loads at EP=32.
+- `--sft --sft_hf_dataset ... --sft_loss_mode assistant` renders each conversation with the model's chat template and computes the loss on assistant tokens only.
+- `--seq_length 32768 --gbs 64`: 600 iterations cover 38,400 records, about 1.07 epochs of the 35,800-record training split.
+- `--recompute_*` and `--no_async_save` are needed to fit: the async-save worker needs its own CUDA context and runs out of memory next to the resident teacher.
+- `--freeze_mtp` keeps the MTP head out of training: its loss is skipped and its parameters are frozen, so the exported MTP head keeps its PTQ weights. Without it, the quantized MTP head's cross-entropy loss trains the head and, through the hidden states, embeddings and output layer it shares, the rest of the student too.
+- Each saved iteration is about 1.6 TB; size `--save_interval` and `--checkpoint_keep_last` to your storage.
+
+
+
+---
+
+### 4. Export
+
+Export the QAD checkpoint to a Hugging Face checkpoint on one 4-GPU node (pipeline parallel 4):
+
+```bash
+torchrun --nproc_per_node 4 /opt/Model-Optimizer/examples/megatron_bridge/export_quantized_megatron_to_hf.py \
+ --hf_model_name_or_path \
+ --trust_remote_code \
+ --megatron_path /path/to/qad_output/checkpoints \
+ --export_unified_hf_path /path/to/nemotron_h_omni_w4a4_hf \
+ --pp_size 4
+```
+
+- The exporter writes the quantized language model, including the MTP head, under `language_model.`, and copies `vision_model.`, `vision_projector.` and `mlp1.` from the source checkpoint.
+- The MTP head exports through the standard language-model mapping; `--export_extra_modules` is for EAGLE and Medusa heads and is not needed here.
+- The MTP head can be replaced after export by a separately trained MTP head in the same format.
+- It exports the latest `iter_*` under `--megatron_path`. To export another iteration, point `--megatron_path` at a directory holding a link to that `iter_*` and a matching `latest_checkpointed_iteration.txt`.
+- The backbone KV cache is cast at the fixed FP8 range, so the export has no backbone `k_scale` / `v_scale`; vLLM then uses a scale of 1.0, which is that range. Only the MTP attention carries a calibrated pair.
+
+---
+
+### 5. Evaluation and Serving
+
+Serve with vLLM on one 4-GPU GB300 node, TP=4 with expert parallelism. Quantized checkpoints use these engine arguments; the BF16 baseline uses `--kv-cache-dtype auto`.
+
+| Setting | Value |
+| --- | --- |
+| `--kv-cache-dtype` | `fp8` |
+| `--max-model-len` | 262144 |
+| `--mamba-ssm-cache-dtype` | `float16` |
+| `--mamba-conv-cache-dtype` | `float32` |
+| Mamba cache rounding | stochastic rounding, 5 Philox rounds |
+| Reasoning parser | `nemotron_v3` |
+| Chat template | the checkpoint's `chat_template.jinja` |
diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/data_blend.yaml b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/data_blend.yaml
new file mode 100644
index 00000000000..183ff660882
--- /dev/null
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/data_blend.yaml
@@ -0,0 +1,68 @@
+# Weighted text-only chat blend for this tutorial's QAD, distilled on assistant tokens only
+# (`distill.py --sft --sft_hf_dataset ... --sft_loss_mode assistant`). Prepare it with
+# `python -m modelopt.torch.utils.plugins.prepare_megatron_data_blend --config data_blend.yaml`;
+# see examples/dataset/MEGATRON_DATA_PREP.md. Set `output_dir` first.
+
+output_dir: /path/to/qad_blend
+output_format: chat_jsonl
+# 600 iterations x 64 conversations read the 35,800-row training split about once.
+target_samples: 36_800
+# Take the first usable rows of each split in order, then mix them with `seed`.
+shuffle: false
+seed: 1234
+sources:
+ - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1
+ config: Nemotron-SFT-Code
+ split: train
+ content_field: text
+ weight: 5
+ - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1
+ config: Nemotron-SFT-General
+ split: train
+ content_field: text
+ weight: 20
+ - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1
+ config: Nemotron-SFT-MATH
+ split: train
+ content_field: text
+ weight: 5
+ - hf_dataset: nvidia/Nemotron-Math-v2
+ config: default
+ split: high_part00
+ content_field: messages
+ weight: 10
+ - hf_dataset: nvidia/Nemotron-SFT-Math-v3
+ config: default
+ split: train
+ content_field: messages
+ weight: 17
+ - hf_dataset: nvidia/Nemotron-Competitive-Programming-v1
+ config: default
+ split: competitive_coding_python_part00
+ content_field: messages
+ weight: 15
+ - hf_dataset: nvidia/Nemotron-Competitive-Programming-v1
+ config: default
+ split: competitive_coding_cpp_part00
+ content_field: messages
+ weight: 5
+ - hf_dataset: nvidia/Nemotron-Post-Training-Dataset-v1
+ config: default
+ split: stem
+ content_field: messages
+ weight: 8
+ - hf_dataset: nvidia/Nemotron-Science-v1
+ config: default
+ split: MCQ
+ content_field: messages
+ weight: 3
+ - hf_dataset: nvidia/Nemotron-Science-v1
+ config: default
+ split: RQA
+ content_field: messages
+ weight: 2
+ - hf_dataset: nvidia/Nemotron-SFT-Instruction-Following-Chat-v2
+ config: default
+ split: reasoning_off
+ content_field: messages
+ weight: 2
diff --git a/examples/megatron_bridge/tutorials/README.md b/examples/megatron_bridge/tutorials/README.md
index ced00fde2a5..454403083e8 100644
--- a/examples/megatron_bridge/tutorials/README.md
+++ b/examples/megatron_bridge/tutorials/README.md
@@ -8,4 +8,5 @@ Each one walks through a complete workflow using the scripts in [examples/megatr
| Tutorial | What it covers |
| --- | --- |
| [NVIDIA-Nemotron-3-Nano-30B-A3B-BF16](NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md) | End-to-end optimization of the Nemotron-3-Nano-30B-A3B-BF16 (MoE + Mamba-Transformer hybrid) model: Minitron structured pruning (31.6B/A3.6B → 22B/A3.0B) → two-phase knowledge distillation (100B tokens, 8K then 32K seq length) → quantization → vLLM deployment. Includes data-blend preparation, evaluation setup, and detailed pruning / data-blend / long-context ablations. |
+| [NemotronH-Omni-W4A4-QAD](NemotronH-Omni-W4A4-QAD/README.md) | **W4A4** quantization of the language model of a NemotronH omni vision-language checkpoint (`NemotronH_Omni_Reasoning_V3`, Mamba-Transformer MoE): mixed NVFP4/FP8 PTQ → KD-only QAD on chat data → Hugging Face export → vLLM serving. |
| [Qwen3.6-35B-A3B](Qwen3.6-35B-A3B/README.md) | End-to-end **W4A4 NVFP4** optimization of the Qwen3.6-35B-A3B (MoE + hybrid linear/full attention VLM) model: NVFP4 W4A4 post-training quantization → quantization-aware distillation (QAD) to recover the accuracy W4A4 costs → evaluation → vLLM deployment. |
diff --git a/modelopt_recipes/model_type/nemotron_h/ptq/nvfp4-aggressive-mse.yaml b/modelopt_recipes/model_type/nemotron_h/ptq/nvfp4-aggressive-mse.yaml
index 9445dd3b25f..e539eb881f9 100644
--- a/modelopt_recipes/model_type/nemotron_h/ptq/nvfp4-aggressive-mse.yaml
+++ b/modelopt_recipes/model_type/nemotron_h/ptq/nvfp4-aggressive-mse.yaml
@@ -29,6 +29,10 @@
# quant_cfg is ordered and the LAST match wins, so the blanket disable comes first and the
# vision / MTP rules come last. Only configs/numerics/* is imported, so every rule that
# applies to this model is visible in this file.
+#
+# Rules are written for both Hugging Face (hf_ptq) and Megatron-Core (Megatron-Bridge) module
+# names. Where the two differ -- attention, lm_head and the MTP block -- both spellings are
+# listed; each is inert on the other framework's model.
imports:
nvfp4: configs/numerics/nvfp4
@@ -84,6 +88,19 @@ quantize:
- quantizer_name: '*mixer.out_proj.input_quantizer'
cfg:
$import: fp8
+ # Megatron-Core attention: q/k/v_proj are fused into linear_qkv, o_proj is linear_proj.
+ - quantizer_name: '*self_attention.linear_qkv.weight_quantizer'
+ cfg:
+ $import: fp8
+ - quantizer_name: '*self_attention.linear_qkv.input_quantizer'
+ cfg:
+ $import: fp8
+ - quantizer_name: '*self_attention.linear_proj.weight_quantizer'
+ cfg:
+ $import: fp8
+ - quantizer_name: '*self_attention.linear_proj.input_quantizer'
+ cfg:
+ $import: fp8
# Backbone KV casts at a fixed FP8 range: no _amax buffer, so no k_scale/v_scale is
# exported, matching configs/ptq/units/kv_fp8_cast. MTP KV below stays calibrated and
# keeps its pair.
@@ -105,7 +122,22 @@ quantize:
- quantizer_name: '*mtp.layers.1.mixer.shared_experts.*input_quantizer'
cfg:
$import: nvfp4
- # Restated so this rule cannot inherit the backbone setting above.
+ # Megatron-Core MTP block 1 is mtp.layers.0.mtp_model_layer.layers.1.
+ - quantizer_name: '*mtp_model_layer.layers.1.mlp.experts.*weight_quantizer'
+ cfg:
+ $import: nvfp4_static
+ - quantizer_name: '*mtp_model_layer.layers.1.mlp.experts.*input_quantizer'
+ cfg:
+ $import: nvfp4
+ - quantizer_name: '*mtp_model_layer.layers.1.mlp.shared_experts.*weight_quantizer'
+ cfg:
+ $import: nvfp4_static
+ - quantizer_name: '*mtp_model_layer.layers.1.mlp.shared_experts.*input_quantizer'
+ cfg:
+ $import: nvfp4
+ # Restated so this rule cannot inherit the backbone setting above. It also matches the
+ # Megatron-Core MTP attention (mtp.layers.0.mtp_model_layer.layers.0...), so it needs no
+ # separate Megatron-Core rule.
- quantizer_name: '*mtp.layers.0.*[kv]_bmm_quantizer'
cfg:
$import: fp8
@@ -120,10 +152,23 @@ quantize:
enable: true
cfg:
$import: fp8
+ # Megatron-Core lm_head. Effective only with untied embeddings: for a tied output_layer the
+ # Megatron checkpoint save drops its quantizer state, so it reloads and exports BF16.
+ - quantizer_name: '*output_layer.weight_quantizer'
+ enable: true
+ cfg:
+ $import: fp8
+ - quantizer_name: '*output_layer.input_quantizer'
+ enable: true
+ cfg:
+ $import: fp8
- quantizer_name: '*vision_model*'
enable: false
- quantizer_name: '*vision_projector*'
enable: false
+ # Megatron-Bridge names the projector vision_projection.
+ - quantizer_name: '*vision_projection*'
+ enable: false
- quantizer_name: '*radio_model*'
enable: false
- quantizer_name: '*vision_tower*'
@@ -142,6 +187,8 @@ quantize:
# quantization like any other. Image calibration supplies no audio, so those input
# quantizers collect no amax and the export drops the input scale while still writing an
# FP8 weight scale -- a layer that claims W8A8 with no activation scale. Keep it BF16.
+ # Megatron-Bridge builds no audio branch for the omni wrapper; these guards also cover the
+ # Megatron-Core attention rules should it name one sound_encoder / sound_projector.
- quantizer_name: '*audio_tower*'
enable: false
- quantizer_name: '*sound_encoder*'
diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md
index d7eb46775ea..ee37ec72dcf 100644
--- a/modelopt_recipes/ptq.md
+++ b/modelopt_recipes/ptq.md
@@ -423,6 +423,11 @@ recipes, for Step-3.7 checkpoints; Step-3.5 has its own recipe above.
`static`. Use `--calib_with_images` on VL checkpoints: the vision tower must see image batches
even though it stays BF16.
+ Rules list both Hugging Face and Megatron-Core spellings where they differ (attention
+ `linear_qkv`/`linear_proj`, `output_layer`, the MTP `mtp_model_layer.layers`), so the same
+ recipe applies under `examples/megatron_bridge/quantize.py`; there, static NVFP4 weights need
+ `--tp_size 1`.
+
### Algorithm overrides — `gemma`, `gemma4`, `mpt`
These quantize the **same layers** as the general recipes; only the
diff --git a/tests/unit/recipe/test_nemotron_h_recipe.py b/tests/unit/recipe/test_nemotron_h_recipe.py
index b29dd01795f..2d5d531ab60 100644
--- a/tests/unit/recipe/test_nemotron_h_recipe.py
+++ b/tests/unit/recipe/test_nemotron_h_recipe.py
@@ -66,12 +66,22 @@ def _setup(self):
self.v_bmm_quantizer = TensorQuantizer()
+class _CoreAttention(nn.Module):
+ """Megatron-Core's ``self_attention.core_attention``, the module that holds the KV quantizers."""
+
+
+class _QuantCoreAttention(_QuantAttention):
+ pass
+
+
@pytest.fixture(autouse=True)
def _register_kv_attention():
mtq.register(original_cls=_Attention, quantized_cls=_QuantAttention)
+ mtq.register(original_cls=_CoreAttention, quantized_cls=_QuantCoreAttention)
yield
- if QuantModuleRegistry.get(_Attention) is not None:
- mtq.unregister(_Attention)
+ for cls in (_Attention, _CoreAttention):
+ if QuantModuleRegistry.get(cls) is not None:
+ mtq.unregister(cls)
class _Mamba(nn.Module):
@@ -130,9 +140,61 @@ def _nemotron_h(with_mtp: bool) -> nn.Module:
return model
-def _quantize(with_mtp: bool):
+def _mcore_attention() -> nn.Module:
+ layer = nn.Module()
+ layer.self_attention = nn.Module()
+ layer.self_attention.linear_qkv = nn.Linear(_WIDTH, 3 * _WIDTH, bias=False)
+ layer.self_attention.linear_proj = nn.Linear(_WIDTH, _WIDTH, bias=False)
+ layer.self_attention.core_attention = _CoreAttention()
+ return layer
+
+
+def _mcore_moe() -> nn.Module:
+ layer = nn.Module()
+ layer.mlp = nn.Module()
+ layer.mlp.router = nn.Linear(_WIDTH, 2, bias=False)
+ for name in ("experts", "shared_experts"):
+ mlp = nn.Module()
+ mlp.linear_fc1 = nn.Linear(_WIDTH, _WIDTH, bias=False)
+ mlp.linear_fc2 = nn.Linear(_WIDTH, _WIDTH, bias=False)
+ setattr(layer.mlp, name, mlp)
+ return layer
+
+
+def _nemotron_h_mcore() -> nn.Module:
+ """Megatron-Core-named omni VL skeleton, as built by Megatron-Bridge.
+
+ Names follow a real PTQ checkpoint: the language model sits under ``language_model``,
+ attention is ``self_attention.linear_qkv`` / ``linear_proj``, lm_head is ``output_layer``,
+ and the MTP block nests its attention and MoE layers under ``mtp_model_layer.layers``. The
+ projector is ``vision_projection``. Megatron-Bridge builds no audio branch, but one is
+ included so the test pins the audio guards against the Megatron-Core attention rules too.
+ """
+ model = nn.Module()
+ lm = model.language_model = nn.Module()
+ lm.embedding = nn.Linear(_WIDTH, _WIDTH, bias=False)
+ lm.decoder = nn.Module()
+ lm.decoder.layers = nn.ModuleList([_mcore_attention(), _Block(_Mamba()), _mcore_moe()])
+ lm.output_layer = nn.Linear(_WIDTH, _WIDTH, bias=False)
+ mtp_block = nn.Module()
+ mtp_block.eh_proj = nn.Linear(2 * _WIDTH, _WIDTH, bias=False)
+ mtp_block.mtp_model_layer = nn.Module()
+ mtp_block.mtp_model_layer.layers = nn.ModuleList([_mcore_attention(), _mcore_moe()])
+ lm.mtp = nn.Module()
+ lm.mtp.layers = nn.ModuleList([mtp_block])
+ model.vision_model = nn.Module()
+ model.vision_model.blocks = nn.ModuleList([_Expert(), _mcore_attention()])
+ # An attention layer stands in under the projector to prove its guard.
+ model.vision_projection = nn.ModuleList([_mcore_attention()])
+ model.sound_projector = nn.Module()
+ model.sound_projector.sound_encoder = nn.Module()
+ model.sound_projector.sound_encoder.layers = nn.ModuleList([_mcore_attention()])
+ return model
+
+
+def _quantize(with_mtp: bool, mcore: bool = False):
"""Apply the recipe to a fresh skeleton and return it with its named modules."""
- model = _nemotron_h(with_mtp)
+ model = _nemotron_h_mcore() if mcore else _nemotron_h(with_mtp)
config = load_recipe(_RECIPE).quantize.model_dump()
# Calibration settings the recipe promises, asserted before they are dropped: weight-MSE
@@ -314,3 +376,48 @@ def test_no_quantizer_outside_the_intended_set_is_enabled():
expected.add(f"lm_head.{kind}")
assert enabled == expected
+
+
+def test_megatron_core_names_get_the_same_formats():
+ """Under Megatron-Bridge the recipe selects the same layers as on a Hub checkpoint.
+
+ Attention, lm_head and the MTP block have Megatron-Core names of their own; without those
+ spellings attention and lm_head would silently stay BF16. Everything else is shared.
+ """
+ _, modules = _quantize(with_mtp=True, mcore=True)
+ lm, mtp = "language_model.decoder.layers", "language_model.mtp.layers.0.mtp_model_layer.layers"
+
+ expected = set()
+ for kind in ("weight_quantizer", "input_quantizer"):
+ for stem in (
+ f"{lm}.0.self_attention.linear_qkv",
+ f"{lm}.0.self_attention.linear_proj",
+ f"{lm}.1.mixer.in_proj",
+ f"{lm}.1.mixer.out_proj",
+ "language_model.output_layer",
+ ):
+ _assert_fp8(modules[f"{stem}.{kind}"], stem)
+ expected.add(f"{stem}.{kind}")
+ for prefix in (f"{lm}.2", f"{mtp}.1"):
+ for mlp in ("experts", "shared_experts"):
+ for fc in ("linear_fc1", "linear_fc2"):
+ stem = f"{prefix}.mlp.{mlp}.{fc}"
+ block_type = "static" if kind == "weight_quantizer" else "dynamic"
+ _assert_nvfp4(modules[f"{stem}.{kind}"], stem, block_type=block_type)
+ expected.add(f"{stem}.{kind}")
+
+ for which in ("k", "v"):
+ backbone = f"{lm}.0.self_attention.core_attention.{which}_bmm_quantizer"
+ _assert_fp8(modules[backbone], backbone)
+ assert modules[backbone]._use_constant_amax is True, f"{backbone} should cast"
+ mtp_kv = f"{mtp}.0.self_attention.core_attention.{which}_bmm_quantizer"
+ _assert_fp8(modules[mtp_kv], mtp_kv)
+ assert modules[mtp_kv]._use_constant_amax is False, f"{mtp_kv} should stay calibrated"
+ expected |= {backbone, mtp_kv}
+
+ enabled = {
+ name
+ for name, module in modules.items()
+ if isinstance(module, TensorQuantizer) and module.is_enabled
+ }
+ assert enabled == expected