From 9ac568e30abbe4ebcbfc23a8b1598357a89466b4 Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Fri, 9 Oct 2026 05:42:30 +0000
Subject: [PATCH 01/24] Add a NemotronH omni VL W4A4 PTQ + QAD tutorial for
 Megatron-Bridge

examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD: the data blend, PTQ,
KD-only QAD, HF export and vLLM serving settings behind the reported results for a
NemotronH_Omni_Reasoning_V3 checkpoint, with two helpers: build_qad_blend.py
(materializes the public chat blend) and add_nvfp4_kv_scales.py (writes the 1/6
k/v scales that constant-amax NVFP4 KV quantizers imply, which the export does
not emit). Linked from the Megatron-Bridge READMEs.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 examples/megatron_bridge/README.md            |   2 +
 .../NemotronH-Omni-W4A4-QAD/README.md         | 198 ++++++++++++++++++
 .../add_nvfp4_kv_scales.py                    |  70 +++++++
 .../build_qad_blend.py                        | 129 ++++++++++++
 examples/megatron_bridge/tutorials/README.md  |   1 +
 5 files changed, 400 insertions(+)
 create mode 100644 examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
 create mode 100644 examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/add_nvfp4_kv_scales.py
 create mode 100644 examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/build_qad_blend.py

diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md
index 3e711e3618d..87243070506 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 with an NVFP4 KV cache.
 
 ## 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..82f60391ae5
