[kevalmorabia97]

What does this PR do?

Type of change: new feature

examples/megatron_bridge/export_quantized_megatron_to_hf.py --vllm_fakequant_export now exports through Megatron-Bridge instead of ModelOpt's Megatron exporter: it checks the quantizers, folds the weight quantizers into BF16 weights (mtq.fold_weight) and calls bridge.save_hf_pretrained. Megatron-Bridge writes the enabled activation and KV-cache quantizers to quantizer_state.pth and a versioned quant_recipe.yaml under HF names, which vllm_serve_fakequant.py <dir> loads automatically. This is how the numerics team evaluates QAD checkpoints in vLLM (OMNIML-6019).

What the Megatron-Bridge path adds over the old exporter:

• any TP / EP / PP (--tp_size is new and only allowed with this flag), and every architecture Megatron-Bridge maps;
• VLMs: a QAD checkpoint holds only the language model, so the vision tower is loaded from HF first;
• weight-only formats (e.g. GGML IQ2) leave no quantizer enabled after folding, so no sidecars are written and stock vllm serve loads the export.

mtq.fold_weight now drops every buffer of a folded weight quantizer (e.g. a pinned GGML payload), not only _amax / _pre_quant_scale; otherwise Megatron-Bridge could not save a folded IQ2 model. SVDQuantLinear.fold_weight reads its low-rank factors (also buffers of the weight quantizer) before folding, so SVDQuant folds keep their correction term. mtq.fold_weight also folds sequential weight quantizers (e.g. W4A8 INT4 + FP8 weights), which it used to skip, leaving them unfolded and enabled.

check_mcore_vllm_fakequant_quantizers collects the existing exporter's checks (unsupported activation settings, MTP and lm_head quantization) so both paths reject the same inputs before anything is written.

Builds on #2888 (vLLM reload of the versioned recipe; this PR targets its branch); #2873 added the flag on the old exporter. Needs a Megatron-Bridge with NVIDIA-NeMo/Megatron-Bridge#6565; the script raises a clear error and the new tests skip otherwise (the CI's nemo:26.08 container is older).

Usage

torchrun --nproc_per_node 8 export_quantized_megatron_to_hf.py \
    --hf_model_name_or_path Qwen/Qwen3-8B \
    --megatron_path /output/qwen3_8b_nvfp4_qad/checkpoints \
    --tp_size 8 --vllm_fakequant_export \
    --export_unified_hf_path /tmp/Qwen3-8B-NVFP4-QAD-vllm-fq
python examples/vllm_serve/vllm_serve_fakequant.py /tmp/Qwen3-8B-NVFP4-QAD-vllm-fq -tp 8

Testing

Ran in nemo:26.10.rc3 with NVIDIA-NeMo/Megatron-Bridge#6565, 2× RTX 6000 Ada, unless noted:

