Skip to content

Commit af6d7d9

Browse files
Leave MTP layers unquantized at any depth in the default PTQ exclusions (#2911)
### What does this PR do? Type of change: Bug fix The default PTQ exclusions (`configs/ptq/units/default_disabled_quantizers`, imported by the `general/ptq/*` recipes and `mtq.*_CFG` presets) disable multi-token-prediction layers with `mtp.*`, 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 with `examples/megatron_bridge/quantize.py` and a general recipe, `export_quantized_megatron_to_hf.py --vllm_fakequant_export` then 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 - New `tests/unit/recipe/test_presets.py::test_default_exclusions_disable_mtp_at_any_depth`: with `general/ptq/fp8_default-kv_fp8`, MTP quantizers at the root and under `language_model.` are disabled and decoder quantizers stay enabled; fails without the change. - `tests/unit/recipe`: 492 passed. - Found while validating vLLM fake-quant export of Qwen3.6-35B-A3B (#2889): the default `nvfp4_default-kv_fp8` checkpoint was refused on `language_model.mtp.layers.0.eh_proj.input_quantizer`. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ (nested MTP layers that a general recipe used to quantize now stay in BF16, as the exclusion intends) - 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 Related to #2889 (vLLM fake-quant export through Megatron-Bridge). 🤖 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 * **Bug Fixes** * Corrected default post-training quantization recipes for Megatron vision-language models so MTP layers remain unquantized consistently, including nested language-model layers. * Preserved quantization behavior for unrelated decoder projection layers. * Added regression coverage to verify both root-level and nested MTP exclusions in the default FP8 recipe. * **Documentation** * Added a changelog entry for the 0.49.0 bug fix. <!-- 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>
1 parent 754d3bc commit af6d7d9

3 files changed

Lines changed: 22 additions & 1 deletion

File tree

‎CHANGELOG.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ Changelog
3737

3838
**Bug Fixes**
3939

40+
- The default PTQ recipes now leave multi-token-prediction (MTP) layers unquantized in Megatron vision-language models too (``language_model.mtp.*``), as they already did at the model root.
4041
- Fix ONNX quantization placing a per-tensor Q/DQ pair on a rank-0 (scalar) constant input, which TensorRT's ONNX parser rejects (``convertAxis`` asserts the quantization axis lies in ``[0, nbDims]`` but ``nbDims == 0``). Rank-0 constant operands are now kept unquantized in the int8, fp8, and autotune export paths; the existing single-element Mul/Div/Pad stripping stays gated under ``target_dla``.
4142
- Fix PTQ of native Transformers FP8 experts to apply source weight scales during calibration and export. Re-quantize affected checkpoints from the original FP8 source.
4243
- Fix ONNX preprocessing with custom TensorRT plugin libraries on TensorRT 11.

‎modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@
4646
enable: false
4747
- quantizer_name: '*router*'
4848
enable: false
49-
- quantizer_name: 'mtp.*'
49+
- quantizer_name: '*mtp.*'
5050
enable: false
5151
- quantizer_name: 'output.*'
5252
enable: false

‎tests/unit/recipe/test_presets.py‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import argparse
2525

2626
import pytest
27+
import torch
2728

2829
import modelopt.torch.quantization as mtq
2930
from modelopt.recipe import load_recipe, presets
@@ -96,6 +97,25 @@ def test_w4a16_nvfp4_preset_disables_vllm_marlin_incompatible_projections():
9697
} <= disabled_quantizers
9798

9899

100+
def test_default_exclusions_disable_mtp_at_any_depth():
101+
# Megatron VLMs nest the MTP block under ``language_model.``; HF and Megatron LLMs keep it at the root.
102+
model = torch.nn.Module()
103+
for path in ("mtp", "language_model.mtp", "language_model.decoder"):
104+
parent = model
105+
for part in path.split("."):
106+
if not hasattr(parent, part):
107+
parent.add_module(part, torch.nn.Module())
108+
parent = getattr(parent, part)
109+
parent.add_module("proj", torch.nn.Linear(16, 16))
110+
111+
quant_cfg = load_recipe("general/ptq/fp8_default-kv_fp8").quantize.model_dump()
112+
mtq.quantize(model, {**quant_cfg, "algorithm": None})
113+
114+
assert model.language_model.decoder.proj.input_quantizer.is_enabled
115+
assert not model.mtp.proj.input_quantizer.is_enabled
116+
assert not model.language_model.mtp.proj.input_quantizer.is_enabled
117+
118+
99119
def test_local_hessian_layerwise_is_scoped_to_the_model_recipe():
100120
# Layerwise calibration needs identifiable decoder layers, so it belongs to the model
101121
# recipe the published numbers come from -- not to the shared preset, which also backs

0 commit comments

Comments
 (0)