From 6373a69f698c9cf3518fabbda3ad7124386bfe0a Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Tue, 20 Oct 2026 23:58:05 -0700 Subject: [PATCH] Leave MTP layers unquantized at any depth in the default PTQ exclusions 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 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- CHANGELOG.rst | 1 + .../units/default_disabled_quantizers.yaml | 2 +- tests/unit/recipe/test_presets.py | 20 +++++++++++++++++++ 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 9ee5a31b096..7974cad044f 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -37,6 +37,7 @@ Changelog **Bug Fixes** +- 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. - 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``. - 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. - Fix ONNX preprocessing with custom TensorRT plugin libraries on TensorRT 11. diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml index bd017766937..8402587dcc2 100644 --- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml @@ -46,7 +46,7 @@ enable: false - quantizer_name: '*router*' enable: false - - quantizer_name: 'mtp.*' + - quantizer_name: '*mtp.*' enable: false - quantizer_name: 'output.*' enable: false diff --git a/tests/unit/recipe/test_presets.py b/tests/unit/recipe/test_presets.py index ed9b0e71388..16f5d31a2d6 100644 --- a/tests/unit/recipe/test_presets.py +++ b/tests/unit/recipe/test_presets.py @@ -24,6 +24,7 @@ import argparse import pytest +import torch import modelopt.torch.quantization as mtq from modelopt.recipe import load_recipe, presets @@ -96,6 +97,25 @@ def test_w4a16_nvfp4_preset_disables_vllm_marlin_incompatible_projections(): } <= disabled_quantizers +def test_default_exclusions_disable_mtp_at_any_depth(): + # Megatron VLMs nest the MTP block under ``language_model.``; HF and Megatron LLMs keep it at the root. + model = torch.nn.Module() + for path in ("mtp", "language_model.mtp", "language_model.decoder"): + parent = model + for part in path.split("."): + if not hasattr(parent, part): + parent.add_module(part, torch.nn.Module()) + parent = getattr(parent, part) + parent.add_module("proj", torch.nn.Linear(16, 16)) + + quant_cfg = load_recipe("general/ptq/fp8_default-kv_fp8").quantize.model_dump() + mtq.quantize(model, {**quant_cfg, "algorithm": None}) + + assert model.language_model.decoder.proj.input_quantizer.is_enabled + assert not model.mtp.proj.input_quantizer.is_enabled + assert not model.language_model.mtp.proj.input_quantizer.is_enabled + + def test_local_hessian_layerwise_is_scoped_to_the_model_recipe(): # Layerwise calibration needs identifiable decoder layers, so it belongs to the model # recipe the published numbers come from -- not to the shared preset, which also backs