• tests/gpu_megatron/torch/export/test_vllm_fakequant_mbridge_export.py (new): tiny Nemotron-H (NVFP4 + FP8 KV) and Qwen3 (W4A8, sequential weight quantizers) exported both ways match the old exporter (weights bitwise, ranges equal, recipe a superset of the old recipe's settings).
• No new example tests (they are slow): test_quantize_export.py::test_quantize_and_export also exports each checkpoint it already quantizes with --vllm_fakequant_export, at EP (Qwen3-MoE, per-expert names), TP (Qwen3-VL, vision tower bit-exact) and PP (Nemotron-H). test_qad.py adds that export only for the Qwen3.5-VL-MoE case (packed experts), since only QAD saves a VLM's language model alone. With 2 GPUs: 6 passed in 25 min (slowest case 386 s).
• test_tensor_quant_cpu.py::test_fold_weight_folds_sequential_weight_quantizer (new) fails without the sequential fold.
• test_tensor_quant_cpu.py::test_fold_weight_folds_svdquant_lora (new) fails without the SVDQuant change; test_unified_hf_export_and_check_safetensors (incl. nvfp4_svdquant) passes.
• Regression sweep: tests/unit, tests/gpu/torch/{quantization,export}, tests/gpu_megatron/torch/{export,quantization}, tests/gpu_vllm, tests/examples/{megatron_bridge,vllm_serve}: no new failures; the remaining failures also fail on main in the same container (missing optional packages, transformers/torch API drift, pytest timeouts). The CI image (nemo:26.08, older Megatron-Bridge) passes tests/gpu_megatron/torch/{export,quantization}, test_qad.py, test_quantize_export.py and tests/examples/vllm_serve, with the new tests skipped.
• End to end, scoring 64 CNN/DailyMail test articles (first 512 tokens each) with the per-token negative log-likelihood (NLL) in each engine:

Model (NVFP4 W4A4 + FP8 KV)	Export check	Megatron vs vLLM fake-quant NLL
Qwen3-0.6B	TP1 == TP2; served output identical to the old exporter's checkpoint (8/8 greedy, Δ logprob 0.0)	–
NVIDIA-Nemotron-Nano-9B-v2	TP1 == TP2 == PP2 (341 tensors, 128 ranges); weights bitwise equal to the old exporter; served output identical to it (8/8)	2.3465 vs 2.3467 (BF16: 2.2987 vs 2.2990)
Qwen3-30B-A3B	EP2 == PP2 (18,867 tensors, 18,720 ranges); vLLM TP2 serves correctly	–
Qwen3.6-35B-A3B (GB300, vLLM 0.31; NVFP4 MoE, FP8 attention / GatedDeltaNet, FP8 KV)	EP4 == TP2×EP2 == PP2×EP2 (1,045 tensors, 390 ranges, 410 recipe entries incl. constant-amax KV); vision and MTP weights bitwise equal to the source; text and image prompts answer correctly	each quantizer group alone within ±0.004; all together 2.2844 vs 2.3008 (HF transformers with the same quantizers: 2.2901)

On Qwen3.6 the three engines see the same 410 quantizers and ranges; with the MoE layers quantized to NVFP4, tiny kernel differences flip expert routing, which is where the 0.016 spread comes from (per-token |Δ| ≈ 0.2 with the MoE quantized vs 0.07–0.1 without).

• IQ2_XXS Qwen3-0.6B: no sidecars, stock vllm serve loads the export; Megatron logits are bitwise identical before and after folding.
• An internal quantization format with a custom KV-cache backend and backend_extra_args: both survive export and vLLM reload.

Before your PR is "Ready for review"

• Is this change backward compatible?: ❌ --vllm_fakequant_export (added in Megatron-Bridge: GGML PTQ, QAD and fake-quant export for large MoE and VLM models #2873, unreleased) now needs a Megatron-Bridge with the quantizer sidecars.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅
• Did you get Claude approval on this PR?: ❌

Additional Information

Merge order: #2888 (versioned recipe reload) → NVIDIA-NeMo/Megatron-Bridge#6565 → this PR. OMNIML-6019.

Known limits found while testing (unchanged by this PR): MTP and lm_head activation quantization cannot be reloaded by vLLM fake quant, so the export refuses them. The general PTQ recipes quantized the MTP layers of Megatron VLMs by mistake; #2911 fixes that. vllm_serve_fakequant.py still runs a one-sample calibration before loading the sidecars, so offline clusters need --modelopt-quant-dataset <local jsonl>.

🤖 Generated with Claude Code

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review: 3

• main
• release/.*
• feature/.*

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: eff6bcc3-20c7-4da3-86b9-f6841830d4e2

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[cjluo-nv]

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Changes requested: sequential weight quantizers pass validation but are not folded by the new Bridge export path.

Needs action:

• Fix sequential weight folding in export_vllm_fakequant, or reject it collectively before mutation; add a regression test comparing exported weights with the sequential QDQ result (see inline comment).

No action needed:

• The Bridge design is justified by architecture, VLM, and parallelism coverage versus extending the existing exporter.
• Existing test edits add coverage without weakening assertions. The new test’s NVIDIA header matches LICENSE_HEADER.
• Tests were inspected, not executed.

[kevalmorabia97]

/claude review

[cjluo-nv]

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

LGTM: the prior sequential-folding bug is resolved with regression coverage; human sign-off remains for the justified existing-test edits.

Needs action:

• Confirm the existing-test extensions in tests/examples/megatron_bridge: assertions remain intact, sidecar checks require a capable Bridge, and the timeout increase accommodates the added export.

No action needed:

• ✔️ Resolved since the last review: sequential weights now fold through both generic and TE grouped-linear paths; CPU QDQ and bitwise exporter-comparison tests cover the fix.
• Bridge reuse is justified by architecture, VLM, and TP/EP/PP coverage versus extending the old exporter.
• Buffer cleanup preserves SVDQuant corrections; the new NVIDIA header matches LICENSE_HEADER.
• Tests were inspected, not executed.