--- /dev/null
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -0,0 +1,198 @@
+# 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
+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 is quantized; the vision tower, projector and MTP head stay BF16 and are copied from the source checkpoint.
+
+## Results
+
+All checkpoints were served with the same vLLM settings ([Section 5](#5-evaluation-and-serving)); quantized ones with an NVFP4 KV cache.
+
+| Model | AA-LCR | SciCode (subtask) | GPQA Diamond pass@1 (maj@8) |
+| --- | --- | --- | --- |
+| **BF16** (teacher) | 63.31 | 46.48 | 89.52 (90.66) |
+| **W4A4 PTQ** (QAD student) | 62.00 | 46.92 | 89.27 (90.66) |
+| ↳ QAD 200 iters | 61.69 | 43.62 | 88.70 (89.65) |
+| ↳ QAD 300 iters | 62.56 | 45.53 | 88.70 (88.13) |
+| ↳ QAD 400 iters | 64.06 | 45.36 | 88.89 (90.40) |
+| ↳ QAD 500 iters | 62.94 | 46.22 | 89.84 (90.66) |
+| **↳ QAD 600 iters** | **63.75** | **46.35** | **89.77 (91.67)** |
+
+Sample counts: AA-LCR 1,600 (100 questions × 16), SciCode 520, GPQA Diamond 1,584 (198 questions × 8).
+
+The recipe is already close to lossless after PTQ: every PTQ score is within 1.4 points of BF16. The first few hundred QAD iterations dip before recovering, and QAD 600 lands within 0.5 points of BF16 on all three benchmarks. The exported W4A4 checkpoint is **74 GB**.
+
+> [!NOTE]
+> The numbers above come from runs whose code was functionally equivalent to the commands below but predates them: the Omni export mapping and chat-data distillation were local patches then, and the student was built without its MTP head (`--mtp_num_layers 0`). This code path was checked end to end with a short run (a 32-sample PTQ, 20 QAD iterations and the export); QAD keeps the MTP head unchanged either way (see [Section 3](#3-quantization-aware-distillation-qad)).
+
+---
+
+## 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` at [`d23c1f98`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/commit/d23c1f981596c227b13d98c95b09afb865d80258) (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.
+
+Every command below runs with `srun`, one task per GPU, and sets `RANK` / `WORLD_SIZE` / `LOCAL_RANK` from `SLURM_PROCID` / `SLURM_NTASKS` / `SLURM_LOCALID` in the `srun` body. Pass `--trust_remote_code`, since the model ships its own modeling code.
+
+### 1. Data Preparation
+
+QAD distills on text chat data, read directly from jsonl by `distill.py --sft_hf_dataset`. [`build_qad_blend.py`](build_qad_blend.py) materializes a 36,800-record sample of public Nemotron post-training datasets, 400 records per weight point:
+
+| Dataset / split | Weight | Records |
+| --- | --- | --- |
+| Nemotron-Pretraining-SFT-v1 / Nemotron-SFT-General | 20 | 8,000 |
+| Nemotron-SFT-Math-v3 / train | 17 | 6,800 |
+| Nemotron-Competitive-Programming-v1 / python_part00 | 15 | 6,000 |
+| Nemotron-Math-v2 / high_part00 | 10 | 4,000 |
+| Nemotron-Post-Training-Dataset-v1 / stem | 8 | 3,200 |
+| Nemotron-Pretraining-SFT-v1 / Nemotron-SFT-Code | 5 | 2,000 |
+| Nemotron-Pretraining-SFT-v1 / Nemotron-SFT-MATH | 5 | 2,000 |
+| Nemotron-Competitive-Programming-v1 / cpp_part00 | 5 | 2,000 |
+| Nemotron-Science-v1 / MCQ | 3 | 1,200 |
+| Nemotron-Science-v1 / RQA | 2 | 800 |
+| Nemotron-SFT-Instruction-Following-Chat-v2 / reasoning_off | 2 | 800 |
+
+The script also lists `Nemotron-SFT-Instruction-Following-Chat-v2 / reasoning_on` and `Nemotron-Agentic-v1 / tool_calling`; in our run they contributed no records (tool-calling rows are not text-only chat), so the blend has no agentic data. After a seeded shuffle, 1,000 records are held out for validation and 35,800 are used for training.
+
+```bash
+python examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/build_qad_blend.py \
+    --output_dir /path/to/qad_blend
+```
+
+Records follow the `{"messages": [{"role": ..., "content": ...}, ...]}` schema and are rendered with the model's own chat template during QAD. Records longer than the 32K sequence length used below are truncated.
+
+---
+
+### 2. Quantization
+
+The recipe is a fixed mixed-precision scheme; it will be published under `modelopt_recipes`, and its path goes to `--recipe` below. Counts are read from the exported checkpoint's `hf_quant_config.json`:
+
+| Component | Precision | Tensors |
+| --- | --- | --- |
+| MoE routed experts (`up`/`down_proj`) | **NVFP4 W4A4**, block 16, MSE-selected static weight scales | 40,960 |
+| MoE shared experts | **NVFP4 W4A4** | 80 |
+| Mamba `in_proj` / `out_proj` | **FP8 W8A8** per-tensor | 80 |
+| Attention `q`/`k`/`v`/`o_proj` (8 layers) | **FP8 W8A8** per-tensor | 32 |
+| `lm_head` | **FP8 W8A8** | 1 |
+| KV cache | **NVFP4** cast, constant amax 448 | — |
+| Vision tower, projector, MTP head, routers, embeddings | BF16 | — |
+
+PTQ takes **16 minutes on one 4-GPU GB300 node** with 512 calibration samples.
+
+<details>
+<summary>PTQ command (click to expand)</summary>
+
+```bash
+# SBATCH --nodes=1 --ntasks-per-node=4 --gpus-per-node=4
+srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/quantize.py \
+    --hf_model_name_or_path <nemotron_h_omni-checkpoint> \
+    --trust_remote_code \
+    --recipe /path/to/recipe.yaml \
+    --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**: static NVFP4 weight scales are calibrated per full weight, so the recipe needs 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 MoE experts are built as grouped GEMMs (`TEGroupedMLP`), which is what the recipe's `mlp.experts.linear_fc*` rules match. The log reports `moe_grouped_gemm: true`; do not pass `--no_moe_grouped_gemm`.
+- The MTP head is built and stays BF16 (the recipe disables `language_model.mtp.*`).
+
+</details>
+
+---
+
+### 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. 600 iterations took **7.7 hours** (about 44 s/iteration at steady state). The validation KD loss fell from 1.36e-2 at iteration 100 to 1.01e-2 at iteration 600.
+
+<details>
+<summary>QAD command (click to expand)</summary>
+
+```bash
+# SBATCH --nodes=8 --ntasks-per-node=4 --gpus-per-node=4
+srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/distill.py \
+    --teacher_hf_path <nemotron_h_omni-checkpoint> \
+    --student_hf_path <nemotron_h_omni-checkpoint> \
+    --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 \
+    --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.
+- The MTP head is built but receives no loss under KD-only training (`mtp_1 loss` and `mtp_2 loss` log as 0), so it stays bit-identical to the BF16 source.
+- Each saved iteration is about 1.6 TB; size `--save_interval` and `--checkpoint_keep_last` to your storage.
+
+</details>
+
+---
+
+### 4. Export
+
+Export the QAD checkpoint to a Hugging Face checkpoint on one 4-GPU node (pipeline parallel 4, about 26 minutes), then add the NVFP4 KV-cache scales:
+
+```bash
+# SBATCH --nodes=1 --ntasks-per-node=4 --gpus-per-node=4
+srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/export_quantized_megatron_to_hf.py \
+    --hf_model_name_or_path <nemotron_h_omni-checkpoint> \
+    --trust_remote_code \
+    --megatron_path /path/to/qad_output/checkpoints \
+    --export_unified_hf_path /path/to/nemotron_h_omni_w4a4_hf \
+    --pp_size 4
+
+python examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/add_nvfp4_kv_scales.py \
+    /path/to/nemotron_h_omni_w4a4_hf
+```
+
+- The exporter writes the quantized language model under `language_model.`, the MTP head from the model in BF16, and copies `vision_model.`, `vision_projector.` and `mlp1.` from the source checkpoint.
+- 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`.
+- [`add_nvfp4_kv_scales.py`](add_nvfp4_kv_scales.py) is needed because the recipe's KV quantizers use a constant amax, so the export carries NVFP4 KV metadata but no `k_scale` / `v_scale` tensors. The script writes the scale that amax implies, 448 / (6 × 448) = 1/6, for each of the 8 attention layers.
+
+---
+
+### 5. Evaluation and Serving
+
+All checkpoints were served with vLLM on one 4-GPU GB300 node, TP=4 with expert parallelism. Quantized checkpoints used these engine arguments; the BF16 baseline used `--kv-cache-dtype auto`.
+
+| Setting | Value |
+| --- | --- |
+| `--kv-cache-dtype` | `nvfp4` |
+| `--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` |
+
+Benchmarks: AA-LCR (mean reward × 100, 16 repeats), SciCode (subtask accuracy), and GPQA Diamond without tools (pass@1 averaged over 8 samples, plus majority@8).
+
+## Potential further improvements
+
+- **Long reasoning:** on a separate long-reasoning benchmark (APEX shortlist, no tools), PTQ stayed close to BF16 (87.8 vs 88.4) while QAD 600 scored 83.2 and produced about 18% fewer output tokens. One hypothesis is that distilling on 32K-truncated data shortens long reasoning; a longer `--seq_length` or longer-context data would test it.
+- **Data:** the blend is mostly math and code (64% of records) and has no agentic or tool-calling data.
+- **Multimodal data:** QAD here is text-only. Image-text distillation would exercise the language model on vision-conditioned inputs as well.
diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/add_nvfp4_kv_scales.py b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/add_nvfp4_kv_scales.py
new file mode 100644
index 00000000000..0d2cde52d67
--- /dev/null
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/add_nvfp4_kv_scales.py
@@ -0,0 +1,70 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Add the NVFP4 KV-cache scales the recipe implies to an exported HF checkpoint.
+
+The recipe's K/V quantizers use a constant amax (448) and keep no calibrated amax, so the export
+writes NVFP4 KV metadata but no ``k_scale`` / ``v_scale`` tensors. The scale that amax implies is
+448 / (6 * 448) = 1/6. This writes it for every attention layer into one extra shard and adds the
+shard to the index. Re-running is a no-op.
+
+Usage: python add_nvfp4_kv_scales.py <exported_hf_dir>
+"""
+
+import json
+import os
+import re
+import sys
+
+import torch
+from safetensors.torch import save_file
+
+SHARD = "model-kv-scales.safetensors"
+# Backbone attention only: the MTP head stays BF16 and has no quantized KV cache.
+K_PROJ = re.compile(r"(.*\.backbone\.layers\.\d+\.mixer)\.k_proj\.weight$")
+
+
+def main(hf_dir: str):
+    with open(os.path.join(hf_dir, "hf_quant_config.json")) as f:
+        kv_algo = json.load(f)["quantization"].get("kv_cache_quant_algo")
+    if kv_algo != "NVFP4":
+        sys.exit(f"expected kv_cache_quant_algo NVFP4, found {kv_algo}")
+
+    index_path = os.path.join(hf_dir, "model.safetensors.index.json")
+    with open(index_path) as f:
+        index = json.load(f)
+    weight_map = index["weight_map"]
+    attention = sorted({m.group(1) for k in weight_map if (m := K_PROJ.match(k))})
+    scales = {
+        f"{prefix}.{proj}.{name}": torch.tensor(1.0 / 6.0, dtype=torch.float32)
+        for prefix in attention
+        for proj, name in (("k_proj", "k_scale"), ("v_proj", "v_scale"))
+    }
+    scales = {k: v for k, v in scales.items() if k not in weight_map}
+    if scales:
+        save_file(scales, os.path.join(hf_dir, SHARD), metadata={"format": "pt"})
+        weight_map.update(dict.fromkeys(scales, SHARD))
+        index.setdefault("metadata", {})["total_size"] = index.get("metadata", {}).get(
+            "total_size", 0
+        ) + 4 * len(scales)
+        with open(index_path, "w") as f:
+            json.dump(index, f, indent=2)
+    print(f"{len(attention)} attention layers; added {len(scales)} k/v scales")
+
+
+if __name__ == "__main__":
+    if len(sys.argv) != 2:
+        sys.exit(__doc__)
+    main(sys.argv[1])
diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/build_qad_blend.py b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/build_qad_blend.py
new file mode 100644
index 00000000000..be56a738bda
--- /dev/null
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/build_qad_blend.py
@@ -0,0 +1,129 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Materialize the tutorial's QAD blend as chat jsonl (train.jsonl / validation.jsonl).
+
+``distill.py --sft_hf_dataset`` reads one source per split and has no blend weights, so the
+weighted mixture is written out here. Each source contributes ``total * weight / sum(weights)``
+rows, taken from the start of its Hugging Face split. Rows come in three shapes:
+
+* ``messages``: used as-is.
+* ``text`` with ``<extra_id_1>User`` / ``Assistant`` markers (SFT-General): parsed into messages.
+* ``text`` without markers (SFT-Code, SFT-MATH): one assistant turn, so the whole text is trained.
+
+Rows Megatron-Bridge does not accept as text-only chat are skipped. Run inside the NeMo container.
+"""
+
+import argparse
+import json
+import os
+import random
+import re
+
+from datasets import load_dataset
+from megatron.bridge.data.conversation_processing import is_text_only_chat_example
+
+# (weight, repo, config, split, field)
+SOURCES = [
+    (5, "nvidia/Nemotron-Pretraining-SFT-v1", "Nemotron-SFT-Code", "train", "text"),
+    (20, "nvidia/Nemotron-Pretraining-SFT-v1", "Nemotron-SFT-General", "train", "text"),
+    (5, "nvidia/Nemotron-Pretraining-SFT-v1", "Nemotron-SFT-MATH", "train", "text"),
+    (10, "nvidia/Nemotron-Math-v2", "default", "high_part00", "messages"),
+    (17, "nvidia/Nemotron-SFT-Math-v3", "default", "train", "messages"),
+    (
+        15,
+        "nvidia/Nemotron-Competitive-Programming-v1",
+        "default",
+        "competitive_coding_python_part00",
+        "messages",
+    ),
+    (
+        5,
+        "nvidia/Nemotron-Competitive-Programming-v1",
+        "default",
+        "competitive_coding_cpp_part00",
+        "messages",
+    ),
+    (8, "nvidia/Nemotron-Post-Training-Dataset-v1", "default", "stem", "messages"),
+    (3, "nvidia/Nemotron-Science-v1", "default", "MCQ", "messages"),
+    (2, "nvidia/Nemotron-Science-v1", "default", "RQA", "messages"),
+    (3, "nvidia/Nemotron-SFT-Instruction-Following-Chat-v2", "default", "reasoning_on", "messages"),
+    (
+        2,
+        "nvidia/Nemotron-SFT-Instruction-Following-Chat-v2",
+        "default",
+        "reasoning_off",
+        "messages",
+    ),
+    (5, "nvidia/Nemotron-Agentic-v1", "default", "tool_calling", "messages"),
+]
+TURN = re.compile(
+    r"<extra_id_\d+>(User|Assistant|System)\n?(.*?)(?=<extra_id_\d+>(?:User|Assistant|System)|\Z)",
+    re.DOTALL,
+)
+ROLE = {"User": "user", "Assistant": "assistant", "System": "system"}
+
+
+def to_chat(row: dict, field: str) -> dict | None:
+    if field == "messages":
+        return {"messages": row["messages"]} if row.get("messages") else None
+    text = row.get("text") or ""
+    if not text:
+        return None
+    turns = [{"role": ROLE[r], "content": c.strip()} for r, c in TURN.findall(text)]
+    return {"messages": turns or [{"role": "assistant", "content": text}]}
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    parser.add_argument("--output_dir", required=True)
+    parser.add_argument("--total", type=int, default=40000, help="Rows to request across sources.")
+    parser.add_argument("--num_validation", type=int, default=1000)
+    parser.add_argument("--seed", type=int, default=1234)
+    args = parser.parse_args()
+
+    total_weight = sum(w for w, *_ in SOURCES)
+    rows = []
+    for weight, repo, config, split, field in SOURCES:
+        want = max(1, round(args.total * weight / total_weight))
+        try:
+            ds = load_dataset(repo, config, split=f"{split}[:{want * 3}]")
+        except Exception as e:  # e.g. a gated or renamed split
+            print(f"{repo} {split}: skipped ({type(e).__name__}: {e})")
+            continue
+        got, dropped = [], 0
+        for row in ds:
+            chat = to_chat(row, field)
+            if chat is None or not is_text_only_chat_example(chat):
+                dropped += 1
+                continue
+            got.append(chat)
+            if len(got) == want:
+                break
+        rows.extend(got)
+        print(f"{repo} {split}: {len(got)}/{want} rows ({dropped} dropped)")
+
+    random.Random(args.seed).shuffle(rows)
+    os.makedirs(args.output_dir, exist_ok=True)
+    splits = {"validation": rows[: args.num_validation], "train": rows[args.num_validation :]}
+    for name, chunk in splits.items():
+        path = os.path.join(args.output_dir, f"{name}.jsonl")
+        with open(path, "w") as f:
+            f.writelines(json.dumps(r) + "\n" for r in chunk)
+        print(f"wrote {path}: {len(chunk)} rows")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/megatron_bridge/tutorials/README.md b/examples/megatron_bridge/tutorials/README.md
index ced00fde2a5..6c6c6fd0628 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 with an NVFP4 KV cache → KD-only QAD on chat data → Hugging Face export → vLLM serving. QAD 600 lands within 0.5 points of BF16 on AA-LCR, SciCode and GPQA Diamond. |
 | [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. |

From 61d9d8a1edf5368ef463c3aaa57f63aacae86178 Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Mon, 12 Oct 2026 20:43:37 +0000
Subject: [PATCH 02/24] Add Megatron-Core module names to the NemotronH
 aggressive-mse recipe

The rules used Hugging Face names only, so under Megatron-Bridge attention
(linear_qkv / linear_proj), lm_head (output_layer) and the MTP block-1 experts
(mtp.layers.0.mtp_model_layer.layers.1) silently stayed BF16. List both
spellings where they differ, as the Lightning recipes do; each is inert on the
other framework. Adds a Megatron-Core-named case to the recipe test.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../nemotron_h/ptq/nvfp4-aggressive-mse.yaml  |  42 ++++++-
 modelopt_recipes/ptq.md                       |   5 +
 tests/unit/recipe/test_nemotron_h_recipe.py   | 109 +++++++++++++++++-
 3 files changed, 151 insertions(+), 5 deletions(-)

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..aeed620645c 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,21 @@ 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. Megatron-Core MTP block 0
+    # (mtp.layers.0.mtp_model_layer.layers.0) falls under the same glob.
     - quantizer_name: '*mtp.layers.0.*[kv]_bmm_quantizer'
       cfg:
         $import: fp8
@@ -120,6 +151,15 @@ quantize:
       enable: true
       cfg:
         $import: fp8
+    # Megatron-Core lm_head.
+    - 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*'
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..470b1500a2e 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,55 @@ 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``.
+    """
+    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()])
+    model.vision_projector = nn.Linear(_WIDTH, _WIDTH, bias=False)
+    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 +370,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

From 19ffa620de4100a70336c1c9ec17e2200626bcf6 Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Mon, 12 Oct 2026 20:43:37 +0000
Subject: [PATCH 03/24] Point the NemotronH omni VL tutorial at the
 nvfp4-aggressive-mse recipe

Use model_type/nemotron_h/ptq/nvfp4-aggressive-mse for PTQ: FP8 KV cast and
NVFP4 MTP block-1 experts. Drop add_nvfp4_kv_scales.py, which only served the
NVFP4 KV cast, and serve with an FP8 KV cache. The reported results come from the
earlier NVFP4-KV variant and are marked as such until they are refreshed.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 examples/megatron_bridge/README.md            |  2 +-
 .../NemotronH-Omni-W4A4-QAD/README.md         | 49 ++++++-------
 .../add_nvfp4_kv_scales.py                    | 70 -------------------
 examples/megatron_bridge/tutorials/README.md  |  2 +-
 4 files changed, 27 insertions(+), 96 deletions(-)
 delete mode 100644 examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/add_nvfp4_kv_scales.py

diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md
index 87243070506..a6428d79022 100644
--- a/examples/megatron_bridge/README.md
+++ b/examples/megatron_bridge/README.md
@@ -20,7 +20,7 @@ This directory contains examples of using Model Optimizer with the [NeMo Megatro
 >
 > 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 with an NVFP4 KV cache.
+> 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
index 82f60391ae5..4799e669e81 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -8,11 +8,14 @@ End-to-end W4A4 quantization of a NemotronH omni vision-language checkpoint (arc
 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 is quantized; the vision tower, projector and MTP head stay BF16 and are copied from the source checkpoint.
+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 ([Section 5](#5-evaluation-and-serving)); quantized ones with an NVFP4 KV cache.
+> [!IMPORTANT]
+> These results come from an earlier variant of the recipe: an NVFP4 KV-cache cast instead of FP8, and the whole MTP head in BF16. They will be refreshed with [`nvfp4-aggressive-mse`](../../../../modelopt_recipes/model_type/nemotron_h/ptq/nvfp4-aggressive-mse.yaml), the recipe the steps below use; until then they show the expected order of magnitude, not exact targets.
+
+All checkpoints were served with the same vLLM settings; quantized ones with an NVFP4 KV cache.
 
 | Model | AA-LCR | SciCode (subtask) | GPQA Diamond pass@1 (maj@8) |
 | --- | --- | --- | --- |
@@ -70,17 +73,18 @@ Records follow the `{"messages": [{"role": ..., "content": ...}, ...]}` schema a
 
 ### 2. Quantization
 
-The recipe is a fixed mixed-precision scheme; it will be published under `modelopt_recipes`, and its path goes to `--recipe` below. Counts are read from the exported checkpoint's `hf_quant_config.json`:
+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 | Tensors |
-| --- | --- | --- |
-| MoE routed experts (`up`/`down_proj`) | **NVFP4 W4A4**, block 16, MSE-selected static weight scales | 40,960 |
-| MoE shared experts | **NVFP4 W4A4** | 80 |
-| Mamba `in_proj` / `out_proj` | **FP8 W8A8** per-tensor | 80 |
-| Attention `q`/`k`/`v`/`o_proj` (8 layers) | **FP8 W8A8** per-tensor | 32 |
-| `lm_head` | **FP8 W8A8** | 1 |
-| KV cache | **NVFP4** cast, constant amax 448 | — |
-| Vision tower, projector, MTP head, routers, embeddings | BF16 | — |
+| 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 takes **16 minutes on one 4-GPU GB300 node** with 512 calibration samples.
 
@@ -92,7 +96,7 @@ PTQ takes **16 minutes on one 4-GPU GB300 node** with 512 calibration samples.
 srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/quantize.py \
     --hf_model_name_or_path <nemotron_h_omni-checkpoint> \
     --trust_remote_code \
-    --recipe /path/to/recipe.yaml \
+    --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 \
@@ -103,8 +107,8 @@ srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/quantize.py \
 
 - `--tp_size 1` is **required**: static NVFP4 weight scales are calibrated per full weight, so the recipe needs 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 MoE experts are built as grouped GEMMs (`TEGroupedMLP`), which is what the recipe's `mlp.experts.linear_fc*` rules match. The log reports `moe_grouped_gemm: true`; do not pass `--no_moe_grouped_gemm`.
-- The MTP head is built and stays BF16 (the recipe disables `language_model.mtp.*`).
+- The MoE experts are built as grouped GEMMs (`TEGroupedMLP`). The log reports `moe_grouped_gemm: true`; do not pass `--no_moe_grouped_gemm`.
+- The MTP head is built. In Megatron-Core its MoE block is `language_model.mtp.layers.0.mtp_model_layer.layers.1`, which the recipe quantizes like the backbone experts; the rest of the head stays BF16.
 
 </details>
 
@@ -145,7 +149,7 @@ srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/distill.py \
 - `--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.
-- The MTP head is built but receives no loss under KD-only training (`mtp_1 loss` and `mtp_2 loss` log as 0), so it stays bit-identical to the BF16 source.
+- The MTP head is built but receives no loss under KD-only training (`mtp_1 loss` and `mtp_2 loss` log as 0), so its weights stay as PTQ left them.
 - Each saved iteration is about 1.6 TB; size `--save_interval` and `--checkpoint_keep_last` to your storage.
 
 </details>
@@ -154,7 +158,7 @@ srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/distill.py \
 
 ### 4. Export
 
-Export the QAD checkpoint to a Hugging Face checkpoint on one 4-GPU node (pipeline parallel 4, about 26 minutes), then add the NVFP4 KV-cache scales:
+Export the QAD checkpoint to a Hugging Face checkpoint on one 4-GPU node (pipeline parallel 4, about 26 minutes):
 
 ```bash
 # SBATCH --nodes=1 --ntasks-per-node=4 --gpus-per-node=4
@@ -164,24 +168,21 @@ srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/export_quantize
     --megatron_path /path/to/qad_output/checkpoints \
     --export_unified_hf_path /path/to/nemotron_h_omni_w4a4_hf \
     --pp_size 4
-
-python examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/add_nvfp4_kv_scales.py \
-    /path/to/nemotron_h_omni_w4a4_hf
 ```
 
-- The exporter writes the quantized language model under `language_model.`, the MTP head from the model in BF16, and copies `vision_model.`, `vision_projector.` and `mlp1.` from the source checkpoint.
+- 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.
 - 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`.
-- [`add_nvfp4_kv_scales.py`](add_nvfp4_kv_scales.py) is needed because the recipe's KV quantizers use a constant amax, so the export carries NVFP4 KV metadata but no `k_scale` / `v_scale` tensors. The script writes the scale that amax implies, 448 / (6 × 448) = 1/6, for each of the 8 attention layers.
+- 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
 
-All checkpoints were served with vLLM on one 4-GPU GB300 node, TP=4 with expert parallelism. Quantized checkpoints used these engine arguments; the BF16 baseline used `--kv-cache-dtype auto`.
+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` | `nvfp4` |
+| `--kv-cache-dtype` | `fp8` |
 | `--max-model-len` | 262144 |
 | `--mamba-ssm-cache-dtype` | `float16` |
 | `--mamba-conv-cache-dtype` | `float32` |
diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/add_nvfp4_kv_scales.py b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/add_nvfp4_kv_scales.py
deleted file mode 100644
index 0d2cde52d67..00000000000
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/add_nvfp4_kv_scales.py
+++ /dev/null
@@ -1,70 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Add the NVFP4 KV-cache scales the recipe implies to an exported HF checkpoint.
-
-The recipe's K/V quantizers use a constant amax (448) and keep no calibrated amax, so the export
-writes NVFP4 KV metadata but no ``k_scale`` / ``v_scale`` tensors. The scale that amax implies is
-448 / (6 * 448) = 1/6. This writes it for every attention layer into one extra shard and adds the
-shard to the index. Re-running is a no-op.
-
-Usage: python add_nvfp4_kv_scales.py <exported_hf_dir>
-"""
-
-import json
-import os
-import re
-import sys
-
-import torch
-from safetensors.torch import save_file
-
-SHARD = "model-kv-scales.safetensors"
-# Backbone attention only: the MTP head stays BF16 and has no quantized KV cache.
-K_PROJ = re.compile(r"(.*\.backbone\.layers\.\d+\.mixer)\.k_proj\.weight$")
-
-
-def main(hf_dir: str):
-    with open(os.path.join(hf_dir, "hf_quant_config.json")) as f:
-        kv_algo = json.load(f)["quantization"].get("kv_cache_quant_algo")
-    if kv_algo != "NVFP4":
-        sys.exit(f"expected kv_cache_quant_algo NVFP4, found {kv_algo}")
-
-    index_path = os.path.join(hf_dir, "model.safetensors.index.json")
-    with open(index_path) as f:
-        index = json.load(f)
-    weight_map = index["weight_map"]
-    attention = sorted({m.group(1) for k in weight_map if (m := K_PROJ.match(k))})
-    scales = {
-        f"{prefix}.{proj}.{name}": torch.tensor(1.0 / 6.0, dtype=torch.float32)
-        for prefix in attention
-        for proj, name in (("k_proj", "k_scale"), ("v_proj", "v_scale"))
-    }
-    scales = {k: v for k, v in scales.items() if k not in weight_map}
-    if scales:
-        save_file(scales, os.path.join(hf_dir, SHARD), metadata={"format": "pt"})
-        weight_map.update(dict.fromkeys(scales, SHARD))
-        index.setdefault("metadata", {})["total_size"] = index.get("metadata", {}).get(
-            "total_size", 0
-        ) + 4 * len(scales)
-        with open(index_path, "w") as f:
-            json.dump(index, f, indent=2)
-    print(f"{len(attention)} attention layers; added {len(scales)} k/v scales")
-
-
-if __name__ == "__main__":
-    if len(sys.argv) != 2:
-        sys.exit(__doc__)
-    main(sys.argv[1])
diff --git a/examples/megatron_bridge/tutorials/README.md b/examples/megatron_bridge/tutorials/README.md
index 6c6c6fd0628..454403083e8 100644
--- a/examples/megatron_bridge/tutorials/README.md
+++ b/examples/megatron_bridge/tutorials/README.md
@@ -8,5 +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 with an NVFP4 KV cache → KD-only QAD on chat data → Hugging Face export → vLLM serving. QAD 600 lands within 0.5 points of BF16 on AA-LCR, SciCode and GPQA Diamond. |
+| [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. |

From 8929e7c6550d7818afa3246521d27dab899985c7 Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Mon, 12 Oct 2026 20:49:13 +0000
Subject: [PATCH 04/24] Correct the reason the NemotronH omni VL tutorial needs
 TP=1

Static NVFP4 weight scales are per-shard under tensor parallelism, but the Megatron
sharded checkpoint stores them as replicated, so calibration raises for TP > 1
(_check_nvfp4_static_tp_supported).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
index 4799e669e81..8fec55fc886 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -105,7 +105,7 @@ srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/quantize.py \
     --skip_generate
 ```
 
-- `--tp_size 1` is **required**: static NVFP4 weight scales are calibrated per full weight, so the recipe needs TP=1. Expert parallelism spreads the experts instead.
+- `--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 MoE experts are built as grouped GEMMs (`TEGroupedMLP`). The log reports `moe_grouped_gemm: true`; do not pass `--no_moe_grouped_gemm`.
 - The MTP head is built. In Megatron-Core its MoE block is `language_model.mtp.layers.0.mtp_model_layer.layers.1`, which the recipe quantizes like the backbone experts; the rest of the head stays BF16.

From eff8d28e2c46928304ed8852f04b7a3a448df471 Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Mon, 12 Oct 2026 20:54:49 +0000
Subject: [PATCH 05/24] Clarify the MTP KV comment in the NemotronH
 aggressive-mse recipe

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../model_type/nemotron_h/ptq/nvfp4-aggressive-mse.yaml      | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

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 aeed620645c..73987b08155 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
@@ -135,8 +135,9 @@ quantize:
     - 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. Megatron-Core MTP block 0
-    # (mtp.layers.0.mtp_model_layer.layers.0) falls under the same glob.
+    # 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

From 7d62553988c1c8af87e60e22feefca571b2eb9ec Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Mon, 12 Oct 2026 21:29:13 +0000
Subject: [PATCH 06/24] Make build_qad_blend.py the script that built the QAD
 blend

Replace the rewrite with a cleaned-up copy of the script that produced the blend:
stream each split in order, stop a source after 10x rejected rows, skip a source
that fails to load, drop empty <extra_id_> turns, cache each source under parts/,
and hold out min(1000, len // 20) rows. Paths become flags. Drop the README's
unverified reason for the two empty sources.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../NemotronH-Omni-W4A4-QAD/README.md         |  2 +-
 .../build_qad_blend.py                        | 71 +++++++++++++------
 2 files changed, 50 insertions(+), 23 deletions(-)

diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
index 8fec55fc886..5a11e5eecdf 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -60,7 +60,7 @@ QAD distills on text chat data, read directly from jsonl by `distill.py --sft_hf
 | Nemotron-Science-v1 / RQA | 2 | 800 |
 | Nemotron-SFT-Instruction-Following-Chat-v2 / reasoning_off | 2 | 800 |
 
-The script also lists `Nemotron-SFT-Instruction-Following-Chat-v2 / reasoning_on` and `Nemotron-Agentic-v1 / tool_calling`; in our run they contributed no records (tool-calling rows are not text-only chat), so the blend has no agentic data. After a seeded shuffle, 1,000 records are held out for validation and 35,800 are used for training.
+The script also lists `Nemotron-SFT-Instruction-Following-Chat-v2 / reasoning_on` and `Nemotron-Agentic-v1 / tool_calling`; in our run they contributed no records, so the blend has no agentic data. After a seeded shuffle, 1,000 records are held out for validation and 35,800 are used for training.
 
 ```bash
 python examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/build_qad_blend.py \
diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/build_qad_blend.py b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/build_qad_blend.py
index be56a738bda..6dc80211f5d 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/build_qad_blend.py
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/build_qad_blend.py
@@ -13,17 +13,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 """Materialize the tutorial's QAD blend as chat jsonl (train.jsonl / validation.jsonl).
 
 ``distill.py --sft_hf_dataset`` reads one source per split and has no blend weights, so the
 weighted mixture is written out here. Each source contributes ``total * weight / sum(weights)``
-rows, taken from the start of its Hugging Face split. Rows come in three shapes:
+rows: the first ones of its Hugging Face split, streamed in order, that Megatron-Bridge accepts
+as text-only chat. A source is skipped if it fails to load or yields too few usable rows. Rows come
+in three shapes:
 
 * ``messages``: used as-is.
-* ``text`` with ``<extra_id_1>User`` / ``Assistant`` markers (SFT-General): parsed into messages.
+* ``text`` with ``<extra_id_1>User`` / ``Assistant`` markers (SFT-General): parsed into messages,
+  dropping empty turns.
 * ``text`` without markers (SFT-Code, SFT-MATH): one assistant turn, so the whole text is trained.
 
-Rows Megatron-Bridge does not accept as text-only chat are skipped. Run inside the NeMo container.
+Each source is cached under ``<output_dir>/parts/`` so an interrupted run resumes. Run inside the
+NeMo container.
 """
 
 import argparse
@@ -82,42 +87,64 @@ def to_chat(row: dict, field: str) -> dict | None:
     text = row.get("text") or ""
     if not text:
         return None
-    turns = [{"role": ROLE[r], "content": c.strip()} for r, c in TURN.findall(text)]
-    return {"messages": turns or [{"role": "assistant", "content": text}]}
+    if "<extra_id_" in text:
+        turns = [
+            {"role": ROLE[r], "content": c.strip()} for r, c in TURN.findall(text) if c.strip()
+        ]
+        if turns:
+            return {"messages": turns}
+    return {"messages": [{"role": "assistant", "content": text}]}
+
+
+def take(repo: str, config: str, split: str, field: str, want: int) -> tuple[list[dict], int]:
+    """Stream the split and keep the first ``want`` usable rows; give up after ``10 * want`` rejects."""
+    got, dropped = [], 0
+    try:
+        for row in load_dataset(repo, config, split=split, streaming=True):
+            chat = to_chat(row, field)
+            if chat is None or not is_text_only_chat_example(chat):
+                dropped += 1
+                if dropped > want * 10:
+                    break
+                continue
+            got.append(chat)
+            if len(got) >= want:
+                break
+    except Exception as e:  # e.g. a gated split, or a schema the streaming reader cannot cast
+        print(f"{repo} {split}: failed ({type(e).__name__}: {str(e)[:80]})", flush=True)
+    return got, dropped
 
 
 def main():
     parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
     parser.add_argument("--output_dir", required=True)
     parser.add_argument("--total", type=int, default=40000, help="Rows to request across sources.")
-    parser.add_argument("--num_validation", type=int, default=1000)
     parser.add_argument("--seed", type=int, default=1234)
     args = parser.parse_args()
 
+    parts = os.path.join(args.output_dir, "parts")
+    os.makedirs(parts, exist_ok=True)
     total_weight = sum(w for w, *_ in SOURCES)
     rows = []
     for weight, repo, config, split, field in SOURCES:
         want = max(1, round(args.total * weight / total_weight))
-        try:
-            ds = load_dataset(repo, config, split=f"{split}[:{want * 3}]")
-        except Exception as e:  # e.g. a gated or renamed split
-            print(f"{repo} {split}: skipped ({type(e).__name__}: {e})")
-            continue
-        got, dropped = [], 0
-        for row in ds:
-            chat = to_chat(row, field)
-            if chat is None or not is_text_only_chat_example(chat):
-                dropped += 1
+        part = os.path.join(parts, f"{repo.split('/')[-1]}__{config}__{split}.jsonl")
+        if os.path.exists(part):
+            with open(part) as f:
+                got = [json.loads(line) for line in f]
+            if len(got) >= want:
+                rows.extend(got[:want])
+                print(f"{repo} {split}: {want} rows (cached)", flush=True)
                 continue
-            got.append(chat)
-            if len(got) == want:
-                break
+        got, dropped = take(repo, config, split, field, want)
+        with open(part, "w") as f:
+            f.writelines(json.dumps(r) + "\n" for r in got)
         rows.extend(got)
-        print(f"{repo} {split}: {len(got)}/{want} rows ({dropped} dropped)")
+        print(f"{repo} {split}: {len(got)}/{want} rows ({dropped} dropped)", flush=True)
 
     random.Random(args.seed).shuffle(rows)
-    os.makedirs(args.output_dir, exist_ok=True)
-    splits = {"validation": rows[: args.num_validation], "train": rows[args.num_validation :]}
+    num_validation = min(1000, max(1, len(rows) // 20))
+    splits = {"validation": rows[:num_validation], "train": rows[num_validation:]}
     for name, chunk in splits.items():
         path = os.path.join(args.output_dir, f"{name}.jsonl")
         with open(path, "w") as f:

From df8ea80a11a70ac0dba6cb40c592cd59b5b09ef4 Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Mon, 12 Oct 2026 21:49:12 +0000
Subject: [PATCH 07/24] Shorten the data preparation step of the NemotronH omni
 VL tutorial

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../NemotronH-Omni-W4A4-QAD/README.md         | 22 ++-----------------
 1 file changed, 2 insertions(+), 20 deletions(-)

diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
index 5a11e5eecdf..b1545ddc432 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -44,31 +44,13 @@ Every command below runs with `srun`, one task per GPU, and sets `RANK` / `WORLD
 
 ### 1. Data Preparation
 
-QAD distills on text chat data, read directly from jsonl by `distill.py --sft_hf_dataset`. [`build_qad_blend.py`](build_qad_blend.py) materializes a 36,800-record sample of public Nemotron post-training datasets, 400 records per weight point:
-
-| Dataset / split | Weight | Records |
-| --- | --- | --- |
-| Nemotron-Pretraining-SFT-v1 / Nemotron-SFT-General | 20 | 8,000 |
-| Nemotron-SFT-Math-v3 / train | 17 | 6,800 |
-| Nemotron-Competitive-Programming-v1 / python_part00 | 15 | 6,000 |
-| Nemotron-Math-v2 / high_part00 | 10 | 4,000 |
-| Nemotron-Post-Training-Dataset-v1 / stem | 8 | 3,200 |
-| Nemotron-Pretraining-SFT-v1 / Nemotron-SFT-Code | 5 | 2,000 |
-| Nemotron-Pretraining-SFT-v1 / Nemotron-SFT-MATH | 5 | 2,000 |
-| Nemotron-Competitive-Programming-v1 / cpp_part00 | 5 | 2,000 |
-| Nemotron-Science-v1 / MCQ | 3 | 1,200 |
-| Nemotron-Science-v1 / RQA | 2 | 800 |
-| Nemotron-SFT-Instruction-Following-Chat-v2 / reasoning_off | 2 | 800 |
-
-The script also lists `Nemotron-SFT-Instruction-Following-Chat-v2 / reasoning_on` and `Nemotron-Agentic-v1 / tool_calling`; in our run they contributed no records, so the blend has no agentic data. After a seeded shuffle, 1,000 records are held out for validation and 35,800 are used for training.
+Generate the QAD data with [`build_qad_blend.py`](build_qad_blend.py):
 
 ```bash
 python examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/build_qad_blend.py \
     --output_dir /path/to/qad_blend
 ```
 
-Records follow the `{"messages": [{"role": ..., "content": ...}, ...]}` schema and are rendered with the model's own chat template during QAD. Records longer than the 32K sequence length used below are truncated.
-
 ---
 
 ### 2. Quantization
@@ -195,5 +177,5 @@ Benchmarks: AA-LCR (mean reward × 100, 16 repeats), SciCode (subtask accuracy),
 ## Potential further improvements
 
 - **Long reasoning:** on a separate long-reasoning benchmark (APEX shortlist, no tools), PTQ stayed close to BF16 (87.8 vs 88.4) while QAD 600 scored 83.2 and produced about 18% fewer output tokens. One hypothesis is that distilling on 32K-truncated data shortens long reasoning; a longer `--seq_length` or longer-context data would test it.
-- **Data:** the blend is mostly math and code (64% of records) and has no agentic or tool-calling data.
+- **Data:** the blend is mostly math and code (64% of records).
 - **Multimodal data:** QAD here is text-only. Image-text distillation would exercise the language model on vision-conditioned inputs as well.

From 49425bf7c10449034f5508a9742009921b4ea714 Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Mon, 12 Oct 2026 22:05:41 +0000
Subject: [PATCH 08/24] Shorten the launch note of the NemotronH omni VL
 tutorial

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
index b1545ddc432..cee483ae36b 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -40,7 +40,7 @@ The recipe is already close to lossless after PTQ: every PTQ score is within 1.4
 
 **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` at [`d23c1f98`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/commit/d23c1f981596c227b13d98c95b09afb865d80258) (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.
 
-Every command below runs with `srun`, one task per GPU, and sets `RANK` / `WORLD_SIZE` / `LOCAL_RANK` from `SLURM_PROCID` / `SLURM_NTASKS` / `SLURM_LOCALID` in the `srun` body. Pass `--trust_remote_code`, since the model ships its own modeling code.
+Commands run under `srun`, one task per GPU, with `RANK` / `WORLD_SIZE` / `LOCAL_RANK` set from `SLURM_PROCID` / `SLURM_NTASKS` / `SLURM_LOCALID`.
 
 ### 1. Data Preparation
 

From 34ca4b60863efb2a2878994e2cad09d88865c079 Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Mon, 12 Oct 2026 22:07:12 +0000
Subject: [PATCH 09/24] Drop the results caveat from the NemotronH omni VL
 tutorial

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../tutorials/NemotronH-Omni-W4A4-QAD/README.md                | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
index cee483ae36b..f2df55bdcf7 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -12,9 +12,6 @@ Only the language model (including its MTP head) is quantized; the vision tower
 
 ## Results
 
-> [!IMPORTANT]
-> These results come from an earlier variant of the recipe: an NVFP4 KV-cache cast instead of FP8, and the whole MTP head in BF16. They will be refreshed with [`nvfp4-aggressive-mse`](../../../../modelopt_recipes/model_type/nemotron_h/ptq/nvfp4-aggressive-mse.yaml), the recipe the steps below use; until then they show the expected order of magnitude, not exact targets.
-
 All checkpoints were served with the same vLLM settings; quantized ones with an NVFP4 KV cache.
 
 | Model | AA-LCR | SciCode (subtask) | GPQA Diamond pass@1 (maj@8) |

From 9e3ecbb6b392584d65f8a938b476f061c043653e Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Mon, 12 Oct 2026 22:08:37 +0000
Subject: [PATCH 10/24] Report the full text panel for BF16 and QAD 600 in the
 NemotronH omni VL tutorial

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../NemotronH-Omni-W4A4-QAD/README.md         | 29 +++++++++----------
 1 file changed, 14 insertions(+), 15 deletions(-)

diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
index f2df55bdcf7..facdc0a2d9c 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -14,19 +14,20 @@ Only the language model (including its MTP head) is quantized; the vision tower
 
 All checkpoints were served with the same vLLM settings; quantized ones with an NVFP4 KV cache.
 
-| Model | AA-LCR | SciCode (subtask) | GPQA Diamond pass@1 (maj@8) |
-| --- | --- | --- | --- |
-| **BF16** (teacher) | 63.31 | 46.48 | 89.52 (90.66) |
-| **W4A4 PTQ** (QAD student) | 62.00 | 46.92 | 89.27 (90.66) |
-| ↳ QAD 200 iters | 61.69 | 43.62 | 88.70 (89.65) |
-| ↳ QAD 300 iters | 62.56 | 45.53 | 88.70 (88.13) |
-| ↳ QAD 400 iters | 64.06 | 45.36 | 88.89 (90.40) |
-| ↳ QAD 500 iters | 62.94 | 46.22 | 89.84 (90.66) |
-| **↳ QAD 600 iters** | **63.75** | **46.35** | **89.77 (91.67)** |
-
-Sample counts: AA-LCR 1,600 (100 questions × 16), SciCode 520, GPQA Diamond 1,584 (198 questions × 8).
-
-The recipe is already close to lossless after PTQ: every PTQ score is within 1.4 points of BF16. The first few hundred QAD iterations dip before recovering, and QAD 600 lands within 0.5 points of BF16 on all three benchmarks. The exported W4A4 checkpoint is **74 GB**.
+| Benchmark | BF16 (teacher) | W4A4 QAD 600 iters |
+| --- | --- | --- |
+| AA-LCR | 68.63 | 69.31 |
+| AA-Omniscience | 26.21 | 25.60 |
+| AA-Omniscience (non-hallucination) | 70.03 | 69.13 |
+| APEX shortlist (no tools) | 88.36 | 83.24 |
+| GPQA Diamond | 89.52 | 89.58 |
+| HLE (no tools) | 31.79 | 30.26 |
+| MMLU-ProX | 76.91 | 80.30 |
+| SciCode (subtask) | 47.13 | 46.43 |
+| SciCode (problem) | 15.10 | 14.62 |
+| τ³-bench banking | 18.88 | 18.04 |
+
+BF16 is the mean of two runs. The exported W4A4 checkpoint is **74 GB**.
 
 > [!NOTE]
 > The numbers above come from runs whose code was functionally equivalent to the commands below but predates them: the Omni export mapping and chat-data distillation were local patches then, and the student was built without its MTP head (`--mtp_num_layers 0`). This code path was checked end to end with a short run (a 32-sample PTQ, 20 QAD iterations and the export); QAD keeps the MTP head unchanged either way (see [Section 3](#3-quantization-aware-distillation-qad)).
@@ -169,8 +170,6 @@ Serve with vLLM on one 4-GPU GB300 node, TP=4 with expert parallelism. Quantized
 | Reasoning parser | `nemotron_v3` |
 | Chat template | the checkpoint's `chat_template.jinja` |
 
-Benchmarks: AA-LCR (mean reward × 100, 16 repeats), SciCode (subtask accuracy), and GPQA Diamond without tools (pass@1 averaged over 8 samples, plus majority@8).
-
 ## Potential further improvements
 
 - **Long reasoning:** on a separate long-reasoning benchmark (APEX shortlist, no tools), PTQ stayed close to BF16 (87.8 vs 88.4) while QAD 600 scored 83.2 and produced about 18% fewer output tokens. One hypothesis is that distilling on 32K-truncated data shortens long reasoning; a longer `--seq_length` or longer-context data would test it.

From fcc75a3ad6cf1029ff28a160578ba2fae847c999 Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Mon, 12 Oct 2026 22:08:48 +0000
Subject: [PATCH 11/24] Use Megatron-Bridge main in the NemotronH omni VL
 tutorial environment

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
index facdc0a2d9c..6edd0ca098c 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -36,7 +36,7 @@ BF16 is the mean of two runs. The exported W4A4 checkpoint is **74 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` at [`d23c1f98`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/commit/d23c1f981596c227b13d98c95b09afb865d80258) (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.
+**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.
 
 Commands run under `srun`, one task per GPU, with `RANK` / `WORLD_SIZE` / `LOCAL_RANK` set from `SLURM_PROCID` / `SLURM_NTASKS` / `SLURM_LOCALID`.
 

From 40f4d76aee5103ab3efe06e702eec26e51320e1c Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Mon, 12 Oct 2026 22:10:09 +0000
Subject: [PATCH 12/24] Drop the PTQ wall time from the NemotronH omni VL
 tutorial

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
index 6edd0ca098c..c7e90109e8e 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -66,7 +66,7 @@ The recipe is [`model_type/nemotron_h/ptq/nvfp4-aggressive-mse`](../../../../mod
 | MTP block-0 KV cache | **FP8**, calibrated |
 | Vision tower, projector, the rest of the MTP head, latent MoE, routers, embeddings | BF16 |
 
-PTQ takes **16 minutes on one 4-GPU GB300 node** with 512 calibration samples.
+PTQ runs on **one 4-GPU GB300 node** with 512 calibration samples.
 
 <details>
 <summary>PTQ command (click to expand)</summary>

From 0b789aa51ab31ac99b4614e896598dbf85517f3f Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Mon, 12 Oct 2026 22:11:08 +0000
Subject: [PATCH 13/24] Drop the grouped-GEMM note from the NemotronH omni VL
 tutorial

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md  | 1 -
 1 file changed, 1 deletion(-)

diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
index c7e90109e8e..1edbe67432d 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -87,7 +87,6 @@ srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/quantize.py \
 
 - `--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 MoE experts are built as grouped GEMMs (`TEGroupedMLP`). The log reports `moe_grouped_gemm: true`; do not pass `--no_moe_grouped_gemm`.
 - The MTP head is built. In Megatron-Core its MoE block is `language_model.mtp.layers.0.mtp_model_layer.layers.1`, which the recipe quantizes like the backbone experts; the rest of the head stays BF16.
 
 </details>

From 6c9da7189f94d77747845e62e83fe71c2ce8b37b Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Mon, 12 Oct 2026 22:19:48 +0000
Subject: [PATCH 14/24] Clarify the MTP head's role in the NemotronH omni VL
 tutorial

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../tutorials/NemotronH-Omni-W4A4-QAD/README.md              | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
index 1edbe67432d..406f01b265d 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -87,7 +87,7 @@ srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/quantize.py \
 
 - `--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 built. In Megatron-Core its MoE block is `language_model.mtp.layers.0.mtp_model_layer.layers.1`, which the recipe quantizes like the backbone experts; the rest of the head stays BF16.
+- 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.
 
 </details>
 
@@ -128,7 +128,7 @@ srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/distill.py \
 - `--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.
-- The MTP head is built but receives no loss under KD-only training (`mtp_1 loss` and `mtp_2 loss` log as 0), so its weights stay as PTQ left them.
+- QAD does not train the MTP head, so the exported MTP head keeps its PTQ weights.
 - Each saved iteration is about 1.6 TB; size `--save_interval` and `--checkpoint_keep_last` to your storage.
 
 </details>
@@ -150,6 +150,7 @@ srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/export_quantize
 ```
 
 - 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 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.
 

From f901d2c8d226c2c816efaab46435a8f240f5247c Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Mon, 12 Oct 2026 22:20:17 +0000
Subject: [PATCH 15/24] Drop the QAD wall time from the NemotronH omni VL
 tutorial

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
index 406f01b265d..8c3b6fedec1 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -97,7 +97,7 @@ srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/quantize.py \
 
 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. 600 iterations took **7.7 hours** (about 44 s/iteration at steady state). The validation KD loss fell from 1.36e-2 at iteration 100 to 1.01e-2 at iteration 600.
+We used **8 nodes × 4 GB300 (32 GPUs)** with context parallelism 4 and expert parallelism 32. The validation KD loss fell from 1.36e-2 at iteration 100 to 1.01e-2 at iteration 600.
 
 <details>
 <summary>QAD command (click to expand)</summary>

From 55d52afca5432a9dfe71bdf5de5e9fd27e32dadd Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Mon, 12 Oct 2026 22:22:32 +0000
Subject: [PATCH 16/24] Drop the export wall time from the NemotronH omni VL
 tutorial

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
index 8c3b6fedec1..5b04920b129 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -137,7 +137,7 @@ srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/distill.py \
 
 ### 4. Export
 
-Export the QAD checkpoint to a Hugging Face checkpoint on one 4-GPU node (pipeline parallel 4, about 26 minutes):
+Export the QAD checkpoint to a Hugging Face checkpoint on one 4-GPU node (pipeline parallel 4):
 
 ```bash
 # SBATCH --nodes=1 --ntasks-per-node=4 --gpus-per-node=4

From 5a2246bb72dd509006df83f5550483b63874b44c Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Mon, 12 Oct 2026 22:31:08 +0000
Subject: [PATCH 17/24] Drop the further-improvements section from the
 NemotronH omni VL tutorial

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../tutorials/NemotronH-Omni-W4A4-QAD/README.md             | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
index 5b04920b129..ad3aa082c49 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -169,9 +169,3 @@ Serve with vLLM on one 4-GPU GB300 node, TP=4 with expert parallelism. Quantized
 | Mamba cache rounding | stochastic rounding, 5 Philox rounds |
 | Reasoning parser | `nemotron_v3` |
 | Chat template | the checkpoint's `chat_template.jinja` |
-
-## Potential further improvements
-
-- **Long reasoning:** on a separate long-reasoning benchmark (APEX shortlist, no tools), PTQ stayed close to BF16 (87.8 vs 88.4) while QAD 600 scored 83.2 and produced about 18% fewer output tokens. One hypothesis is that distilling on 32K-truncated data shortens long reasoning; a longer `--seq_length` or longer-context data would test it.
-- **Data:** the blend is mostly math and code (64% of records).
-- **Multimodal data:** QAD here is text-only. Image-text distillation would exercise the language model on vision-conditioned inputs as well.

From 1a12542c64f6ca3e9aaeac5fc44d4b5f3b60969a Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Tue, 13 Oct 2026 17:50:50 +0000
Subject: [PATCH 18/24] Remove the results section from the NemotronH omni VL
 tutorial

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../NemotronH-Omni-W4A4-QAD/README.md         | 24 +------------------
 1 file changed, 1 insertion(+), 23 deletions(-)

diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
index ad3aa082c49..3031eedecca 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -6,32 +6,10 @@ End-to-end W4A4 quantization of a NemotronH omni vision-language checkpoint (arc
 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
+5. **[Evaluation and Serving](#5-evaluation-and-serving)**: the vLLM settings
 
 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; quantized ones with an NVFP4 KV cache.
-
-| Benchmark | BF16 (teacher) | W4A4 QAD 600 iters |
-| --- | --- | --- |
-| AA-LCR | 68.63 | 69.31 |
-| AA-Omniscience | 26.21 | 25.60 |
-| AA-Omniscience (non-hallucination) | 70.03 | 69.13 |
-| APEX shortlist (no tools) | 88.36 | 83.24 |
-| GPQA Diamond | 89.52 | 89.58 |
-| HLE (no tools) | 31.79 | 30.26 |
-| MMLU-ProX | 76.91 | 80.30 |
-| SciCode (subtask) | 47.13 | 46.43 |
-| SciCode (problem) | 15.10 | 14.62 |
-| τ³-bench banking | 18.88 | 18.04 |
-
-BF16 is the mean of two runs. The exported W4A4 checkpoint is **74 GB**.
-
-> [!NOTE]
-> The numbers above come from runs whose code was functionally equivalent to the commands below but predates them: the Omni export mapping and chat-data distillation were local patches then, and the student was built without its MTP head (`--mtp_num_layers 0`). This code path was checked end to end with a short run (a 32-sample PTQ, 20 QAD iterations and the export); QAD keeps the MTP head unchanged either way (see [Section 3](#3-quantization-aware-distillation-qad)).
-
 ---
 
 ## Steps to Reproduce

From bba811a9c22cea5261564ad3a374c4884bd3484a Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Tue, 13 Oct 2026 17:52:55 +0000
Subject: [PATCH 19/24] Drop the validation KD loss sentence from the NemotronH
 omni VL tutorial

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
index 3031eedecca..28bfc0cddab 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -75,7 +75,7 @@ srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/quantize.py \
 
 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. The validation KD loss fell from 1.36e-2 at iteration 100 to 1.01e-2 at iteration 600.
+We used **8 nodes × 4 GB300 (32 GPUs)** with context parallelism 4 and expert parallelism 32.
 
 <details>
 <summary>QAD command (click to expand)</summary>

From c5342347a349a5421e83e39df04282e7003b9d7a Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Tue, 13 Oct 2026 18:31:07 +0000
Subject: [PATCH 20/24] Make the QAD blend builder fail closed and resume
 safely

A source that fails to load or comes up short now stops the run unless
--allow_partial is given, and the realized per-source counts go to
blend_summary.json. Cached parts and the output splits are written
atomically, and an unreadable part is rebuilt. The two sources that never
yielded usable rows are dropped and --total defaults to 36800, which
reproduces the blend used for QAD byte for byte.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../build_qad_blend.py                        | 147 +++++++++-----
 tests/unit/examples/test_build_qad_blend.py   | 182 ++++++++++++++++++
 2 files changed, 280 insertions(+), 49 deletions(-)
 create mode 100644 tests/unit/examples/test_build_qad_blend.py

diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/build_qad_blend.py b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/build_qad_blend.py
index 6dc80211f5d..388f0700f1b 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/build_qad_blend.py
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/build_qad_blend.py
@@ -19,16 +19,19 @@
 ``distill.py --sft_hf_dataset`` reads one source per split and has no blend weights, so the
 weighted mixture is written out here. Each source contributes ``total * weight / sum(weights)``
 rows: the first ones of its Hugging Face split, streamed in order, that Megatron-Bridge accepts
-as text-only chat. A source is skipped if it fails to load or yields too few usable rows. Rows come
-in three shapes:
+as text-only chat. Rows come in three shapes:
 
 * ``messages``: used as-is.
 * ``text`` with ``<extra_id_1>User`` / ``Assistant`` markers (SFT-General): parsed into messages,
   dropping empty turns.
 * ``text`` without markers (SFT-Code, SFT-MATH): one assistant turn, so the whole text is trained.
 
-Each source is cached under ``<output_dir>/parts/`` so an interrupted run resumes. Run inside the
-NeMo container.
+A source that fails to load or yields fewer rows than its share stops the run, so the mixture
+cannot change silently; ``--allow_partial`` keeps what was collected instead. Either way the
+realized row count of every source is written to ``blend_summary.json``.
+
+Each source is cached under ``<output_dir>/parts/``, written atomically, so an interrupted run
+resumes and an unreadable part is rebuilt. Run inside the NeMo container.
 """
 
 import argparse
@@ -37,9 +40,6 @@
 import random
 import re
 
-from datasets import load_dataset
-from megatron.bridge.data.conversation_processing import is_text_only_chat_example
-
 # (weight, repo, config, split, field)
 SOURCES = [
     (5, "nvidia/Nemotron-Pretraining-SFT-v1", "Nemotron-SFT-Code", "train", "text"),
@@ -64,7 +64,6 @@
     (8, "nvidia/Nemotron-Post-Training-Dataset-v1", "default", "stem", "messages"),
     (3, "nvidia/Nemotron-Science-v1", "default", "MCQ", "messages"),
     (2, "nvidia/Nemotron-Science-v1", "default", "RQA", "messages"),
-    (3, "nvidia/Nemotron-SFT-Instruction-Following-Chat-v2", "default", "reasoning_on", "messages"),
     (
         2,
         "nvidia/Nemotron-SFT-Instruction-Following-Chat-v2",
@@ -72,7 +71,6 @@
         "reasoning_off",
         "messages",
     ),
-    (5, "nvidia/Nemotron-Agentic-v1", "default", "tool_calling", "messages"),
 ]
 TURN = re.compile(
     r"<extra_id_\d+>(User|Assistant|System)\n?(.*?)(?=<extra_id_\d+>(?:User|Assistant|System)|\Z)",
@@ -81,6 +79,20 @@
 ROLE = {"User": "user", "Assistant": "assistant", "System": "system"}
 
 
+def stream(repo: str, config: str, split: str):
+    """Yield the rows of one Hugging Face split, in order, without downloading it."""
+    from datasets import load_dataset
+
+    yield from load_dataset(repo, config, split=split, streaming=True)
+
+
+def is_usable(chat: dict) -> bool:
+    """Whether Megatron-Bridge's chat path accepts the conversation as text-only chat."""
+    from megatron.bridge.data.conversation_processing import is_text_only_chat_example
+
+    return is_text_only_chat_example(chat)
+
+
 def to_chat(row: dict, field: str) -> dict | None:
     if field == "messages":
         return {"messages": row["messages"]} if row.get("messages") else None
@@ -99,57 +111,94 @@ def to_chat(row: dict, field: str) -> dict | None:
 def take(repo: str, config: str, split: str, field: str, want: int) -> tuple[list[dict], int]:
     """Stream the split and keep the first ``want`` usable rows; give up after ``10 * want`` rejects."""
     got, dropped = [], 0
-    try:
-        for row in load_dataset(repo, config, split=split, streaming=True):
-            chat = to_chat(row, field)
-            if chat is None or not is_text_only_chat_example(chat):
-                dropped += 1
-                if dropped > want * 10:
-                    break
-                continue
-            got.append(chat)
-            if len(got) >= want:
+    for row in stream(repo, config, split):
+        chat = to_chat(row, field)
+        if chat is None or not is_usable(chat):
+            dropped += 1
+            if dropped > want * 10:
                 break
-    except Exception as e:  # e.g. a gated split, or a schema the streaming reader cannot cast
-        print(f"{repo} {split}: failed ({type(e).__name__}: {str(e)[:80]})", flush=True)
+            continue
+        got.append(chat)
+        if len(got) >= want:
+            break
     return got, dropped
 
 
-def main():
-    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
-    parser.add_argument("--output_dir", required=True)
-    parser.add_argument("--total", type=int, default=40000, help="Rows to request across sources.")
-    parser.add_argument("--seed", type=int, default=1234)
-    args = parser.parse_args()
+def write_jsonl(path: str, rows: list[dict]) -> None:
+    """Write ``rows`` to ``path`` atomically: an interrupted write never leaves a partial file."""
+    tmp = f"{path}.tmp"
+    with open(tmp, "w") as f:
+        f.writelines(json.dumps(r) + "\n" for r in rows)
+    os.replace(tmp, path)
+
 
-    parts = os.path.join(args.output_dir, "parts")
+def read_part(path: str) -> list[dict] | None:
+    """Read a cached part, or ``None`` if it is missing or unreadable (it is then rebuilt)."""
+    try:
+        with open(path) as f:
+            return [json.loads(line) for line in f]
+    except (OSError, ValueError):
+        return None
+
+
+def build(output_dir: str, total: int, seed: int, allow_partial: bool = False) -> dict[str, int]:
+    """Build ``train.jsonl`` and ``validation.jsonl`` under ``output_dir``; return the split sizes."""
+    parts = os.path.join(output_dir, "parts")
     os.makedirs(parts, exist_ok=True)
     total_weight = sum(w for w, *_ in SOURCES)
-    rows = []
+    rows, summary = [], []
     for weight, repo, config, split, field in SOURCES:
-        want = max(1, round(args.total * weight / total_weight))
+        want = max(1, round(total * weight / total_weight))
         part = os.path.join(parts, f"{repo.split('/')[-1]}__{config}__{split}.jsonl")
-        if os.path.exists(part):
-            with open(part) as f:
-                got = [json.loads(line) for line in f]
-            if len(got) >= want:
-                rows.extend(got[:want])
-                print(f"{repo} {split}: {want} rows (cached)", flush=True)
-                continue
-        got, dropped = take(repo, config, split, field, want)
-        with open(part, "w") as f:
-            f.writelines(json.dumps(r) + "\n" for r in got)
-        rows.extend(got)
-        print(f"{repo} {split}: {len(got)}/{want} rows ({dropped} dropped)", flush=True)
-
-    random.Random(args.seed).shuffle(rows)
-    num_validation = min(1000, max(1, len(rows) // 20))
+        got = read_part(part)
+        if got is not None and len(got) >= want:
+            print(f"{repo} {split}: {want} rows (cached)", flush=True)
+        else:
+            try:
+                got, dropped = take(repo, config, split, field, want)
+            except (
+                Exception
+            ) as e:  # e.g. a gated split, or a schema the streaming reader cannot cast
+                if not allow_partial:
+                    raise RuntimeError(f"{repo} {split}: failed to load") from e
+                print(f"{repo} {split}: failed ({type(e).__name__}: {str(e)[:80]})", flush=True)
+                got, dropped = [], 0
+            write_jsonl(part, got)
+            print(f"{repo} {split}: {len(got)}/{want} rows ({dropped} dropped)", flush=True)
+            if len(got) < want and not allow_partial:
+                raise RuntimeError(
+                    f"{repo} {split}: {len(got)}/{want} usable rows; pass --allow_partial to keep them"
+                )
+        rows.extend(got[:want])
+        summary.append(
+            {"repo": repo, "config": config, "split": split, "want": want, "rows": len(got[:want])}
+        )
+
+    random.Random(seed).shuffle(rows)
+    num_validation = min(1000, len(rows) // 20)
     splits = {"validation": rows[:num_validation], "train": rows[num_validation:]}
+    if not splits["train"] or not splits["validation"]:
+        raise RuntimeError(f"too few rows ({len(rows)}) for a train/validation split")
     for name, chunk in splits.items():
-        path = os.path.join(args.output_dir, f"{name}.jsonl")
-        with open(path, "w") as f:
-            f.writelines(json.dumps(r) + "\n" for r in chunk)
-        print(f"wrote {path}: {len(chunk)} rows")
+        write_jsonl(os.path.join(output_dir, f"{name}.jsonl"), chunk)
+        print(f"wrote {name}.jsonl: {len(chunk)} rows", flush=True)
+    with open(os.path.join(output_dir, "blend_summary.json"), "w") as f:
+        json.dump({"total": total, "seed": seed, "sources": summary}, f, indent=2)
+    return {name: len(chunk) for name, chunk in splits.items()}
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    parser.add_argument("--output_dir", required=True)
+    parser.add_argument("--total", type=int, default=36800, help="Rows to request across sources.")
+    parser.add_argument("--seed", type=int, default=1234)
+    parser.add_argument(
+        "--allow_partial",
+        action="store_true",
+        help="Keep going when a source fails or is short, instead of stopping.",
+    )
+    args = parser.parse_args()
+    build(args.output_dir, args.total, args.seed, args.allow_partial)
 
 
 if __name__ == "__main__":
diff --git a/tests/unit/examples/test_build_qad_blend.py b/tests/unit/examples/test_build_qad_blend.py
new file mode 100644
index 00000000000..e719936f83e
--- /dev/null
+++ b/tests/unit/examples/test_build_qad_blend.py
@@ -0,0 +1,182 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Offline tests for the NemotronH omni VL tutorial's QAD blend builder.
+
+The Hugging Face reader and Megatron-Bridge's chat filter are replaced by fakes, so these run on
+CPU with neither ``datasets`` nor Megatron-Bridge installed.
+"""
+
+import importlib.util
+import json
+from pathlib import Path
+
+import pytest
+
+_SCRIPT = (
+    Path(__file__).resolve().parents[3]
+    / "examples"
+    / "megatron_bridge"
+    / "tutorials"
+    / "NemotronH-Omni-W4A4-QAD"
+    / "build_qad_blend.py"
+)
+_SPEC = importlib.util.spec_from_file_location("build_qad_blend", _SCRIPT)
+assert _SPEC is not None and _SPEC.loader is not None
+blend = importlib.util.module_from_spec(_SPEC)
+_SPEC.loader.exec_module(blend)
+
+
+def _chat(i: int) -> dict:
+    return {
+        "messages": [
+            {"role": "user", "content": f"q{i}"},
+            {"role": "assistant", "content": f"a{i}"},
+        ]
+    }
+
+
+@pytest.fixture
+def two_sources(monkeypatch):
+    """Two ``messages`` sources of equal weight, each with 100 rows; returns the per-source stream counts."""
+    monkeypatch.setattr(
+        blend,
+        "SOURCES",
+        [
+            (1, "org/a", "default", "train", "messages"),
+            (1, "org/b", "default", "train", "messages"),
+        ],
+    )
+    monkeypatch.setattr(blend, "is_usable", lambda chat: True)
+    calls = {"org/a": 0, "org/b": 0}
+
+    def stream(repo, config, split):
+        calls[repo] += 1
+        yield from (_chat(i) for i in range(100))
+
+    monkeypatch.setattr(blend, "stream", stream)
+    return calls
+
+
+def _read(path: Path) -> list[dict]:
+    return [json.loads(line) for line in path.read_text().splitlines()]
+
+
+def test_to_chat_handles_all_three_row_shapes():
+    msgs = [{"role": "user", "content": "hi"}]
+    assert blend.to_chat({"messages": msgs}, "messages") == {"messages": msgs}
+    assert blend.to_chat({"messages": []}, "messages") is None
+
+    text = "<extra_id_0>System\n<extra_id_1>User\nhello\n<extra_id_1>Assistant\nworld\n"
+    assert blend.to_chat({"text": text}, "text") == {
+        "messages": [
+            {"role": "user", "content": "hello"},
+            {"role": "assistant", "content": "world"},
+        ]
+    }
+    assert blend.to_chat({"text": "def f(): pass"}, "text") == {
+        "messages": [{"role": "assistant", "content": "def f(): pass"}]
+    }
+    assert blend.to_chat({"text": ""}, "text") is None
+
+
+def test_take_skips_rejected_rows_and_gives_up_after_too_many(monkeypatch):
+    monkeypatch.setattr(blend, "stream", lambda *a: iter([_chat(i) for i in range(50)]))
+    monkeypatch.setattr(blend, "is_usable", lambda chat: chat["messages"][0]["content"] != "q1")
+    got, dropped = blend.take("org/a", "default", "train", "messages", want=3)
+    assert [c["messages"][0]["content"] for c in got] == ["q0", "q2", "q3"]
+    assert dropped == 1
+
+    monkeypatch.setattr(blend, "is_usable", lambda chat: False)
+    got, dropped = blend.take("org/a", "default", "train", "messages", want=2)
+    assert got == [] and dropped == 21  # stops after 10 * want rejects
+
+
+def test_build_splits_weighted_rows_and_records_the_mixture(tmp_path, two_sources):
+    sizes = blend.build(str(tmp_path), total=100, seed=0)
+    assert sizes == {"validation": 5, "train": 95}
+    rows = _read(tmp_path / "train.jsonl") + _read(tmp_path / "validation.jsonl")
+    assert len(rows) == 100 and all(set(r) == {"messages"} for r in rows)
+    summary = json.loads((tmp_path / "blend_summary.json").read_text())
+    assert [(s["repo"], s["want"], s["rows"]) for s in summary["sources"]] == [
+        ("org/a", 50, 50),
+        ("org/b", 50, 50),
+    ]
+
+
+def test_build_reuses_cached_parts(tmp_path, two_sources):
+    blend.build(str(tmp_path), total=100, seed=0)
+    first = (tmp_path / "train.jsonl").read_text()
+    blend.build(str(tmp_path), total=100, seed=0)
+    assert two_sources == {"org/a": 1, "org/b": 1}  # second run streams nothing
+    assert (tmp_path / "train.jsonl").read_text() == first
+
+
+def test_build_rebuilds_a_truncated_part(tmp_path, two_sources):
+    blend.build(str(tmp_path), total=100, seed=0)
+    first = (tmp_path / "train.jsonl").read_text()
+    part = tmp_path / "parts" / "a__default__train.jsonl"
+    part.write_text(part.read_text()[:-20])  # an interrupted write, cut mid-record
+    blend.build(str(tmp_path), total=100, seed=0)
+    assert two_sources == {"org/a": 2, "org/b": 1}
+    assert len(_read(part)) == 50
+    assert (tmp_path / "train.jsonl").read_text() == first
+
+
+def test_build_writes_parts_atomically(tmp_path, two_sources, monkeypatch):
+    real_dumps = blend.json.dumps
+    calls = {"n": 0}
+
+    def flaky_dumps(obj, *args, **kwargs):
+        calls["n"] += 1
+        if calls["n"] == 75:  # partway through the second source's part
+            raise KeyboardInterrupt
+        return real_dumps(obj, *args, **kwargs)
+
+    monkeypatch.setattr(blend.json, "dumps", flaky_dumps)
+    with pytest.raises(KeyboardInterrupt):
+        blend.build(str(tmp_path), total=100, seed=0)
+    assert not (tmp_path / "parts" / "b__default__train.jsonl").exists()
+    assert len(_read(tmp_path / "parts" / "a__default__train.jsonl")) == 50
+
+
+def test_build_fails_on_a_source_that_cannot_load(tmp_path, two_sources, monkeypatch):
+    real_stream = blend.stream
+
+    def stream(repo, config, split):
+        if repo == "org/b":
+            raise TypeError("cannot cast")
+        yield from real_stream(repo, config, split)
+
+    monkeypatch.setattr(blend, "stream", stream)
+    with pytest.raises(RuntimeError, match="org/b train: failed to load"):
+        blend.build(str(tmp_path), total=100, seed=0)
+    assert not (tmp_path / "train.jsonl").exists()
+
+    sizes = blend.build(str(tmp_path), total=100, seed=0, allow_partial=True)
+    assert sizes == {"validation": 2, "train": 48}
+    summary = json.loads((tmp_path / "blend_summary.json").read_text())
+    assert [s["rows"] for s in summary["sources"]] == [50, 0]
+
+
+def test_build_fails_on_a_short_source(tmp_path, two_sources):
+    with pytest.raises(RuntimeError, match="org/a train: 100/250 usable rows"):
+        blend.build(str(tmp_path), total=500, seed=0)
+
+
+def test_build_rejects_an_empty_blend(tmp_path, two_sources, monkeypatch):
+    monkeypatch.setattr(blend, "is_usable", lambda chat: False)
+    with pytest.raises(RuntimeError, match="too few rows"):
+        blend.build(str(tmp_path), total=100, seed=0, allow_partial=True)

From 0076140fbdf90bc1a1ef85743e122c8f30bd4e1f Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Tue, 13 Oct 2026 18:36:10 +0000
Subject: [PATCH 21/24] Add the PTQ and QAD results to the NemotronH omni VL
 tutorial and freeze MTP in QAD

The QAD command passes --freeze_mtp, as the run behind the results did.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../NemotronH-Omni-W4A4-QAD/README.md         | 22 +++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
index 28bfc0cddab..b853597d394 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -6,10 +6,27 @@ End-to-end W4A4 quantization of a NemotronH omni vision-language checkpoint (arc
 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
+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 |
+| SciCode (subtask) | 47.13 | 46.48 | 46.58 |
+| SciCode (problem) | 15.10 | 13.65 | 15.38 |
+| τ³-bench banking | 18.88 | 17.40 | 20.10 |
+
+The exported W4A4 checkpoint is **70 GB**.
+
 ---
 
 ## Steps to Reproduce
@@ -91,6 +108,7 @@ srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/distill.py \
     --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 \
@@ -106,7 +124,7 @@ srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/distill.py \
 - `--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.
-- QAD does not train the MTP head, so the exported MTP head keeps its PTQ weights.
+- `--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.
 
 </details>

From 551257692c5102d4b02142e9537a2a7dff0685f0 Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Tue, 13 Oct 2026 18:54:07 +0000
Subject: [PATCH 22/24] Use runnable launch commands in the NemotronH omni VL
 tutorial; always fail closed in the blend builder

PTQ and export launch with torchrun, and QAD runs with python under Slurm
as the Nano tutorial does; the blend builder is invoked by its mounted
path. The builder no longer has a partial mode: a source that fails or
comes up short always stops the run, and its imports stay local so the
unit tests run without datasets or Megatron-Bridge.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../NemotronH-Omni-W4A4-QAD/README.md         | 15 ++++----
 .../build_qad_blend.py                        | 36 +++++++------------
 tests/unit/examples/test_build_qad_blend.py   | 33 ++++++++---------
 3 files changed, 34 insertions(+), 50 deletions(-)

diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
index b853597d394..20719b3c18d 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -33,14 +33,14 @@ The exported W4A4 checkpoint is **70 GB**.
 
 **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.
 
-Commands run under `srun`, one task per GPU, with `RANK` / `WORLD_SIZE` / `LOCAL_RANK` set from `SLURM_PROCID` / `SLURM_NTASKS` / `SLURM_LOCALID`.
+Single-node steps launch with `torchrun`; QAD runs on 8 nodes under Slurm with `python` (see [Slurm Usage](../../README.md#slurm-usage)).
 
 ### 1. Data Preparation
 
 Generate the QAD data with [`build_qad_blend.py`](build_qad_blend.py):
 
 ```bash
-python examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/build_qad_blend.py \
+python /opt/Model-Optimizer/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/build_qad_blend.py \
     --output_dir /path/to/qad_blend
 ```
 
@@ -67,8 +67,7 @@ PTQ runs on **one 4-GPU GB300 node** with 512 calibration samples.
 <summary>PTQ command (click to expand)</summary>
 
 ```bash
-# SBATCH --nodes=1 --ntasks-per-node=4 --gpus-per-node=4
-srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/quantize.py \
+torchrun --nproc_per_node 4 /opt/Model-Optimizer/examples/megatron_bridge/quantize.py \
     --hf_model_name_or_path <nemotron_h_omni-checkpoint> \
     --trust_remote_code \
     --recipe model_type/nemotron_h/ptq/nvfp4-aggressive-mse \
@@ -97,9 +96,10 @@ We used **8 nodes × 4 GB300 (32 GPUs)** with context parallelism 4 and expert p
 <details>
 <summary>QAD command (click to expand)</summary>
 
+> NOTE: We use `python -u` for slurm multi-node run here: `#SBATCH --nodes=8 --ntasks-per-node=4 --gpus-per-node=4`. See [Slurm Usage](../../README.md#slurm-usage).
+
 ```bash
-# SBATCH --nodes=8 --ntasks-per-node=4 --gpus-per-node=4
-srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/distill.py \
+python -u /opt/Model-Optimizer/examples/megatron_bridge/distill.py \
     --teacher_hf_path <nemotron_h_omni-checkpoint> \
     --student_hf_path <nemotron_h_omni-checkpoint> \
     --student_megatron_path /path/to/nemotron_h_omni_w4a4_megatron \
@@ -136,8 +136,7 @@ srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/distill.py \
 Export the QAD checkpoint to a Hugging Face checkpoint on one 4-GPU node (pipeline parallel 4):
 
 ```bash
-# SBATCH --nodes=1 --ntasks-per-node=4 --gpus-per-node=4
-srun ... python -u /opt/Model-Optimizer/examples/megatron_bridge/export_quantized_megatron_to_hf.py \
+torchrun --nproc_per_node 4 /opt/Model-Optimizer/examples/megatron_bridge/export_quantized_megatron_to_hf.py \
     --hf_model_name_or_path <nemotron_h_omni-checkpoint> \
     --trust_remote_code \
     --megatron_path /path/to/qad_output/checkpoints \
diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/build_qad_blend.py b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/build_qad_blend.py
index 388f0700f1b..49acaddcc44 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/build_qad_blend.py
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/build_qad_blend.py
@@ -27,8 +27,7 @@
 * ``text`` without markers (SFT-Code, SFT-MATH): one assistant turn, so the whole text is trained.
 
 A source that fails to load or yields fewer rows than its share stops the run, so the mixture
-cannot change silently; ``--allow_partial`` keeps what was collected instead. Either way the
-realized row count of every source is written to ``blend_summary.json``.
+cannot change silently.
 
 Each source is cached under ``<output_dir>/parts/``, written atomically, so an interrupted run
 resumes and an unreadable part is rebuilt. Run inside the NeMo container.
@@ -79,6 +78,10 @@
 ROLE = {"User": "user", "Assistant": "assistant", "System": "system"}
 
 
+# ``datasets`` and Megatron-Bridge are imported inside the two functions below, not at module
+# level, so the unit tests can load this script without either installed.
+
+
 def stream(repo: str, config: str, split: str):
     """Yield the rows of one Hugging Face split, in order, without downloading it."""
     from datasets import load_dataset
@@ -141,12 +144,12 @@ def read_part(path: str) -> list[dict] | None:
         return None
 
 
-def build(output_dir: str, total: int, seed: int, allow_partial: bool = False) -> dict[str, int]:
+def build(output_dir: str, total: int, seed: int) -> dict[str, int]:
     """Build ``train.jsonl`` and ``validation.jsonl`` under ``output_dir``; return the split sizes."""
     parts = os.path.join(output_dir, "parts")
     os.makedirs(parts, exist_ok=True)
     total_weight = sum(w for w, *_ in SOURCES)
-    rows, summary = [], []
+    rows = []
     for weight, repo, config, split, field in SOURCES:
         want = max(1, round(total * weight / total_weight))
         part = os.path.join(parts, f"{repo.split('/')[-1]}__{config}__{split}.jsonl")
@@ -159,20 +162,12 @@ def build(output_dir: str, total: int, seed: int, allow_partial: bool = False) -
             except (
                 Exception
             ) as e:  # e.g. a gated split, or a schema the streaming reader cannot cast
-                if not allow_partial:
-                    raise RuntimeError(f"{repo} {split}: failed to load") from e
-                print(f"{repo} {split}: failed ({type(e).__name__}: {str(e)[:80]})", flush=True)
-                got, dropped = [], 0
+                raise RuntimeError(f"{repo} {split}: failed to load") from e
+            if len(got) < want:
+                raise RuntimeError(f"{repo} {split}: only {len(got)}/{want} usable rows")
             write_jsonl(part, got)
-            print(f"{repo} {split}: {len(got)}/{want} rows ({dropped} dropped)", flush=True)
-            if len(got) < want and not allow_partial:
-                raise RuntimeError(
-                    f"{repo} {split}: {len(got)}/{want} usable rows; pass --allow_partial to keep them"
-                )
+            print(f"{repo} {split}: {want} rows ({dropped} dropped)", flush=True)
         rows.extend(got[:want])
-        summary.append(
-            {"repo": repo, "config": config, "split": split, "want": want, "rows": len(got[:want])}
-        )
 
     random.Random(seed).shuffle(rows)
     num_validation = min(1000, len(rows) // 20)
@@ -182,8 +177,6 @@ def build(output_dir: str, total: int, seed: int, allow_partial: bool = False) -
     for name, chunk in splits.items():
         write_jsonl(os.path.join(output_dir, f"{name}.jsonl"), chunk)
         print(f"wrote {name}.jsonl: {len(chunk)} rows", flush=True)
-    with open(os.path.join(output_dir, "blend_summary.json"), "w") as f:
-        json.dump({"total": total, "seed": seed, "sources": summary}, f, indent=2)
     return {name: len(chunk) for name, chunk in splits.items()}
 
 
@@ -192,13 +185,8 @@ def main():
     parser.add_argument("--output_dir", required=True)
     parser.add_argument("--total", type=int, default=36800, help="Rows to request across sources.")
     parser.add_argument("--seed", type=int, default=1234)
-    parser.add_argument(
-        "--allow_partial",
-        action="store_true",
-        help="Keep going when a source fails or is short, instead of stopping.",
-    )
     args = parser.parse_args()
-    build(args.output_dir, args.total, args.seed, args.allow_partial)
+    build(args.output_dir, args.total, args.seed)
 
 
 if __name__ == "__main__":
diff --git a/tests/unit/examples/test_build_qad_blend.py b/tests/unit/examples/test_build_qad_blend.py
index e719936f83e..33a0f04169c 100644
--- a/tests/unit/examples/test_build_qad_blend.py
+++ b/tests/unit/examples/test_build_qad_blend.py
@@ -21,6 +21,7 @@
 
 import importlib.util
 import json
+from collections import Counter
 from pathlib import Path
 
 import pytest
@@ -104,16 +105,14 @@ def test_take_skips_rejected_rows_and_gives_up_after_too_many(monkeypatch):
     assert got == [] and dropped == 21  # stops after 10 * want rejects
 
 
-def test_build_splits_weighted_rows_and_records_the_mixture(tmp_path, two_sources):
+def test_build_splits_weighted_rows(tmp_path, two_sources):
     sizes = blend.build(str(tmp_path), total=100, seed=0)
     assert sizes == {"validation": 5, "train": 95}
     rows = _read(tmp_path / "train.jsonl") + _read(tmp_path / "validation.jsonl")
     assert len(rows) == 100 and all(set(r) == {"messages"} for r in rows)
-    summary = json.loads((tmp_path / "blend_summary.json").read_text())
-    assert [(s["repo"], s["want"], s["rows"]) for s in summary["sources"]] == [
-        ("org/a", 50, 50),
-        ("org/b", 50, 50),
-    ]
+    # Equal weights: each source contributes its first 50 rows, q0..q49.
+    counts = Counter(r["messages"][0]["content"] for r in rows)
+    assert counts == {f"q{i}": 2 for i in range(50)}
 
 
 def test_build_reuses_cached_parts(tmp_path, two_sources):
@@ -157,26 +156,24 @@ def test_build_fails_on_a_source_that_cannot_load(tmp_path, two_sources, monkeyp
 
     def stream(repo, config, split):
         if repo == "org/b":
-            raise TypeError("cannot cast")
-        yield from real_stream(repo, config, split)
+            yield from real_stream(repo, config, split)
+            return
+        yield from (_chat(i) for i in range(10))  # yields some rows, then the reader fails
+        raise TypeError("cannot cast")
 
     monkeypatch.setattr(blend, "stream", stream)
-    with pytest.raises(RuntimeError, match="org/b train: failed to load"):
+    with pytest.raises(RuntimeError, match="org/a train: failed to load"):
         blend.build(str(tmp_path), total=100, seed=0)
+    assert not (tmp_path / "parts" / "a__default__train.jsonl").exists()
     assert not (tmp_path / "train.jsonl").exists()
 
-    sizes = blend.build(str(tmp_path), total=100, seed=0, allow_partial=True)
-    assert sizes == {"validation": 2, "train": 48}
-    summary = json.loads((tmp_path / "blend_summary.json").read_text())
-    assert [s["rows"] for s in summary["sources"]] == [50, 0]
-
 
 def test_build_fails_on_a_short_source(tmp_path, two_sources):
-    with pytest.raises(RuntimeError, match="org/a train: 100/250 usable rows"):
+    with pytest.raises(RuntimeError, match="org/a train: only 100/250 usable rows"):
         blend.build(str(tmp_path), total=500, seed=0)
+    assert not (tmp_path / "parts" / "a__default__train.jsonl").exists()
 
 
-def test_build_rejects_an_empty_blend(tmp_path, two_sources, monkeypatch):
-    monkeypatch.setattr(blend, "is_usable", lambda chat: False)
+def test_build_rejects_a_blend_too_small_to_split(tmp_path, two_sources):
     with pytest.raises(RuntimeError, match="too few rows"):
-        blend.build(str(tmp_path), total=100, seed=0, allow_partial=True)
+        blend.build(str(tmp_path), total=10, seed=0)

From 628ff2a481b34f0fa3d6f3f7d56ff3f8bd57c09a Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Tue, 13 Oct 2026 18:59:25 +0000
Subject: [PATCH 23/24] Add the MMLU-ProX result to the NemotronH omni VL
 tutorial

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md  | 1 +
 1 file changed, 1 insertion(+)

diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
index 20719b3c18d..40f6816a3c1 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -21,6 +21,7 @@ All checkpoints were served with the same vLLM settings (see [Section 5](#5-eval
 | 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 |

From 727e08324f2f005f812b00ad1c4a41836fdae353 Mon Sep 17 00:00:00 2001
From: Yue <yueshen@nvidia.com>
Date: Tue, 13 Oct 2026 19:03:43 +0000
Subject: [PATCH 24/24] Give the NemotronH omni VL tutorial a complete
 multi-node QAD launch

distill.py initializes torch.distributed from RANK / WORLD_SIZE / LOCAL_RANK
and MASTER_ADDR / MASTER_PORT, which Slurm does not set. The QAD command is
now an sbatch script whose srun tasks map the Slurm rank variables to them,
as in the 8-node runs behind the results.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yue <yueshen@nvidia.com>
---
 .../tutorials/NemotronH-Omni-W4A4-QAD/README.md        | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
index 40f6816a3c1..607e4e95e62 100644
--- a/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
+++ b/examples/megatron_bridge/tutorials/NemotronH-Omni-W4A4-QAD/README.md
@@ -34,7 +34,7 @@ The exported W4A4 checkpoint is **70 GB**.
 
 **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.
 
-Single-node steps launch with `torchrun`; QAD runs on 8 nodes under Slurm with `python` (see [Slurm Usage](../../README.md#slurm-usage)).
+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
 
@@ -97,9 +97,13 @@ We used **8 nodes × 4 GB300 (32 GPUs)** with context parallelism 4 and expert p
 <details>
 <summary>QAD command (click to expand)</summary>
 
-> NOTE: We use `python -u` for slurm multi-node run here: `#SBATCH --nodes=8 --ntasks-per-node=4 --gpus-per-node=4`. See [Slurm Usage](../../README.md#slurm-usage).
+`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 <nemotron_h_omni-checkpoint> \
     --student_hf_path <nemotron_h_omni-checkpoint> \
@@ -118,7 +122,7 @@ python -u /opt/Model-Optimizer/examples/megatron_bridge/distill.py \
     --save_interval 100 --checkpoint_keep_last -1 \
     --eval_iters 4 --eval_interval 100 \
     --log_interval 10 \
-    --output_dir /path/to/qad_output
+    --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.