Repository navigation
Leave MTP layers unquantized at any depth in the default PTQ exclusions - #2911
Merged
Merged
Conversation
The default exclusions matched only a root-level mtp.* block, so Megatron vision-language models (language_model.mtp.*) had their MTP layers quantized by the general recipes, and the vLLM fake-quant export then refused the checkpoint. Match *mtp.* instead. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 1 files. (2 skipped: 2 … | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly and concisely describes the main change: excluding MTP layers from default PTQ quantization at any nesting depth. Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Security Anti-Patterns PASS. The pull request changes only CHANGELOG.rst, one YAML recipe, and a unit test. It adds no Python changes under modelopt or examples and no dependency-manifest changes. The added test contains no… Full details: Security Anti-Patterns
-
Explanation
PASS. The pull request changes only CHANGELOG.rst, one YAML recipe, and a unit test. It adds no Python changes under modelopt or examples and no dependency-manifest changes. The added test contains no torch.load, numpy.load(..., allow_pickle=True), trust_remote_code=True, external-input eval/exec, or # nosec pattern. SECURITY.md also excludes tests from these contributor rules.
-
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 1 files. (2 skipped: 2 unsupported.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
✨ Finishing Touches: 💡 1
-
📝 Generate docstrings: 💡
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2911 +/- ##
==========================================
+ Coverage 70.08% 78.82% +8.74%
==========================================
Files 659 659
Lines 74081 74081
==========================================
+ Hits 51917 58392 +6475
+ Misses 22164 15689 -6475
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
cjluo-nv
approved these changes
Oct 21, 2026
cjluo-nv
left a comment
cjluo-nv
left a comment
Collaborator
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Approve: the shared exclusion now matches root and nested MTP layers, and the regression test verifies decoder quantization remains enabled.
No action needed.
Contributor
|
kevalmorabia97
added a commit
that referenced
this pull request
Oct 24, 2026
### 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, merged); #2873 added the flag on the old exporter. Needs a Megatron-Bridge with NVIDIA-NeMo/Megatron-Bridge#6565 (merged, not yet in a NeMo container); the script raises a clear error otherwise. The tests do not skip on an older Megatron-Bridge, so CI cannot silently stop covering the export. CI: the new `.github/actions/megatron-from-source` action replaces the nemo:26.08 container's Megatron-Bridge and Megatron-Core sources with Megatron-Bridge `1d9caf7324` (#6565's merge commit) and its pinned Megatron-LM in every CI job that runs in a NeMo container (today the `gpu_megatron` GPU tests and the `megatron_bridge` example tests), so the new tests run in CI. Remove it once a NeMo container includes #6565. The Pre-Requisites in `examples/megatron_bridge/README.md` show the same checkout and mounts for running locally. ### Usage ```bash 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). - In nemo:26.08 with Megatron-Bridge `1d9caf7324` installed the way the new action does it: `tests/gpu_megatron/torch/export` 139 passed; `tests/examples/megatron_bridge` (`test_quantize_export.py`, `test_qad.py`, `test_mlflow_utils.py`) 81 passed and `test_distill.py` 21 passed; `tests/gpu_megatron/torch/{prune,nas}` 62 passed; `tests/examples/vllm_serve` 120 passed, with `test_convert_ggml_to_int8` erroring as it does on main (the container's vLLM has no WNA16 kernel). - 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 #2873, unreleased) now needs a Megatron-Bridge with the quantizer sidecars (NVIDIA-NeMo/Megatron-Bridge#6565). - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ❌ ### Additional Information NVIDIA-NeMo/Megatron-Bridge#6565 and #2888 are merged. 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 fixed 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](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **New Features** * Export Megatron checkpoints for vLLM fake-quant serving, including LLMs and VLMs. Weight quantization is folded into BF16 weights; enabled activation and KV-cache settings are saved separately. * Fake-quant exports support tensor, expert, and pipeline parallel configurations, including expert parallelism for checkpoints built without grouped GEMM. * **Bug Fixes** * Sequential weight quantizers can now be folded into weights, preventing them from remaining active after folding. * Improved folding support for quantized weights with LoRA factors. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Type of change: Bug fix
The default PTQ exclusions (
configs/ptq/units/default_disabled_quantizers, imported by thegeneral/ptq/*recipes andmtq.*_CFGpresets) disable multi-token-prediction layers withmtp.*, which only matches an MTP block at the model root (HF checkpoints, Megatron LLMs). Megatron vision-language models nest it under the language model (language_model.mtp.layers.0...), so their MTP layers were quantized. For Qwen3.5/3.6 VLMs quantized withexamples/megatron_bridge/quantize.pyand a general recipe,export_quantized_megatron_to_hf.py --vllm_fakequant_exportthen refuses the checkpoint ("MTP quantization is not supported by vLLM fakequant export/reload"). The pattern is now*mtp.*, matching the model-specific recipes that already use*mtp*.Models whose MTP block was nested and previously quantized by a general recipe now keep it in BF16, which is what the exclusion intends. Recipes that quantize MTP on purpose (e.g. the Nemotron 3.5 Super ones) do not import this unit after their MTP entries, so they are unaffected.
Usage
No API change.
Testing
tests/unit/recipe/test_presets.py::test_default_exclusions_disable_mtp_at_any_depth: withgeneral/ptq/fp8_default-kv_fp8, MTP quantizers at the root and underlanguage_model.are disabled and decoder quantizers stay enabled; fails without the change.tests/unit/recipe: 492 passed.nvfp4_default-kv_fp8checkpoint was refused onlanguage_model.mtp.layers.0.eh_proj.input_quantizer.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Related to #2889 (vLLM fake-quant export through Megatron-Bridge).
🤖 Generated with Claude Code
Summary by CodeRabbit
Bug Fixes
Documentation