From 613175aa9e74c086b34d7ac09f2ad15e8ae91990 Mon Sep 17 00:00:00 2001 From: sugunav14 <178320438+sugunav14@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:28:01 -0700 Subject: [PATCH 01/28] Speed up FSDP2 MoE calibration by dropping redundant expert-weight gathers (#2359) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Perf enhancement `promote_static_block_weight_quantizers` runs at the end of `max_calibrate`. All it does is read quantizer state -- it takes the weight the iterator hands it and throws it away. For this it goes through `iter_weights_for_calibration`, and on a fused-MoE module that iterator yields `weight[idx]`, once per expert. Under FSDP2 the fused weight is a DTensor spread across ranks, so `weight[idx]` isn't a cheap view. Slicing it makes PyTorch pull the whole fused expert weight back from every rank, just to hand over one slice that the loop then drops. That happens once per expert, per projection, per layer -- tens of thousands of round-trips on a large MoE. The fix adds `iter_weight_quantizers_for_calibration`: - The base `QuantModule` implementation delegates to `iter_weights_for_calibration` and drops the weight, so every subclass gets a correct implementation for free and the two iterators cannot drift apart. - Only the fused-MoE class — the one where materializing the weight view is itself expensive — overrides it, walking the per-expert quantizer `ModuleList` directly. It keeps the same skip condition as the weight iterator, since fetching an attribute is free and only indexing collectives. `promote_static_block_weight_quantizers` then iterates quantizers instead of `(weight, quantizer)` pairs. No other caller changes: the other four call sites genuinely use the weight. Why not just wrap the promote loop in `enable_weight_access_and_writeback`, the way `weight_only_quantize` does? It would still gather per module for weights the loop never reads. `weight_only_quantize` needs the window because it actually computes amax from the weight; promote only needs the quantizer. Also included: a warn-once check if `_amax` is ever a `DTensor`. It should not be — `_amax` is a registered buffer, and `fully_shard` shards parameters, not buffers — which is why the `reduce_amax` in this loop stays local. If that assumption ever breaks, the reduction becomes a per-quantizer collective too, and the warning says so rather than letting it degrade silently. Measured on Qwen-3.8 2.4T at world 64 (8 nodes, B300), NVFP4: | | pre-export | |---|---| | before | 79.3 min | | after | **19.0 min** | 60.3 min saved, 4.2x. The block is gone rather than shortened -- the run logs 13 silent minutes in total, all of it model load. Calibration itself is unchanged at 2:23, so the saving comes from the promote loop rather than from work moving elsewhere. End-to-end projects to 41.7 min. No API change. `mtq.quantize` picks this up automatically. - Is this change backward compatible?: ✅ Additive; `iter_weights_for_calibration` and all its callers that use the weight are untouched. - 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)?: ✅ Under 0.47 Bug Fixes — the call site dates to 0.46. - Did you get Claude approval on this PR?: ❌ Not yet. - **Performance** - Reduced calibration overhead for FSDP2-sharded fused-MoE quantization by avoiding unnecessary expert-weight gathering when only quantizer state is needed. - Preserved quantizer ordering and projection filtering during calibration. - **Bug Fixes** - Added a warning when global amax reduction requires collective processing for individual quantizers. - **Tests** - Added coverage for gated and non-gated fused-expert calibration, including verification that fused weights are not unnecessarily indexed. --------- Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) Signed-off-by: Chad Voegele --- CHANGELOG.rst | 1 + .../quantization/nn/modules/quant_module.py | 11 ++++ .../torch/quantization/plugins/huggingface.py | 18 +++++++ .../torch/quantization/utils/core_utils.py | 16 +++++- .../plugins/test_fused_experts.py | 54 +++++++++++++++++++ 5 files changed, 99 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d4a29905480..8d0313eb6f3 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -53,6 +53,7 @@ Changelog **Bug Fixes** +- Speed up ``mtq.quantize`` on FSDP2-sharded fused-MoE models. Promoting static-block weight quantizers gathered each expert's slice of the fused weight across ranks even though only quantizer state is read, adding a collective per expert to calibration. - Avoid querying CUDA/Blackwell capability when ``NVFP4QTensor.quantize`` uses its CPU path or has the optional TensorRT-LLM fast path disabled. - Fix NVFP4 ONNX export to quantize FP4 weights with the published FP8 block scales, matching eager ModelOpt packed weights. Block scales below ``2**-9`` are now clamped to that minimum, and non-finite or negative scales raise an error. - Fix Megatron-Bridge Quantization Aware Distillation of a vision-language model silently discarding the ModelOpt state, so the distilled checkpoint restored no quantizers and exported as an unquantized model. Re-run QAD to regenerate any affected checkpoint. diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py index 1b5fb09062e..522f2d0db09 100644 --- a/modelopt/torch/quantization/nn/modules/quant_module.py +++ b/modelopt/torch/quantization/nn/modules/quant_module.py @@ -128,6 +128,17 @@ def iter_weights_for_calibration(self): weight_quantizer = getattr(self, quantizer_attr_names(weight_name).weight_quantizer) yield getattr(self, weight_name), weight_quantizer + def iter_weight_quantizers_for_calibration(self): + """Yield just the weight quantizers, without materializing the weight views. + + Callers that only inspect quantizer state must use this rather than discarding the + weight from :meth:`iter_weights_for_calibration`. Reading a weight is free here, but a + subclass whose weight view costs something -- the fused-MoE modules slice a 3-D DTensor + per expert, one redistribute collective each under FSDP2 -- overrides this to skip it. + """ + for _, weight_quantizer in self.iter_weights_for_calibration(): + yield weight_quantizer + @staticmethod @torch.no_grad() def _fold_weight_quantizer( diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 4acb4d30dfa..693c6c0d02a 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1100,6 +1100,24 @@ def iter_weights_for_calibration(self): for idx, q in enumerate(quantizers): yield weight[idx], q + def iter_weight_quantizers_for_calibration(self): + """Yield the per-expert weight quantizers without slicing the fused 3-D weight. + + Overrides the base, which would go through :meth:`iter_weights_for_calibration` and + evaluate ``weight[idx]``. Under FSDP2 that weight is a DTensor, so every expert slice + dispatches a redistribute collective to produce a value a quantizer-only caller discards. + """ + for weight_name, quantizers_name in ( + (self._first_proj_attr, self._first_proj_weight_quantizers_attr), + ("down_proj", "down_proj_weight_quantizers"), + ): + # Same skip condition as iter_weights_for_calibration, so the two stay in lockstep. + # Fetching the attribute is free; only indexing into it would collective. + quantizers = getattr(self, quantizers_name, None) + if getattr(self, weight_name, None) is None or quantizers is None: + continue + yield from quantizers + def fold_weight(self, keep_attrs: bool = False): """Bake each per-expert weight quantizer into its slice of the fused 3-D weight. diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 3bd31a03128..0adcb35a534 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -17,6 +17,7 @@ import copy import itertools +import warnings from collections import namedtuple from contextlib import ExitStack, contextmanager, nullcontext from typing import TYPE_CHECKING, Any @@ -1107,10 +1108,13 @@ def promote_static_block_weight_quantizers(model: nn.Module) -> int: for quantizer in state._member_quantizers() } converted = 0 + # Quantizer-only iteration: this loop reads quantizer state and never the weight, and asking + # for weights here costs one DTensor collective per expert on a fused-MoE model under FSDP2. + warned_dtensor_amax = False for _name, module in list(model.named_modules()): if not isinstance(module, QuantModule): continue - for _, quantizer in module.iter_weights_for_calibration(): + for quantizer in module.iter_weight_quantizers_for_calibration(): if isinstance(quantizer, SequentialQuantizer): if len(quantizer) == 0: continue @@ -1123,6 +1127,16 @@ def promote_static_block_weight_quantizers(model: nn.Module) -> int: amax = quantizer.amax if amax is None: continue + if isinstance(amax, DTensor) and not warned_dtensor_amax: + # _amax is a buffer, so FSDP2 leaves it replicated and reduce_amax below stays + # local. If that ever stops holding, the reduction becomes a per-quantizer + # collective and this loop needs a batched reduction instead. + warned_dtensor_amax = True + warnings.warn( + "promote_static_block_weight_quantizers: _amax is a DTensor, so the " + "per-quantizer global-amax reduction is a collective. Batch the reduction " + "before running this at scale." + ) if quantizer.is_nvfp4_static: # Grouped siblings share one canonical global_amax (common FP8 grid); otherwise # fall back to this quantizer's own per-block amax. diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 07ef2aaff24..49cd999a205 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -385,6 +385,60 @@ def test_fold_weight_keep_attrs_keeps_amax_disables_rotation(self, monkeypatch): # --------------------------------------------------------------------------- # Tests for export # --------------------------------------------------------------------------- +class TestIterWeightQuantizersForCalibration: + """The quantizer-only iterator must agree with the weight iterator and never index a weight. + + Indexing the fused 3-D weight is what dispatches a redistribute collective per expert under + FSDP2, so callers that only read quantizer state go through the quantizer-only path. + """ + + @staticmethod + def _convert(model): + expert_type = type(model.moe.experts) + TestQuantFusedExperts._cleanup_registry(expert_type) + register_fused_experts_on_the_fly(model) + converted = QuantModuleRegistry.convert(model.moe.experts) + TestQuantFusedExperts._cleanup_registry(expert_type) + return converted + + @pytest.mark.parametrize( + "model_cls", [_TinyMoEModel, _TinyNonGatedMoEModel], ids=["gated", "non_gated"] + ) + def test_yields_the_same_quantizers_in_the_same_order(self, model_cls): + experts = self._convert(model_cls()) + + from_weights = [q for _, q in experts.iter_weights_for_calibration()] + quantizers_only = list(experts.iter_weight_quantizers_for_calibration()) + + assert quantizers_only, "expected per-expert weight quantizers" + assert [id(q) for q in quantizers_only] == [id(q) for q in from_weights] + + @pytest.mark.parametrize( + "model_cls", [_TinyMoEModel, _TinyNonGatedMoEModel], ids=["gated", "non_gated"] + ) + def test_does_not_index_the_fused_weight(self, model_cls): + """The point of the override: no ``weight[idx]``, which is the per-expert collective.""" + experts = self._convert(model_cls()) + + class _NoIndexing(torch.Tensor): + @staticmethod + def __new__(cls, data): + return torch.Tensor._make_subclass(cls, data, False) + + def __getitem__(self, item): + raise AssertionError("iter_weight_quantizers_for_calibration indexed the weight") + + for name in (experts._first_proj_attr, "down_proj"): + weight = getattr(experts, name) + setattr(experts, name, nn.Parameter(_NoIndexing(weight.data), requires_grad=False)) + + assert list(experts.iter_weight_quantizers_for_calibration()) + # The weight iterator is the expensive one; it must still slice, or the guard above is + # not actually testing anything. + with pytest.raises(AssertionError, match="indexed the weight"): + list(experts.iter_weights_for_calibration()) + + class TestExportFusedExperts: @staticmethod def _cleanup_registry(mod_type): From 5d4d3953197352a879f28ebf5bac6e9e603ff171 Mon Sep 17 00:00:00 2001 From: Zhiyu Date: Wed, 9 Sep 2026 18:04:43 -0700 Subject: [PATCH 02/28] feat(quantization): PTQ support for Step-3.7 MoE checkpoints (#2202) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: New feature Adds PTQ support for **Step-3.7** (`stepfun-ai/Step-3.7-Flash`). Follow-up to [NVBug 6518665](https://nvbugspro.nvidia.com/bug/6518665) / [OMNIML-5583](https://jirasw.nvidia.com/browse/OMNIML-5583): with the export crash fixed in #2071 the run completes, but the checkpoint it writes is silently unquantized — ```json {"quantization": {"quant_algo": null, "kv_cache_quant_algo": "FP8", "quantized_layers": {}}} ``` Two independent causes, both from Step's `trust_remote_code` modeling code. **1. The expert weights were invisible to quantization.** Step-3.5 and Step-3.7 ship the same custom `MoELinear`: a plain `nn.Module` holding one 3-D `weight` of `[num_experts, out_features, in_features]`, whose `forward(x, expert_id)` runs `F.linear` against the selected slice. It is not an `nn.Linear`, and the weights sit on the projection submodule rather than on the expert container, so neither the plain-linear path nor `_fused_experts_wrapper_class` (which wants a 3-D `down_proj` *Parameter*) claims it. The `_QuantMoELinear` wrapper that handles exactly this layout has existed since #1063, but its registration was gated on the Step-3.5 class names: ```python if type(model).__name__ not in ("Step3p5ForCausalLM", "Step3p5Model"): return for module in model.modules(): if type(module).__name__ == "Step3p5MoEMLP": ``` Step-3.7's root is `Step3p7ForConditionalGeneration` and its container is `Step3p7MoEMLP`, so it returned immediately and no expert ever got a quantizer. Detection is now **structural** — a 3-D `weight` plus `num_experts` / `in_features` / `out_features` and a two-positional-argument forward — so any Step revision (or another model shipping this layout) is picked up without a third hardcoded name. `_reconstruct_fused_moe_linear` likewise matches the wrapper type instead of the generated `QuantMoELinear` class name; a model whose class is spelled differently would otherwise quantize fine but export unusable per-expert keys. **2. Step's module names don't match the general recipes.** The MoE block is `moe` and the dense sibling is `share_expert`, so `*.experts.*`, `*block_sparse_moe*` and `*mlp*` reach none of the routed experts. This PR ships `huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast` and `huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8`, which select `*moe*` and disable the router (`moe.gate`) and the shared expert — mirroring the existing Step-3.5 recipe — and documents the naming trap in `modelopt_recipes/ptq.md`. ```bash python examples/hf_ptq/hf_ptq.py --model /local/Step-3.7-Flash --trust_remote_code \ --recipe huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast \ --dataset /local/cnn_dailymail --calib_size 32 --export_path /local/Step-3.7-Flash-nvfp4 ``` - `tests/unit/torch/quantization/plugins/test_moe_linear.py` — structural detection (positive plus 2-D-weight / wrong-forward negatives), registration on a Step-3.7-shaped model, per-expert quantizers with calibrated amax, and reconstruction back to the 3-D parameter. Plus two export-dispatch tests added from review: the registry resolves a `Quant_SyntheticMoELinear` to `_export_moe_linear` (this one fails against the old name-keyed registration), and the handler fills an unrouted expert's input amax. - `tests/unit/recipe/test_step3p7_recipes.py` — drives both shipped recipes over a model mirroring Step's real paths (`model.language_model.layers[i].{moe,share_expert,mlp}`): routed experts NVFP4-quantized per expert, router / shared expert / dense MLP / `lm_head` per recipe scope. Ran locally (torch 2.11, transformers 5.5.4 — the version in the bug report): the two new files (12 tests) plus `tests/unit/recipe`, `tests/unit/torch/quantization/plugins/`, `tests/unit/torch/export/test_export_weight.py` and `test_export_registry.py` — 324 passed. Full `tests/unit` (minus onnx/puzzletron): 2490 passed, with 4 pre-existing `test_quant_aware_conversion.py` failures that reproduce unchanged on clean `main`. Not run end-to-end on the real Step-3.7-Flash checkpoint (1.4 TB / 8×B200) — QA can re-run against this branch with the recipe above. - Is this change backward compatible?: ✅ — Step-3.5 keeps working; the name gate is replaced by a superset. - 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?: ❌ - **Export dispatch keyed on the class name** (caught in review): registration was made class-name-independent, but `_export_moe_linear` in `hf_export_handlers.py` was still registered for the literal string `"QuantMoELinear"`, so a compatible class under another name bypassed the input-amax fallback for unrouted experts. The predicate now matches `_QuantMoELinear` through the MRO (lazy import, since the wrapper lives in the optional transformers plugin), keeping the name check as a fallback so the synthetic stand-ins in `test_export_registry.py` / `test_export_weight.py` still match. This was the same class-name coupling already fixed in `_reconstruct_fused_moe_linear`, in the file I hadn't looked at. - Canonical 2026 license header on the new test file. - Recipe comments corrected: `*moe*` matches the router, but **not** `share_expert` (`layers.N.share_expert.*` contains no `moe` segment) — that entry is an explicit guard, not an override. - `ptq.md` recommendation scoped to Step-3.7, since Step-3.5 has its own recipe. Pairs with #2203 (fail fast when a quant config matches no weight quantizer), which turns this class of silent no-op into an error for any model. Independent branches; either can merge first. 🤖 Generated with [Claude Code](https://claude.com/claude-code) - **New Features** - Added post-training quantization (PTQ) support for Step-3.7 Flash models, including per-expert quantization for routed MoE layers. - Added NVFP4 recipes for expert-only and MLP-only quantization, with optional FP8 KV-cache support. - **Bug Fixes** - Improved detection, quantization, reconstruction, and export of expert-indexed MoE layers across Step model revisions. - Added safeguards to prevent quantization of unsupported offloaded expert weights. - **Documentation** - Clarified Step-3.7 recipe selection, calibration guidance, and quantization exclusions. - **Tests** - Added coverage for expert, MLP, router, attention, shared-expert, export, and calibration behavior. --------- Signed-off-by: Zhiyu Cheng Co-authored-by: Claude Opus 5 (1M context) Signed-off-by: Chad Voegele --- CHANGELOG.rst | 2 + modelopt/torch/export/hf_export_handlers.py | 27 +- .../torch/quantization/plugins/huggingface.py | 175 ++++++-- .../ptq/nvfp4_experts_only-kv_fp8_cast.yaml | 56 +++ .../step3p7/ptq/nvfp4_mlp_only-kv_fp8.yaml | 61 +++ modelopt_recipes/ptq.md | 9 + tests/unit/recipe/test_step3p7_recipes.py | 177 ++++++++ .../quantization/plugins/test_moe_linear.py | 395 ++++++++++++++++++ 8 files changed, 871 insertions(+), 31 deletions(-) create mode 100644 modelopt_recipes/huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast.yaml create mode 100644 modelopt_recipes/huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8.yaml create mode 100644 tests/unit/recipe/test_step3p7_recipes.py create mode 100644 tests/unit/torch/quantization/plugins/test_moe_linear.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 8d0313eb6f3..4094d3bc0c6 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -17,6 +17,8 @@ Changelog - Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference and ``mtq.preserve_quantizer_attributes_context`` for restoring temporary quantizer property and type changes. Temporary folding snapshots affected fake-quant weights on a configurable device and restores them with their quantizer state; retained pre-quant scales are inactive, while shared weights, shared quantizers, and ``SequentialQuantizer`` weights are unsupported. - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Instead of setting the global scale from the largest per-block amax seen during calibration (plain ``max``, which leaves no room above it so any larger activation saturates), it anchors the scale to a low percentile of the per-block amax distribution, leaving the rest of the FP8 block-scale range as headroom: ``amax = max(rho * anchor, upper)``, where ``anchor`` and ``upper`` are the per-block amaxes at ``anchor_percentile`` (default 1) and ``upper_percentile`` (default 99.99; set to 100 to never clip calibration data), and ``rho`` (default 16384) is the headroom factor. Applies only to NVFP4 dynamic-block input quantizers; ``SequentialQuantizer`` activation quantizers raise. Weight scales are an orthogonal axis selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian``), so one recipe can combine a weight calibration with this activation policy in a single pass. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` with only the calibration algorithm swapped and exports a standard NVFP4 checkpoint. +- Add PTQ support for Step-3.7 (``stepfun-ai/Step-3.7-Flash``), whose routed experts were previously left unquantized. Quantize with the new ``huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast`` or ``huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8`` recipes rather than the general ones, which select experts by module names Step does not use. + *Megatron Framework (M-LM / M-Bridge)* - Add ``clamp_kv_cache_scales`` to ``export_mcore_gpt_to_hf``. Set it to ``False`` when exporting a QAT Megatron-Core model to preserve its learned FP8 KV-cache scales; the default retains the existing minimum scale of 1.0. diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 21a8a3fe246..8fb8b61e027 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -118,9 +118,30 @@ def _prepare_iterable_experts(name: str, moe_module: nn.Module, ctx: ExportConte # Export handlers are registered in the same precedence as the legacy model walk. -@ExportModuleRegistry.register( - "QuantMoELinear", predicate=lambda module: hasattr(module, "experts") -) +def _is_quant_moe_linear(module: nn.Module) -> bool: + """Whether ``module`` is an expert-indexed ``MoELinear`` expanded by ``_QuantMoELinear``. + + Matched by wrapper type first, not only by the dynamically generated class name + (``Quant`` + the model's own class name): the wrapper is registered structurally, so a + compatible remote-code class under any other name — or a second one, whose generated + name gets uniquified — would bypass this handler and export without the input-amax + fallback. The name check is kept as a fallback so stand-in modules match too. + + The wrapper lives in the optional transformers plugin, hence the lazy import. + """ + if not hasattr(module, "experts"): + return False + try: + from modelopt.torch.quantization.plugins.huggingface import _QuantMoELinear + + if isinstance(module, _QuantMoELinear): + return True + except ImportError: + pass + return any(cls.__name__ == "QuantMoELinear" for cls in type(module).__mro__) + + +@ExportModuleRegistry.register(predicate=_is_quant_moe_linear) def _export_moe_linear(name: str, module: nn.Module, ctx: ExportContext) -> None: """Fill missing input amax before child expert QuantLinears are exported.""" set_expert_quantizer_amax(list(module.experts), quantizer_attrs="input_quantizer") diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 693c6c0d02a..2515910ec6e 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -17,6 +17,7 @@ import inspect import logging +import re import warnings from contextlib import contextmanager from functools import partial @@ -1895,7 +1896,7 @@ def _is_param_grad_enabled_for_auto_quantize(pname, model): class _QuantMoELinear(QuantModule): - """Quantization wrapper for Step3p5 MoELinear modules (fused expert weights). + """Quantization wrapper for expert-indexed MoELinear modules (fused expert weights). MoELinear has weight shape [num_experts, out_features, in_features] with forward(x, expert_id). We expand it into per-expert nn.Linear modules so @@ -1913,6 +1914,19 @@ class _QuantMoELinear(QuantModule): def _setup(self): from accelerate import init_empty_weights + # Accelerate's CPU/disk offload (`device_map="auto"`, `--offload_folder`) leaves + # `weight` as a meta tensor and keeps the real value in the module's offload hook, + # keyed on the original `weight` name. Expanding that would copy meta storage into + # every expert and then delete the key the hook restores into, silently producing a + # checkpoint of zeros. Refuse instead of corrupting. + if self.weight.is_meta or getattr(getattr(self, "_hf_hook", None), "offload", False): + raise NotImplementedError( + f"{type(self).__name__}: expert-indexed MoELinear weights cannot be quantized " + "while offloaded by Accelerate (the weight is a meta tensor whose value lives " + "in the offload hook). Load the model without CPU/disk offload — more GPUs, or " + "a device_map that keeps the MoE layers resident — and re-run." + ) + dtype, device = self.weight.dtype, self.weight.device with init_empty_weights(): @@ -1932,38 +1946,143 @@ def _setup(self): self.experts = experts def forward(self, x, expert_id): - # experts[expert_id] is a _QuantLinear after quantization wrapping, - # providing per-expert input_quantizer and weight_quantizer. - # Cast input to match expert weight dtype before linear operation, - # then cast output to float32 to match original MoELinear forward behavior. + # experts[expert_id] is a _QuantLinear after quantization wrapping, providing + # per-expert input_quantizer and weight_quantizer. + # + # MoELinear.forward always promotes to fp32 for the matmul regardless of storage + # dtype (`F.linear(x.float(), self.weight[expert_id].float())`), so leaving the + # expert's weight at its native storage dtype (e.g. bf16) and downcasting x to + # match before the matmul would compute in bf16 and change the model's output even + # with every quantizer disabled -- the bf16 rounding this class exists to quantize + # *past*, not to reintroduce as a side effect of conversion. + # + # A prior version of this fix instead expanded every expert's weight in fp32 + # permanently in `_setup`. That reproduces Step's fp32 compute but turns a per-call + # transient promotion into persistent model state: on Step-3.7's full routed-expert + # set (42 layers x 3 projections x 288 experts x 4096 x 1280), doubling from bf16 to + # fp32 adds roughly 354 GiB held throughout calibration, on top of device placement + # already sized for bf16 -- a model that loaded successfully can then OOM. It also + # left disabled/unquantized experts reconstructed at fp32 in the exported checkpoint. + # + # Instead, only the one expert actually being called is promoted, transiently, for + # the duration of this one call -- matching Step's own per-call `.float()` memory + # profile instead of Step-3.7's full expert set. `expert.weight` is read here + # outside any `quantize_weight()` context, so `_get_quantized_weight` passes it + # through unchanged and this is the real underlying nn.Parameter (the same pattern + # `_setup` above uses), not a value computed by the quantizer -- so reassigning its + # `.data` genuinely mutates the persisted storage, not a transient wrapper. + # + # This must keep calling `expert(x)` (`__call__`, not `.forward()`) rather than + # reimplementing the input/weight-quantize/output-quantize sequence inline: some + # calibration algorithms (e.g. `local_hessian_calibrate`) register a + # `forward_pre_hook` directly on the quantized Linear module, which only fires + # through standard `nn.Module.__call__` dispatch. expert = self.experts[expert_id] - x = x.to(expert.weight.dtype) - return expert(x).float() + original_weight = expert.weight.data + with torch.no_grad(): + expert.weight.data = original_weight.float() + try: + out = expert(x.float()) + finally: + with torch.no_grad(): + expert.weight.data = original_weight + return out.float() + + +def _is_expert_indexed_moe_linear(module: nn.Module) -> bool: + """Whether ``module`` packs one projection's experts into an expert-indexed 3-D weight. + + The Step family (``stepfun-ai/Step-3.5-Flash``, ``stepfun-ai/Step-3.7-Flash``) ships a + custom ``MoELinear`` via ``trust_remote_code``: a plain ``nn.Module`` holding a single + ``weight`` of shape ``[num_experts, out_features, in_features]``, whose + ``forward(x, expert_id)`` runs ``F.linear`` against the selected expert's slice. The + weights therefore live on the projection submodule rather than on the expert container, + which is what :func:`_fused_experts_wrapper_class` looks for, and the module is not an + ``nn.Linear``, so neither the fused-experts path nor the plain linear path claims it. + + Detection is structural rather than keyed on class names so new Step revisions are picked + up without another hardcoded name, but the shape alone is not a sufficient contract: the + replacement forward indexes ``self.experts[expert_id]``, so it only works for callers that + pass a **scalar expert index**. Grouped-GEMM MoE layers share the exact same 3-D weight and + attribute set while passing a per-expert token-count *tensor* instead (e.g. Moondream3's + ``MoeFusedLinear.forward(input, m_sizes)``, which would raise ``TypeError: only integer + tensors of a single element can be converted to an index`` on the first calibration + forward). The second parameter must therefore be named ``expert_id``, which is the + scalar-index contract both Step revisions declare. + """ + weight = getattr(module, "weight", None) + if not isinstance(weight, (nn.Parameter, Tensor)) or weight.dim() != 3: + return False + if not all(hasattr(module, attr) for attr in ("num_experts", "in_features", "out_features")): + return False + # The wrapper rebuilds the weight as `num_experts` Linears of (out_features, in_features), + # so a 3-D weight laid out any other way would silently copy the wrong slices. + if tuple(weight.shape) != (module.num_experts, module.out_features, module.in_features): + return False + try: + params = list(inspect.signature(type(module).forward).parameters.values())[1:] + except (TypeError, ValueError): + return False + # The replacement forward is exactly `(x, expert_id)`, so anything the caller could pass + # beyond those two — a keyword-only `router_state`, *args, **kwargs — would raise once + # converted. Require the signature to match what the wrapper can honour. + return ( + len(params) == 2 + and all(p.kind is p.POSITIONAL_OR_KEYWORD and p.default is p.empty for p in params) + and params[1].name == "expert_id" + ) + +_STEP_FAMILY_RE = re.compile(r"(?i)^step\d") -def register_step3p5_moe_on_the_fly(model): - """Register Step3p5 MoELinear for quantization. - Step3p5 uses a custom MoELinear class (loaded via trust_remote_code) with - weight shape [num_experts, out_features, in_features] and forward(x, expert_id). - We detect it by model class name, then grab the type from the first MoE layer. +def _is_step_family_model(model: nn.Module) -> bool: + """Whether ``model`` is a Step-family root model (Step-3.5, Step-3.7, or a future revision). + + Matched against the ``step`` convention shared by ``model_type`` (``"step3p5"``, + ``"step3p7"``) and the remote-code class name (``Step3p5ForCausalLM``, + ``Step3p7ForConditionalGeneration``), not an exact revision, so a new Step release is + still picked up without another hardcoded name. This is deliberately narrower than the + shape/signature check in :func:`_is_expert_indexed_moe_linear` alone: that check accepts + any module with a matching 3-D weight and an ``(x, expert_id)`` forward, which is a + coincidence risk on its own -- an unrelated architecture happening to reuse the parameter + name ``expert_id`` with different semantics (a per-expert bias or post-scale, say) would + be claimed and have that behavior silently dropped by the replacement wrapper. Gating on + the model family keeps the shape check doing what it is actually good at: telling + Step revisions apart without a class-name allowlist, rather than distinguishing Step + from arbitrary third-party MoE code. + """ + model_type = str(getattr(getattr(model, "config", None), "model_type", "") or "") + return bool(_STEP_FAMILY_RE.match(model_type) or _STEP_FAMILY_RE.match(type(model).__name__)) + + +def register_moe_linear_on_the_fly(model): + """Register expert-indexed ``MoELinear`` modules (Step-3.5 / Step-3.7) for quantization. + + Without this the routed experts carry no quantizer at all: an experts-only recipe matches + nothing and the export writes a checkpoint with ``quant_algo: null``. """ - if type(model).__name__ not in ("Step3p5ForCausalLM", "Step3p5Model"): + if not _is_step_family_model(model): return - for module in model.modules(): - if type(module).__name__ == "Step3p5MoEMLP": - moe_linear_type = type(module.up_proj) - if QuantModuleRegistry.get(moe_linear_type) is None: - QuantModuleRegistry.register({moe_linear_type: f"hf.{moe_linear_type.__name__}"})( - _QuantMoELinear - ) - break + visited_types = set() + for name, module in model.named_modules(): + mod_type = type(module) + if mod_type in visited_types or QuantModuleRegistry.get(mod_type) is not None: + continue + visited_types.add(mod_type) + + if _is_expert_indexed_moe_linear(module): + print( + f"\033[1mDetected expert-indexed MoE linear '{name}' of type " + f"{mod_type.__name__}, registering with _QuantMoELinear.\033[0m" + ) + QuantModuleRegistry.register({mod_type: f"hf.{mod_type.__name__}"})(_QuantMoELinear) def _reconstruct_fused_moe_linear(model: nn.Module) -> None: - """Reconstruct QuantMoELinear per-expert weights back to original 3D MoELinear format. + """Reconstruct :class:`_QuantMoELinear` per-expert weights back to the 3-D MoELinear format. - After _process_quantized_modules, each expert's nn.Linear inside QuantMoELinear has: + After _process_quantized_modules, each expert's nn.Linear inside the wrapper has: - weight: fp4-quantized tensor [out_features, in_features] - weight_scale, weight_scale_2: per-block / global scales - input_scale: activation scale (if calibrated) @@ -1971,12 +2090,12 @@ def _reconstruct_fused_moe_linear(model: nn.Module) -> None: This stacks them back into the original MoELinear layout so the exported state_dict uses the original key names (e.g. moe.up_proj.weight with shape [N, out, in]). - Note: QuantMoELinear is the dynamically generated class name (Quant + MoELinear), - not _QuantMoELinear which is the implementation class. + Matched by wrapper type rather than by the dynamically generated class name (``Quant`` + + the model's own class name): a model whose class is not spelled ``MoELinear`` would + otherwise quantize normally but export unusable per-expert keys. """ for _name, module in model.named_modules(): - # Match QuantMoELinear (dynamically generated name) not _QuantMoELinear (implementation class) - if type(module).__name__ != "QuantMoELinear": + if not isinstance(module, _QuantMoELinear): continue n = module.num_experts @@ -2006,7 +2125,7 @@ def _reconstruct_fused_moe_linear(model: nn.Module) -> None: [ register_falcon_linears_on_the_fly, register_dbrx_moe_on_the_fly, - register_step3p5_moe_on_the_fly, + register_moe_linear_on_the_fly, register_fused_experts_on_the_fly, force_eager_experts_impl_on_the_fly, register_sparse_moe_on_the_fly, diff --git a/modelopt_recipes/huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast.yaml b/modelopt_recipes/huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast.yaml new file mode 100644 index 00000000000..086bdf4038c --- /dev/null +++ b/modelopt_recipes/huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast.yaml @@ -0,0 +1,56 @@ +# 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. + +# Model-specific PTQ recipe: Step-3.7 routed-experts-only NVFP4 with FP8 KV-cache cast. +# +# Step's remote code names the MoE block `moe` (not `experts` / `block_sparse_moe`), so the +# general experts-only recipe matches nothing on this architecture. The routed experts are +# `moe.{gate,up,down}_proj`, the router is `moe.gate`, and each MoE layer also has a dense +# `share_expert` that stays in BF16. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + nvfp4: configs/numerics/nvfp4 + kv_fp8_cast: configs/ptq/units/kv_fp8_cast + +metadata: + recipe_type: ptq + description: >- + Applies dynamic NVFP4 to Step-3.7 routed-expert weight and input quantizers, plus FP8 + KV-cache cast mode using constant amax; the router, the shared expert and the dense MLPs + stay unquantized. Uses max calibration. + +quantize: + algorithm: + method: max + layerwise: {enable: false} + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*moe*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*moe*input_quantizer' + cfg: + $import: nvfp4 + - $import: kv_fp8_cast + - $import: default_disabled_quantizers + # The router (`moe.gate`) is matched by `*moe*` above, so it must be disabled last + # (later wins). `share_expert` is not matched by any rule above — it is listed as an + # explicit guard so the shared expert stays BF16 if the patterns are ever widened. + - quantizer_name: '*moe.gate.*' + enable: false + - quantizer_name: '*share_expert*' + enable: false diff --git a/modelopt_recipes/huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8.yaml b/modelopt_recipes/huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8.yaml new file mode 100644 index 00000000000..b858d895117 --- /dev/null +++ b/modelopt_recipes/huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8.yaml @@ -0,0 +1,61 @@ +# 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. + +# Model-specific PTQ recipe: Step-3.7 MLP/MoE NVFP4 with calibrated FP8 KV cache. +# +# Same module-naming caveat as the experts-only recipe: Step's MoE block is `moe`, so the +# general `*mlp*` / `*.experts.*` patterns reach only the dense layers' `mlp` submodules and +# leave the routed experts — the bulk of the model — in BF16. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + nvfp4: configs/numerics/nvfp4 + kv_fp8: configs/ptq/units/kv_fp8 + +metadata: + recipe_type: ptq + description: >- + Applies dynamic NVFP4 to Step-3.7 routed-expert and dense-MLP weight and input + quantizers, plus calibrated FP8 KV-cache quantization; the router and the shared expert + stay unquantized. Uses max calibration. + +quantize: + algorithm: + method: max + layerwise: {enable: false} + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*moe*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*moe*input_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*mlp*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*mlp*input_quantizer' + cfg: + $import: nvfp4 + - $import: kv_fp8 + - $import: default_disabled_quantizers + # The router (`moe.gate`) is matched by `*moe*` above, so it must be disabled last + # (later wins). `share_expert` is not matched by any rule above — it is listed as an + # explicit guard so the shared expert stays BF16 if the patterns are ever widened. + - quantizer_name: '*moe.gate.*' + enable: false + - quantizer_name: '*share_expert*' + enable: false diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index ee2ca1f1e2a..8cb2b718bbc 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -301,6 +301,15 @@ A lighter case: **`models/stepfun-ai/Step-3.5-Flash/ptq/nvfp4-mlp-only`** is clo to one released checkpoint and carrying instance-specific disables (`share_expert`, `moe.gate`, the conv1d branches). +**`step3p7/ptq/{nvfp4_experts_only-kv_fp8_cast,nvfp4_mlp_only-kv_fp8}`** are the +Step-3.7 equivalents, and the reason they exist is **module naming**: Step calls +the MoE block `moe` and the dense sibling `share_expert`, so the general +recipes' `*.experts.*`, `*block_sparse_moe*` and `*mlp*` patterns match nothing +on the routed experts — the general recipe would quantize *nothing* and export a +checkpoint with `quant_algo: null`. These select `*moe*` instead and disable the +router (`moe.gate`) and `share_expert` on top. Use them, not the general +recipes, for Step-3.7 checkpoints; Step-3.5 has its own recipe above. + ### Algorithm overrides — `gemma`, `gemma4`, `mpt` These quantize the **same layers** as the general recipes; only the diff --git a/tests/unit/recipe/test_step3p7_recipes.py b/tests/unit/recipe/test_step3p7_recipes.py new file mode 100644 index 00000000000..1f62c96da14 --- /dev/null +++ b/tests/unit/recipe/test_step3p7_recipes.py @@ -0,0 +1,177 @@ +# 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. + +"""Step-3.7 PTQ recipes: what the `moe` / `share_expert` naming does and does not match.""" + +import types + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe +from modelopt.torch.quantization.nn import QuantModuleRegistry + +HIDDEN_SIZE = 32 +MOE_INTERMEDIATE_SIZE = 16 +NUM_EXPERTS = 2 + + +class _MoELinear(nn.Module): + """Step's expert-indexed projection: one 3-D weight, ``forward(x, expert_id)``.""" + + def __init__(self, num_experts, in_features, out_features): + super().__init__() + self.num_experts = num_experts + self.in_features = in_features + self.out_features = out_features + self.weight = nn.Parameter(torch.randn(num_experts, out_features, in_features) * 0.02) + + def forward(self, x, expert_id): + return F.linear(x.float(), self.weight[expert_id].float()) + + +class _StepMoEMLP(nn.Module): + def __init__(self): + super().__init__() + self.num_experts = NUM_EXPERTS + self.gate = nn.Linear(HIDDEN_SIZE, NUM_EXPERTS, bias=False) # router + self.up_proj = _MoELinear(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + self.gate_proj = _MoELinear(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + self.down_proj = _MoELinear(NUM_EXPERTS, MOE_INTERMEDIATE_SIZE, HIDDEN_SIZE) + + +class _StepMLP(nn.Module): + """Dense FFN — used both as the non-MoE layers' ``mlp`` and as ``share_expert``.""" + + def __init__(self): + super().__init__() + self.gate_proj = nn.Linear(HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE, bias=False) + self.up_proj = nn.Linear(HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE, bias=False) + self.down_proj = nn.Linear(MOE_INTERMEDIATE_SIZE, HIDDEN_SIZE, bias=False) + + +class _StepMoELayer(nn.Module): + def __init__(self): + super().__init__() + self.self_attn = nn.Module() + self.self_attn.q_proj = nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=False) + self.moe = _StepMoEMLP() + self.share_expert = _StepMLP() + + +class _StepDenseLayer(nn.Module): + def __init__(self): + super().__init__() + self.self_attn = nn.Module() + self.self_attn.q_proj = nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=False) + self.mlp = _StepMLP() + + +class _StepModel(nn.Module): + """Mirrors Step3p7ForConditionalGeneration's module paths (one MoE, one dense layer).""" + + def __init__(self): + super().__init__() + # register_moe_linear_on_the_fly gates on the Step-family model_type; real Step + # checkpoints carry this in config.json. + self.config = types.SimpleNamespace(model_type="step3p7") + self.model = nn.Module() + self.model.language_model = nn.Module() + self.model.language_model.layers = nn.ModuleList([_StepMoELayer(), _StepDenseLayer()]) + self.lm_head = nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=False) + + +@pytest.fixture(autouse=True) +def _unregister_moe_linear(): + yield + if QuantModuleRegistry.get(_MoELinear) is not None: + QuantModuleRegistry.unregister(_MoELinear) + + +def _quantize_with_recipe(name): + """Convert (no calibration) with a built-in recipe and return the model.""" + model = _StepModel() + config = load_recipe(name).quantize.model_dump() + config["algorithm"] = None + # `mtq.quantize` runs the custom-model plugins itself, which is what registers Step's + # `MoELinear` — no explicit registration here, so this also covers that hook firing. + mtq.quantize(model, config) + return model + + +def _enabled(module, quantizer="weight_quantizer"): + return getattr(module, quantizer).is_enabled + + +@pytest.mark.parametrize( + "recipe_name", + [ + "huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast", + "huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8", + ], +) +def test_routed_experts_are_quantized(recipe_name): + """The routed experts — the bulk of the model — must end up quantized, per expert.""" + model = _quantize_with_recipe(recipe_name) + moe = model.model.language_model.layers[0].moe + + for proj in ("up_proj", "gate_proj", "down_proj"): + experts = getattr(moe, proj).experts + assert len(experts) == NUM_EXPERTS + for expert in experts: + assert _enabled(expert) + assert _enabled(expert, "input_quantizer") + # Dynamic NVFP4: 16-element blocks along the input dim. + assert expert.weight_quantizer.block_sizes[-1] == 16 + assert expert.weight_quantizer.num_bits == (2, 1) + + +@pytest.mark.parametrize( + "recipe_name", + [ + "huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast", + "huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8", + ], +) +def test_router_shared_expert_and_head_stay_bf16(recipe_name): + """`*moe*` also matches the router; the shared expert and lm_head stay unquantized too.""" + model = _quantize_with_recipe(recipe_name) + moe_layer = model.model.language_model.layers[0] + + assert not _enabled(moe_layer.moe.gate) + for proj in ("gate_proj", "up_proj", "down_proj"): + assert not _enabled(getattr(moe_layer.share_expert, proj)) + assert not _enabled(model.lm_head) + + +def test_experts_only_leaves_dense_mlp_bf16(): + model = _quantize_with_recipe("huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast") + dense_mlp = model.model.language_model.layers[1].mlp + + for proj in ("gate_proj", "up_proj", "down_proj"): + assert not _enabled(getattr(dense_mlp, proj)) + + +def test_mlp_only_also_quantizes_dense_mlp(): + model = _quantize_with_recipe("huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8") + dense_mlp = model.model.language_model.layers[1].mlp + + for proj in ("gate_proj", "up_proj", "down_proj"): + assert _enabled(getattr(dense_mlp, proj)) + # Attention projections are out of scope for both recipes. + assert not _enabled(model.model.language_model.layers[1].self_attn.q_proj) diff --git a/tests/unit/torch/quantization/plugins/test_moe_linear.py b/tests/unit/torch/quantization/plugins/test_moe_linear.py new file mode 100644 index 00000000000..cceaacd44d5 --- /dev/null +++ b/tests/unit/torch/quantization/plugins/test_moe_linear.py @@ -0,0 +1,395 @@ +# 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. + +"""Tests for _QuantMoELinear: expert-indexed MoE weights (Step-3.5 / Step-3.7 remote code).""" + +import types + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +pytest.importorskip("transformers") + +import modelopt.torch.quantization as mtq +from modelopt.torch.export.hf_export_handlers import _export_moe_linear +from modelopt.torch.export.registry import ExportContext, ExportModuleRegistry +from modelopt.torch.quantization.nn import QuantModuleRegistry +from modelopt.torch.quantization.plugins.huggingface import ( + _is_expert_indexed_moe_linear, + _QuantMoELinear, + _reconstruct_fused_moe_linear, + register_moe_linear_on_the_fly, +) + +NUM_EXPERTS = 4 +HIDDEN_SIZE = 32 +MOE_INTERMEDIATE_SIZE = 16 +TOP_K = 2 + + +class _SyntheticMoELinear(nn.Module): + """Mimics Step-3.5 / Step-3.7 ``MoELinear`` (verbatim layout from their remote code).""" + + def __init__(self, num_experts, in_features, out_features): + super().__init__() + self.num_experts = num_experts + self.in_features = in_features + self.out_features = out_features + self.weight = nn.Parameter(torch.randn(num_experts, out_features, in_features) * 0.02) + + def forward(self, x, expert_id): + return F.linear(x.float(), self.weight[expert_id].float()) + + +class _SyntheticStepMoEMLP(nn.Module): + """Mimics ``Step3p7MoEMLP``: a router plus three expert-indexed projections.""" + + def __init__(self): + super().__init__() + self.num_experts = NUM_EXPERTS + self.top_k = TOP_K + self.gate = nn.Linear(HIDDEN_SIZE, NUM_EXPERTS, bias=False) + self.act_fn = nn.SiLU() + self.up_proj = _SyntheticMoELinear(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + self.gate_proj = _SyntheticMoELinear(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + self.down_proj = _SyntheticMoELinear(NUM_EXPERTS, MOE_INTERMEDIATE_SIZE, HIDDEN_SIZE) + + def forward(self, hidden_states): + tokens = hidden_states.view(-1, HIDDEN_SIZE) + routing = F.softmax(self.gate(tokens).float(), dim=-1) + weights, indices = torch.topk(routing, self.top_k, dim=-1) + out = torch.zeros_like(tokens) + for expert_id in range(self.num_experts): + pos, token_idx = torch.where(indices == expert_id) + if token_idx.numel() == 0: + continue + current = tokens[pos] + gate = self.act_fn(self.gate_proj(current, expert_id)) + up = self.up_proj(current, expert_id) + expert_out = self.down_proj(gate * up, expert_id) + out.index_add_(0, pos, (expert_out * weights[pos, token_idx, None]).to(out.dtype)) + return out.view_as(hidden_states) + + +class _TinyStepModel(nn.Module): + def __init__(self): + super().__init__() + # register_moe_linear_on_the_fly gates on the Step-family model_type; real Step + # checkpoints carry this in config.json. + self.config = types.SimpleNamespace(model_type="step3p7") + self.moe = _SyntheticStepMoEMLP() + + def forward(self, x): + return self.moe(x) + + +@pytest.fixture(autouse=True) +def _unregister_synthetic_moe_linear(): + """Keep the on-the-fly registration from leaking into other tests.""" + yield + if QuantModuleRegistry.get(_SyntheticMoELinear) is not None: + QuantModuleRegistry.unregister(_SyntheticMoELinear) + + +def _moe_quant_cfg(): + """Per-tensor INT8 on the expert projections only — CPU-friendly, no kernels needed.""" + return { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*moe*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + {"quantizer_name": "*moe*input_quantizer", "cfg": {"num_bits": 8, "axis": None}}, + {"quantizer_name": "*moe.gate.*", "enable": False}, + ], + "algorithm": "max", + } + + +def test_expert_indexed_moe_linear_is_detected(): + assert _is_expert_indexed_moe_linear( + _SyntheticMoELinear(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + ) + + +@pytest.mark.parametrize( + "module", + [ + pytest.param(nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE), id="plain_linear_2d_weight"), + pytest.param(nn.LayerNorm(HIDDEN_SIZE), id="norm_1d_weight"), + ], +) +def test_unrelated_modules_are_not_claimed(module): + assert not _is_expert_indexed_moe_linear(module) + + +def test_module_with_3d_weight_but_other_forward_is_not_claimed(): + """A 3-D weight alone is not enough — the forward must take ``(x, expert_id)``.""" + + class _NotExpertIndexed(_SyntheticMoELinear): + def forward(self, x, top_k_index, top_k_weights): + return x + + assert not _is_expert_indexed_moe_linear( + _NotExpertIndexed(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + ) + + +def test_disabled_quantizers_reproduce_bf16_weight_fp32_compute_parity(): + """Conversion must not change the model's output when every quantizer is disabled, and + must not permanently promote expert storage to fp32 to get there. + + ``MoELinear.forward`` always promotes to fp32 for the matmul regardless of storage + dtype (``F.linear(x.float(), self.weight[expert_id].float())``). A wrapper that + instead downcasts the fp32 activation to the weight's original storage dtype (e.g. + bf16) before the matmul silently changes the model even with quantization off. The + opposite mistake -- promoting every expert's *storage* to fp32 in `_setup` to match -- + reproduces Step's numerics but doubles the model's expert-weight memory footprint for + its entire lifetime (on Step-3.7's full routed-expert set, ~354 GiB); the promotion + must be transient, scoped to the one expert actually being called. + """ + torch.manual_seed(0) + num_experts, in_features, out_features = 2, 4096, 1280 + module = _SyntheticMoELinear(num_experts, in_features, out_features) + module.weight.data = module.weight.data.to(torch.bfloat16) + module.config = types.SimpleNamespace(model_type="step3p7") # satisfy the family gate + x = torch.randn(8, in_features, dtype=torch.bfloat16) + reference = module(x, 0) + + mtq.quantize(module, {"quant_cfg": [{"quantizer_name": "*", "enable": False}]}) + assert isinstance(module, _QuantMoELinear), "conversion did not happen; test is vacuous" + + # Storage stays at the checkpoint's own dtype -- only the matmul promotes, transiently. + for expert in module.experts: + assert expert.weight.dtype == torch.bfloat16 + + converted = module(x, 0) + assert torch.equal(converted, reference) + + # Reconstruction (export) must also see -- and keep -- the original storage dtype, not + # a permanently-promoted one. + _reconstruct_fused_moe_linear(module) + assert module.weight.dtype == torch.bfloat16 + + +def test_grouped_routing_module_is_not_claimed(): + """A grouped-GEMM MoE layer has the identical shape but no scalar-index contract. + + Moondream3's ``MoeFusedLinear`` carries the same 3-D weight and the same three + attributes, but its second argument is a per-expert token-count *tensor*. Claiming it + would make ``_QuantMoELinear.forward`` evaluate ``self.experts[m_sizes]`` and raise + ``TypeError: only integer tensors of a single element can be converted to an index`` + on the first calibration forward. + """ + + class _MoeFusedLinear(nn.Module): + def __init__(self): + super().__init__() + self.num_experts = NUM_EXPERTS + self.in_features = HIDDEN_SIZE + self.out_features = MOE_INTERMEDIATE_SIZE + self.weight = nn.Parameter( + torch.randn(NUM_EXPERTS, MOE_INTERMEDIATE_SIZE, HIDDEN_SIZE) * 0.02 + ) + + def forward(self, input, m_sizes): + return input + + assert not _is_expert_indexed_moe_linear(_MoeFusedLinear()) + + +def test_extra_forward_parameters_are_not_claimed(): + """The replacement forward is exactly ``(x, expert_id)``. + + Anything else the caller could pass — a keyword-only `router_state`, `*args`, + `**kwargs` — would raise `TypeError` once the module is converted. + """ + + class _WithRouterState(_SyntheticMoELinear): + def forward(self, x, expert_id, *, router_state=None): + return x + + class _WithKwargs(_SyntheticMoELinear): + def forward(self, x, expert_id, **kwargs): + return x + + for cls in (_WithRouterState, _WithKwargs): + assert not _is_expert_indexed_moe_linear( + cls(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + ), cls.__name__ + + +def test_transposed_weight_layout_is_not_claimed(): + """`[num_experts, in_features, out_features]` would rebuild each expert from wrong slices.""" + + class _TransposedMoELinear(_SyntheticMoELinear): + def __init__(self, num_experts, in_features, out_features): + super().__init__(num_experts, in_features, out_features) + self.weight = nn.Parameter(torch.randn(num_experts, in_features, out_features) * 0.02) + + assert not _is_expert_indexed_moe_linear( + _TransposedMoELinear(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + ) + + +def test_offloaded_weights_are_refused_not_silently_corrupted(): + """Accelerate offload leaves `weight` on meta, with the value in the module's hook. + + Expanding that would copy meta storage into each expert and delete the key the hook + restores into, exporting a checkpoint of zeros. Conversion must refuse instead. + """ + pytest.importorskip("accelerate") + from accelerate import cpu_offload # local: accelerate is an optional dependency + + model = _TinyStepModel() + cpu_offload(model.moe.up_proj, execution_device=torch.device("cpu")) + assert model.moe.up_proj.weight.is_meta + + with pytest.raises(NotImplementedError, match="offloaded by Accelerate"): + mtq.quantize(model, _moe_quant_cfg(), forward_loop=None) + + +def test_registration_is_not_gated_on_exact_revision_class_name(): + """Any Step-family root (matched by `model_type`/class-name convention, not an exact + revision) registers its `MoELinear` modules -- Step-3.7 as readily as Step-3.5.""" + model = _TinyStepModel() + assert QuantModuleRegistry.get(_SyntheticMoELinear) is None + + register_moe_linear_on_the_fly(model) + + assert issubclass(QuantModuleRegistry.get(_SyntheticMoELinear), _QuantMoELinear) + + +def test_non_step_model_with_identical_signature_is_not_registered(): + """A same-shape, same-signature module is not enough on its own to be claimed. + + The structural check in `_is_expert_indexed_moe_linear` cannot tell a real Step + `MoELinear` apart from unrelated code that happens to reuse the `(x, expert_id)` + parameter names with different semantics (a per-expert bias or post-scale, say) -- + `_QuantMoELinear` would silently drop that behavior. Registration is therefore also + gated on the model being Step-family; a structurally identical module on a model that + is not must not be registered. + """ + + class _ThirdPartyMoEModel(nn.Module): + def __init__(self): + super().__init__() + self.config = types.SimpleNamespace(model_type="not_step") + self.moe = _SyntheticStepMoEMLP() + + def forward(self, x): + return self.moe(x) + + model = _ThirdPartyMoEModel() + register_moe_linear_on_the_fly(model) + + assert QuantModuleRegistry.get(_SyntheticMoELinear) is None + + +def test_local_hessian_calibration_fires_through_the_transient_weight_swap(): + """`forward` must keep calling `expert(x)` (`__call__`), not `.forward()` directly. + + `local_hessian_calibrate` registers a `forward_pre_hook` on each quantized Linear + module and relies on standard `nn.Module.__call__` dispatch to fire it. A `forward` + that bypassed `__call__` -- e.g. to reimplement the input/weight-quantize/output-quantize + sequence inline instead of transiently swapping the expert's weight storage -- would + silently skip this hook and leave the weight quantizer uncalibrated (amax stays None). + """ + torch.manual_seed(0) + model = _TinyStepModel() + cfg = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*moe*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + {"quantizer_name": "*moe*input_quantizer", "cfg": {"num_bits": 8, "axis": None}}, + {"quantizer_name": "*moe.gate.*", "enable": False}, + ], + "algorithm": "local_hessian", + } + + def forward_loop(m): + for _ in range(3): + m(torch.randn(2, 8, HIDDEN_SIZE)) + + mtq.quantize(model, cfg, forward_loop=forward_loop) + + for expert in model.moe.up_proj.experts: + assert expert.weight_quantizer.amax is not None + + +def test_expert_indexed_moe_is_quantized_and_reconstructed(): + """Each expert gets its own quantizers, and export folds them back to the 3-D layout.""" + torch.manual_seed(0) + model = _TinyStepModel() + reference_weight = model.moe.up_proj.weight.detach().clone() + + def forward_loop(m): + m(torch.randn(2, 8, HIDDEN_SIZE)) + + mtq.quantize(model, _moe_quant_cfg(), forward_loop=forward_loop) + + # Every expert of every projection carries its own calibrated quantizer pair. + for proj in ("up_proj", "gate_proj", "down_proj"): + experts = getattr(model.moe, proj).experts + assert len(experts) == NUM_EXPERTS + for expert in experts: + assert expert.weight_quantizer.is_enabled + assert expert.weight_quantizer.amax is not None + # The router stays untouched. + assert not model.moe.gate.weight_quantizer.is_enabled + + _reconstruct_fused_moe_linear(model) + + # Back to the original ``[num_experts, out_features, in_features]`` parameter, so the + # exported keys match the hub checkpoint instead of per-expert names. + up_proj = model.moe.up_proj + assert not hasattr(up_proj, "experts") + assert up_proj.weight.shape == reference_weight.shape + assert torch.equal(up_proj.weight, reference_weight) + + +def test_export_handler_matches_a_differently_named_wrapper(): + """Export dispatch must key on the wrapper type, not the generated class name. + + The registration is structural, so a compatible remote-code class can be named + anything; its generated class is then ``Quant``. If the export handler only + matched the literal name ``QuantMoELinear``, such a module would skip + ``_export_moe_linear`` and export without the input-amax fallback for experts that + calibration never routed to. + """ + model = _TinyStepModel() + mtq.quantize(model, _moe_quant_cfg(), forward_loop=lambda m: m(torch.randn(2, 8, HIDDEN_SIZE))) + + converted = model.moe.up_proj + # The generated name is derived from the model's own class, not from `MoELinear`. + assert type(converted).__name__ == "Quant_SyntheticMoELinear" + assert ExportModuleRegistry.match(converted) is _export_moe_linear + + +def test_export_handler_fills_input_amax_for_unrouted_experts(): + """The handler is what gives never-routed experts an input amax before export.""" + torch.manual_seed(0) + model = _TinyStepModel() + mtq.quantize(model, _moe_quant_cfg(), forward_loop=lambda m: m(torch.randn(2, 8, HIDDEN_SIZE))) + + experts = model.moe.up_proj.experts + # Simulate an expert that calibration never routed a token to. + experts[0].input_quantizer.reset_amax() + assert experts[0].input_quantizer.amax is None + assert any(e.input_quantizer.amax is not None for e in experts), "need a donor amax" + + _export_moe_linear("moe.up_proj", model.moe.up_proj, ExportContext(model, torch.float16)) + + assert experts[0].input_quantizer.amax is not None From 839fba6f66f688bbc6400df228f5b45b57783785 Mon Sep 17 00:00:00 2001 From: Ajinkya Rasane <131806219+ajrasane@users.noreply.github.com> Date: Thu, 10 Sep 2026 12:16:43 -0400 Subject: [PATCH 03/28] [6410139] Fix ONNX AutoCast for large external initializers (#2317) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix Fix ONNX AutoCast for models whose external initializers exceed the in-memory protobuf limit. - Keep external tensor payloads reference-only during graph sanitization and type inference, then materialize them once before value-dependent classification and conversion. - Duplicate shared initializers directly in `GraphProto`, preserving external-data metadata without reading tensor bytes. - Use file-backed ONNX paths for validation, shape inference, reference execution, and custom-operator inspection when required. - Avoid a redundant sanitizer pass in the fully sanitized AutoCast path while preserving existing behavior for direct `PrecisionConverter` and `convert_to_f16()` callers. No CLI flags, dependencies, or public return types change. ```bash python -m modelopt.onnx.autocast \ --onnx_path model.onnx \ --output_path model_bf16.onnx \ --low_precision_type bf16 ``` - Ran the complete CPU-only AutoCast unit suite with no GPU visible: 248 passed. - Ran focused ONNX utility regressions covering shared initializer duplication and file-backed protobuf routing: 8 passed. - Ran pre-commit on all changed files. - Ran CPU-only integration coverage with an exact 2,147,485,696-byte external initializer using protobuf 7.35.1 and 6.33.6. - Verified BF16, FP16, shared-initializer, and aggregate-external-data runtime-cast cases. - Verified every integration output with `onnx.checker.check_model(..., full_check=True)`; outputs expected to remain external-data-backed did so. Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors). - Is this change backward compatible?: ✅ - 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?: ❌ > 🤖 _Generated by Codex (AI agent)._ - **Bug Fixes** - Fixed ONNX AutoCast failures for models with external initializers larger than 2 GiB. - Improved handling of large or external-data models during shape inference, conversion, and runtime validation. - Preserved external initializer metadata while avoiding unnecessary data materialization. - Improved temporary-file cleanup when model loading or inference fails. - Improved processing of shared initializers, nested graphs, and custom nodes. - **Tests** - Added coverage for large models, external initializers, nested graphs, custom nodes, and runtime cleanup. - **Documentation** - Updated the changelog with recent fixes and benchmarking information. --------- Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> Signed-off-by: Chad Voegele --- CHANGELOG.rst | 1 + modelopt/onnx/autocast/convert.py | 8 +- modelopt/onnx/autocast/graphsanitizer.py | 26 ++- modelopt/onnx/autocast/precisionconverter.py | 17 +- modelopt/onnx/autocast/referencerunner.py | 90 ++++----- modelopt/onnx/utils.py | 146 +++++++++----- tests/unit/onnx/autocast/test_autocast.py | 43 +++- .../unit/onnx/autocast/test_graphsanitizer.py | 64 ++++++ .../onnx/autocast/test_referencerunner.py | 51 +++++ tests/unit/onnx/test_onnx_utils.py | 187 ++++++++++++++++++ 10 files changed, 524 insertions(+), 109 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4094d3bc0c6..41349796945 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -56,6 +56,7 @@ Changelog **Bug Fixes** - Speed up ``mtq.quantize`` on FSDP2-sharded fused-MoE models. Promoting static-block weight quantizers gathered each expert's slice of the fused weight across ranks even though only quantizer state is read, adding a collective per expert to calibration. +- Fix ONNX AutoCast failing on models with external initializers larger than 2 GiB. - Avoid querying CUDA/Blackwell capability when ``NVFP4QTensor.quantize`` uses its CPU path or has the optional TensorRT-LLM fast path disabled. - Fix NVFP4 ONNX export to quantize FP4 weights with the published FP8 block scales, matching eager ModelOpt packed weights. Block scales below ``2**-9`` are now clamped to that minimum, and non-finite or negative scales raise an error. - Fix Megatron-Bridge Quantization Aware Distillation of a vision-language model silently discarding the ModelOpt state, so the distilled checkpoint restored no quantizers and exported as an unquantized model. Re-run QAD to regenerate any affected checkpoint. diff --git a/modelopt/onnx/autocast/convert.py b/modelopt/onnx/autocast/convert.py index d9bc8d68c99..e259235a911 100644 --- a/modelopt/onnx/autocast/convert.py +++ b/modelopt/onnx/autocast/convert.py @@ -23,6 +23,7 @@ nodes. """ +import os from copy import deepcopy import numpy as np @@ -108,7 +109,7 @@ def convert_to_mixed_precision( onnx.ModelProto: The converted mixed precision model. """ # Load and process model - model = onnx.load(onnx_path, load_external_data=True) + model = onnx.load(onnx_path, load_external_data=False) assert low_precision_type in ["fp16", "bf16"], "low_precision_type must be either fp16 or bf16" original_network_io_metadata = _capture_network_io_metadata(model, keep_io_types) @@ -146,6 +147,7 @@ def convert_to_mixed_precision( trt_plugins=trt_plugins, trt_plugins_precision=trt_plugins_precision, max_ir_version=LATEST_IR_VERSION_SUPPORTED_BY_ORT, + onnx_path=onnx_path, ) graph_sanitizer.sanitize() model = graph_sanitizer.model @@ -154,6 +156,9 @@ def convert_to_mixed_precision( # as an exception (triggering infer_types' standalone type-inference fallback) instead of # silently leaving tensors untyped, which would break later type lookups. model = onnx_utils.infer_types(model, use_standalone_type_inference, strict_mode=True) + onnx.external_data_helper.load_external_data_for_model( + model, os.path.dirname(os.path.abspath(onnx_path)) + ) value_info_map, initializer_map, node_to_init_map = utils.setup_mappings(model) # Automatically add 'trt' to list of providers if custom ops are detected @@ -187,6 +192,7 @@ def convert_to_mixed_precision( custom_ops=graph_sanitizer.custom_ops, use_standalone_type_inference=use_standalone_type_inference, original_network_io_metadata=original_network_io_metadata, + sanitize_model=False, ) # Obtain reference data diff --git a/modelopt/onnx/autocast/graphsanitizer.py b/modelopt/onnx/autocast/graphsanitizer.py index 2154a42568e..e9be72c0c27 100644 --- a/modelopt/onnx/autocast/graphsanitizer.py +++ b/modelopt/onnx/autocast/graphsanitizer.py @@ -15,6 +15,8 @@ """Graph sanitization and optimization for ONNX models.""" +import os + import numpy as np import onnx import onnx_graphsurgeon as gs @@ -25,7 +27,12 @@ import modelopt.onnx.utils as onnx_utils from modelopt.onnx.autocast.logging_config import logger from modelopt.onnx.quantization.graph_utils import cast_custom_ops -from modelopt.onnx.trt_utils import interpret_trt_plugins_precision_flag +from modelopt.onnx.trt_utils import ( + get_custom_layers, + infer_types_shapes_tensorrt, + interpret_trt_plugins_precision_flag, + set_trt_plugin_domain, +) class GraphSanitizer: @@ -38,6 +45,7 @@ def __init__( max_ir_version: int | None = None, trt_plugins: list[str] | None = [], trt_plugins_precision: list[str] | None = [], + onnx_path: str | None = None, ) -> None: """Initialize GraphSanitizer. @@ -46,6 +54,7 @@ def __init__( min_opset: minimum opset version to use max_ir_version: maximum IR version supported by ORT trt_plugins: list of TensorRT plugin library paths in .so format (compiled shared library). + onnx_path: path to the source ONNX model, used to resolve external data. """ self.model = model self.min_opset = min_opset @@ -55,6 +64,8 @@ def __init__( self.custom_ops_low_precision_nodes = [] self.trt_plugins = trt_plugins self.trt_plugins_precision = trt_plugins_precision or [] + self.onnx_path = os.path.abspath(onnx_path) if onnx_path is not None else None + self.external_data_dir = os.path.dirname(self.onnx_path) if self.onnx_path else "" def sanitize(self) -> None: """Sanitize the model graph. @@ -118,13 +129,14 @@ def find_custom_nodes(self) -> None: node.op_type for node in self.model.graph.node if node.op_type not in self.standard_ops } if self.custom_ops: - from modelopt.onnx.trt_utils import infer_types_shapes_tensorrt, set_trt_plugin_domain - # Set TensorRT plugin domain info in the graph for ORT compatibility self.model = set_trt_plugin_domain(self.model, self.custom_ops) # Infer types and shapes in the graph for ORT compatibility - self.model = infer_types_shapes_tensorrt(self.model, self.trt_plugins) + _, all_tensor_info = get_custom_layers(self.onnx_path or self.model, self.trt_plugins) + self.model = infer_types_shapes_tensorrt( + self.model, self.trt_plugins, all_tensor_info=all_tensor_info + ) def remove_disconnected_outputs(self) -> None: """Remove disconnected outputs from the model.""" @@ -501,7 +513,7 @@ def _get_initializer_value(self, name: str, return_array: bool = False) -> np.nd """Get value from an initializer by name.""" for init in self.model.graph.initializer: if init.name == name: - value = numpy_helper.to_array(init) + value = numpy_helper.to_array(init, base_dir=self.external_data_dir) return value if return_array else value.item() return None @@ -516,7 +528,7 @@ def _convert_fp64_initializers(self) -> bool: for initializer in self.model.graph.initializer: if initializer.data_type == onnx.TensorProto.DOUBLE: # Convert the data to FP32 - fp64_data = numpy_helper.to_array(initializer) + fp64_data = numpy_helper.to_array(initializer, base_dir=self.external_data_dir) fp32_data = fp64_data.astype(np.float32) # Create new initializer with FP32 data @@ -575,7 +587,7 @@ def _convert_fp64_nodes(self) -> bool: for attr in node.attribute: if attr.name == "value" and attr.t.data_type == onnx.TensorProto.DOUBLE: # Convert the tensor value to FP32 - fp64_data = numpy_helper.to_array(attr.t) + fp64_data = numpy_helper.to_array(attr.t, base_dir=self.external_data_dir) fp32_data = fp64_data.astype(np.float32) new_tensor = numpy_helper.from_array(fp32_data) attr.t.CopyFrom(new_tensor) diff --git a/modelopt/onnx/autocast/precisionconverter.py b/modelopt/onnx/autocast/precisionconverter.py index bd5c0fc13f3..7db660b29b4 100644 --- a/modelopt/onnx/autocast/precisionconverter.py +++ b/modelopt/onnx/autocast/precisionconverter.py @@ -99,6 +99,7 @@ def __init__( tensor_block_dict: dict[str, dict[str, list[int]]] = {}, use_standalone_type_inference: bool = False, original_network_io_metadata: dict[str, list[onnx.ValueInfoProto]] | None = None, + sanitize_model: bool = True, ) -> None: """Initialize PrecisionConverter. @@ -118,11 +119,18 @@ def __init__( tensor_block_dict: Dictionary of tensors (operation type and I/O indices) that should remain in FP32. use_standalone_type_inference: Use standalone type inference instead of ONNX's infer_shapes. original_network_io_metadata: Original public input/output metadata captured at the API boundary. + sanitize_model: Whether to sanitize the model before precision conversion. """ self.model = deepcopy(model) - self.value_info_map = value_info_map - self.initializer_map = initializer_map - self.node_to_init_map = node_to_init_map + self.sanitize_model = sanitize_model + if sanitize_model: + self.value_info_map = value_info_map + self.initializer_map = initializer_map + self.node_to_init_map = node_to_init_map + else: + self.value_info_map, self.initializer_map, self.node_to_init_map = utils.setup_mappings( + self.model + ) self.keep_io_types = keep_io_types self.init_conversion_max_bytes = ( np.inf if init_conversion_max_bytes is None else init_conversion_max_bytes @@ -195,7 +203,8 @@ def convert( "AutoCast can only operate on valid ONNX models, but the input model is invalid. See log for details." ) - self._sanitize_model() + if self.sanitize_model: + self._sanitize_model() # Filter out nodes that are not allowed to be in low precision # This is done here and not in NodeClassifier because it is required for the model to be valid diff --git a/modelopt/onnx/autocast/referencerunner.py b/modelopt/onnx/autocast/referencerunner.py index 6dcf7137689..4db54b3b366 100644 --- a/modelopt/onnx/autocast/referencerunner.py +++ b/modelopt/onnx/autocast/referencerunner.py @@ -26,6 +26,7 @@ import copy import io +import os import sys import tempfile from collections import OrderedDict @@ -112,8 +113,6 @@ def _load_inputs_from_npz(self, input_data_path): Returns: List of input dictionaries, one per batch. """ - import os - if os.path.isdir(input_data_path): # Load all NPZ files in the directory as multiple batches npz_files = sorted([f for f in os.listdir(input_data_path) if f.endswith(".npz")]) @@ -135,7 +134,7 @@ def _load_inputs_from_npz(self, input_data_path): def _validate_inputs(self, data_loader): """Validate that input names and shapes match the model.""" if isinstance(data_loader, list) and ( - isinstance(data_loader[0], (dict, np.lib.npyio.NpzFile)) + isinstance(data_loader[0], dict | np.lib.npyio.NpzFile) ): if sorted(self.input_names) != sorted(data_loader[0].keys()): raise ValueError("Input names from ONNX model do not match provided input names.") @@ -165,8 +164,6 @@ def _load_inputs(self, inputs): # If no inputs are provided, use random inputs data_loader = DataLoader(val_range={"": (-1, 1)}) - import os - if inputs is not None: if isinstance(inputs, str): if inputs.endswith(".json"): @@ -178,7 +175,7 @@ def _load_inputs(self, inputs): f"Invalid input file: {inputs}. Supported input types: .json (Polygraphy JSON format), " ".npz (Numpy), or a directory containing .npz files" ) - elif isinstance(inputs, (dict, OrderedDict)): + elif isinstance(inputs, dict | OrderedDict): data_loader = [inputs] else: raise ValueError( @@ -193,32 +190,32 @@ def _get_ort_runner(self, model): from polygraphy.backend.onnx import BytesFromOnnx from polygraphy.backend.onnxrt import OnnxrtRunner, SessionFromOnnx - # Check if model has external data by checking: - # 1. If any initializer has data_location set to EXTERNAL (even if data is loaded) - # 2. If model size would exceed 2GB (indicating need for external data) - needs_external_data = onnx_utils.check_model_uses_external_data( - self.model - ) or self.model.ByteSize() > 2 * (1024**3) - if needs_external_data: - logger.debug("Model has external data, using file-based approach") - # Get the actual ONNX ModelProto from ModifyOutputs wrapper - modified_model = model() - - # Use a persistent temp file, because we need the file to be present in an broader context - tmp_file = tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) - tmp_file.close() - tmp_file_path = tmp_file.name - onnx_utils.save_onnx(modified_model, tmp_file_path, save_as_external_data=True) - logger.debug(f"Model with all outputs saved to {tmp_file_path}") - build_onnxrt_session = SessionFromOnnx(tmp_file_path, providers=self.providers) + # Get the actual ONNX ModelProto from ModifyOutputs wrapper + modified_model = model() - else: - # For models without external data, use the original BytesFromOnnx approach (no tmp files) - logger.debug("Model has no external data, using BytesFromOnnx approach") - serialize_onnx = BytesFromOnnx(model) - build_onnxrt_session = SessionFromOnnx(serialize_onnx, providers=self.providers) - runners = [OnnxrtRunner(build_onnxrt_session)] - return runners + needs_file_backed_model = onnx_utils.check_model_uses_external_data( + modified_model + ) or onnx_utils.is_model_too_large_for_protobuf(modified_model) + model_temp_dir = None + try: + if needs_file_backed_model: + logger.debug("Model has external data, using file-based approach") + model_temp_dir = tempfile.TemporaryDirectory() + tmp_file_path = os.path.join(model_temp_dir.name, "model.onnx") + onnx_utils.save_onnx(modified_model, tmp_file_path, save_as_external_data=True) + logger.debug(f"Model with all outputs saved to {tmp_file_path}") + build_onnxrt_session = SessionFromOnnx(tmp_file_path, providers=self.providers) + else: + # For models without external data, use the original BytesFromOnnx approach (no tmp files) + logger.debug("Model has no external data, using BytesFromOnnx approach") + serialize_onnx = BytesFromOnnx(modified_model) + build_onnxrt_session = SessionFromOnnx(serialize_onnx, providers=self.providers) + runners = [OnnxrtRunner(build_onnxrt_session)] + except Exception: + if model_temp_dir is not None: + model_temp_dir.cleanup() + raise + return runners, model_temp_dir def _aggregate_tensor_stats(self, all_batch_data: list[OrderedDict]) -> OrderedDict: """Aggregate tensor statistics across multiple batches. @@ -300,22 +297,25 @@ def run(self, inputs=None): modify_outputs = ModifyOnnxOutputs(model_copy, outputs=constants.MARK_ALL) # Load the modified model and create an inference session - runners = self._get_ort_runner(modify_outputs) - - # Comparator is used despite the fact that we are using ONNXRuntime - # because it provides the ability to generate random inputs using DataLoader - data_loader = self._load_inputs(inputs) - - # Temporarily redirect stdout to suppress Comparator.run() output - stdout = sys.stdout - string_buffer = io.StringIO() - sys.stdout = string_buffer + runners, model_temp_dir = self._get_ort_runner(modify_outputs) try: - results = Comparator.run(runners, data_loader=data_loader) + # Comparator is used despite the fact that we are using ONNXRuntime + # because it provides the ability to generate random inputs using DataLoader + data_loader = self._load_inputs(inputs) + + # Temporarily redirect stdout to suppress Comparator.run() output + stdout = sys.stdout + string_buffer = io.StringIO() + sys.stdout = string_buffer + try: + results = Comparator.run(runners, data_loader=data_loader) + finally: + # Capture the output before restoring stdout + captured_output = string_buffer.getvalue() + sys.stdout = stdout finally: - # Capture the output before restoring stdout - captured_output = string_buffer.getvalue() - sys.stdout = stdout + if model_temp_dir is not None: + model_temp_dir.cleanup() if not results: logger.error(f"ONNXRuntime execution failed with output:\n{captured_output}") diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py index f8b5a41a41a..aa7d27bc8fd 100644 --- a/modelopt/onnx/utils.py +++ b/modelopt/onnx/utils.py @@ -18,9 +18,9 @@ import copy import io import os +import sys import tempfile import uuid -from collections import defaultdict from typing import Any import numpy as np @@ -527,50 +527,114 @@ def name_onnx_nodes(graph: onnx.GraphProto) -> bool: def duplicate_shared_constants(onnx_model: onnx.ModelProto) -> tuple[onnx.ModelProto, bool]: """Duplicate constant tensors if they are shared.""" - graph = gs.import_onnx(onnx_model) - name_dict = defaultdict(lambda: 0) - - def _get_unique_name(old_name): - name_dict[old_name] += 1 - return old_name + "_" + str(name_dict[old_name]) - - # Get tensors with shared constant inputs - tensors = [] - for node in graph.nodes: - for inp_idx, tensor in enumerate(node.inputs): - # constant is shared across multiple nodes - if isinstance(tensor, Constant) and len(tensor.outputs) > 1: - tensors.append({"tensor": tensor, "inp_node": node, "inp_idx": inp_idx}) - - # Duplicate shared tensors - for tensor_dict in tensors: - tensor = tensor_dict["tensor"] - new_tensor = Constant( - name=_get_unique_name(tensor.name), - values=tensor.values, + graph = onnx_model.graph + initializers = {initializer.name: initializer for initializer in graph.initializer} + use_counts = dict.fromkeys(initializers, 0) + for node in graph.node: + for input_name in node.input: + if input_name in use_counts: + use_counts[input_name] += 1 + + shared_names = {name for name, count in use_counts.items() if count > 1} + if not shared_names: + return onnx_model, False + + used_names = set(initializers) + used_names.update( + sparse_initializer.values.name for sparse_initializer in graph.sparse_initializer + ) + for value_info in list(graph.input) + list(graph.output) + list(graph.value_info): + used_names.add(value_info.name) + for node in graph.node: + used_names.update(name for name in node.input if name) + used_names.update(name for name in node.output if name) + + def _get_nested_graphs(node: onnx.NodeProto): + for attribute in node.attribute: + if attribute.type == onnx.AttributeProto.GRAPH: + yield attribute.g + elif attribute.type == onnx.AttributeProto.GRAPHS: + yield from attribute.graphs + + def _find_captured_names(nested_graph: onnx.GraphProto, outer_names: set[str]) -> set[str]: + local_names = {value_info.name for value_info in nested_graph.input} + local_names.update(initializer.name for initializer in nested_graph.initializer) + local_names.update( + sparse_initializer.values.name for sparse_initializer in nested_graph.sparse_initializer ) - tensor_dict["inp_node"].inputs[tensor_dict["inp_idx"]] = new_tensor + for nested_node in nested_graph.node: + local_names.update(name for name in nested_node.output if name) - onnx_model = gs.export_onnx(graph) - is_modified = bool(tensors) - return onnx_model, is_modified + visible_outer_names = outer_names - local_names + captured_names = { + output.name for output in nested_graph.output if output.name in visible_outer_names + } + for nested_node in nested_graph.node: + captured_names.update( + input_name for input_name in nested_node.input if input_name in visible_outer_names + ) + for child_graph in _get_nested_graphs(nested_node): + captured_names.update(_find_captured_names(child_graph, visible_outer_names)) + return captured_names + captured_names = set() + for node in graph.node: + for nested_graph in _get_nested_graphs(node): + captured_names.update(_find_captured_names(nested_graph, shared_names)) -def check_model(model: onnx.ModelProto) -> None: - """Checks if the given model is valid.""" - save_as_external_data = False + next_suffix: dict[str, int] = {} + + def _get_unique_name(old_name: str) -> str: + suffix = next_suffix.get(old_name, 1) + new_name = f"{old_name}_{suffix}" + while new_name in used_names: + suffix += 1 + new_name = f"{old_name}_{suffix}" + next_suffix[old_name] = suffix + 1 + used_names.add(new_name) + return new_name + + for node in graph.node: + for input_index, input_name in enumerate(node.input): + if input_name not in shared_names: + continue + + duplicated_initializer = graph.initializer.add() + duplicated_initializer.CopyFrom(initializers[input_name]) + duplicated_initializer.name = _get_unique_name(input_name) + node.input[input_index] = duplicated_initializer.name + + retained_names = captured_names | { + output.name for output in graph.output if output.name in shared_names + } + removed_initializer_names = shared_names - retained_names + for initializer in initializers.values(): + if initializer.name in removed_initializer_names: + graph.initializer.remove(initializer) + for graph_input in list(graph.input): + if graph_input.name in removed_initializer_names: + graph.input.remove(graph_input) + + return onnx_model, True + + +def is_model_too_large_for_protobuf(model: onnx.ModelProto) -> bool: + """Return whether a model cannot safely use an in-memory protobuf API.""" try: model_size = model.ByteSize() except Exception as e: logger.warning( "Failed to compute model size with ByteSize (%s). Using external data path.", e ) - save_as_external_data = True - else: - if model_size <= 0 or model_size > (2 * (1024**3)): - save_as_external_data = True + return True - if save_as_external_data: + max_model_size = onnx.checker.MAXIMUM_PROTOBUF - sys.getsizeof(b"") + return model_size <= 0 or model_size > max_model_size + + +def check_model(model: onnx.ModelProto) -> None: + """Checks if the given model is valid.""" + if is_model_too_large_for_protobuf(model): with tempfile.TemporaryDirectory() as temp_dir: # ONNX also looks in CWD, so we need to use a unique id unique_id = str(uuid.uuid4())[:8] @@ -1171,19 +1235,7 @@ def infer_types_verification(model: onnx.ModelProto) -> onnx.ModelProto: def infer_shapes(model: onnx.ModelProto, **kwargs): """Infers shapes of the onnx graph, handles large models.""" - save_as_external_data = False - try: - model_size = model.ByteSize() - except Exception as e: - logger.warning( - "Failed to compute model size with ByteSize (%s). Using external data path.", e - ) - save_as_external_data = True - else: - if model_size <= 0 or model_size > (2 * (1024**3)): - save_as_external_data = True - - if save_as_external_data: + if is_model_too_large_for_protobuf(model): with tempfile.TemporaryDirectory() as temp_dir: # ONNX also looks in CWD, so we need to use a unique id unique_id = str(uuid.uuid4())[:8] diff --git a/tests/unit/onnx/autocast/test_autocast.py b/tests/unit/onnx/autocast/test_autocast.py index ccad0c7201b..1377b20af9f 100644 --- a/tests/unit/onnx/autocast/test_autocast.py +++ b/tests/unit/onnx/autocast/test_autocast.py @@ -27,6 +27,7 @@ from modelopt.onnx.autocast import convert_to_mixed_precision from modelopt.onnx.autocast.__main__ import get_parser, main from modelopt.onnx.autocast.convert import convert_to_f16 +from modelopt.onnx.autocast.graphsanitizer import GraphSanitizer from modelopt.onnx.autocast.logging_config import configure_logging configure_logging("DEBUG") @@ -153,6 +154,42 @@ def test_convert_simple_model(temp_model_path, temp_output_path, keep_io_types): onnx.checker.check_model(loaded_model) +def test_convert_external_data_sanitizes_once_and_materializes_initializers( + tmp_path, simple_model, monkeypatch +): + model_path = tmp_path / "external_model.onnx" + onnx.save_model( + simple_model, + model_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location="external_model.data", + size_threshold=0, + ) + + sanitize_calls = [] + original_sanitize = GraphSanitizer.sanitize + + def record_sanitize(sanitizer): + initializer = sanitizer.model.graph.initializer[0] + sanitize_calls.append( + (sanitizer.onnx_path, initializer.data_location, bool(initializer.raw_data)) + ) + return original_sanitize(sanitizer) + + monkeypatch.setattr(GraphSanitizer, "sanitize", record_sanitize) + + converted_model = convert_to_mixed_precision(onnx_path=str(model_path), data_max=np.inf) + + assert sanitize_calls == [(str(model_path.resolve()), onnx.TensorProto.EXTERNAL, False)] + for initializer in converted_model.graph.initializer: + assert initializer.data_location != onnx.TensorProto.EXTERNAL + assert not initializer.external_data + assert initializer.raw_data + assert initializer.data_type == onnx.TensorProto.FLOAT16 + onnx.checker.check_model(converted_model) + + def assert_input_precision(nodes, dtype="float16"): for node in nodes: for inp in node.inputs: @@ -199,11 +236,6 @@ def test_conv_resize_conversion(tmp_path): # Convert the model converted_model = convert_to_mixed_precision(onnx_path=onnx_path) - # Output model should be produced in the same tmp_path - output_onnx_path = onnx_path.replace(".onnx", ".fp16.onnx") - onnx.save(converted_model, output_onnx_path) - - # Load the output model graph = gs.import_onnx(converted_model) # Check that Resize is correctly converted: @@ -213,6 +245,7 @@ def test_conv_resize_conversion(tmp_path): assert all(inp.dtype == np.float16 for inp in resize_node.inputs[0:2]), ( "Resize data and ROI inputs should be FP16" ) + assert resize_node.inputs[1].name != resize_node.inputs[2].name @pytest.mark.parametrize("target_opset", [13, 17, 19, 21]) diff --git a/tests/unit/onnx/autocast/test_graphsanitizer.py b/tests/unit/onnx/autocast/test_graphsanitizer.py index cb487b56bf0..f42324a6338 100644 --- a/tests/unit/onnx/autocast/test_graphsanitizer.py +++ b/tests/unit/onnx/autocast/test_graphsanitizer.py @@ -13,10 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest.mock import Mock + import numpy as np +import onnx import pytest from onnx import TensorProto, helper, numpy_helper +import modelopt.onnx.autocast.graphsanitizer as graphsanitizer from modelopt.onnx.autocast.graphsanitizer import GraphSanitizer @@ -407,3 +411,63 @@ def test_convert_fp64_no_changes_needed(): assert sanitizer._convert_fp64_initializers() is False assert sanitizer._convert_fp64_io_types() is False assert sanitizer._convert_fp64_nodes() is False + + +def test_sanitize_large_external_initializer_metadata(): + external_data_bytes = onnx.checker.MAXIMUM_PROTOBUF + 2048 + initializer = TensorProto( + name="weight", + data_type=TensorProto.FLOAT, + dims=[external_data_bytes // 4], + data_location=TensorProto.EXTERNAL, + ) + for key, value in ( + ("location", "weight.bin"), + ("offset", "0"), + ("length", str(external_data_bytes)), + ): + initializer.external_data.add(key=key, value=value) + + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, initializer.dims) + identity = helper.make_node("Identity", [initializer.name], [output.name]) + graph = helper.make_graph( + [identity], "large_external_initializer", [], [output], initializer=[initializer] + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 22)]) + model.ir_version = 10 + assert model.ByteSize() < 1024 + + sanitizer = GraphSanitizer(model, min_opset=22) + sanitizer.sanitize() + + initializer = sanitizer.model.graph.initializer[0] + assert initializer.data_location == TensorProto.EXTERNAL + assert not initializer.raw_data + assert {entry.key: entry.value for entry in initializer.external_data} == { + "location": "weight.bin", + "offset": "0", + "length": str(external_data_bytes), + } + assert [node.op_type for node in sanitizer.model.graph.node] == ["Identity"] + + +def test_find_custom_nodes_uses_source_model_path(tmp_path, monkeypatch): + x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1]) + y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1]) + custom_node = helper.make_node("CustomOp", [x.name], [y.name], name="custom") + graph = helper.make_graph([custom_node], "custom_graph", [x], [y]) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 22)]) + model_path = tmp_path / "custom.onnx" + tensor_info = {"Y": {"dtype": TensorProto.FLOAT, "shape": [1]}} + get_custom_layers = Mock(return_value=([custom_node.name], tensor_info)) + infer_types_shapes = Mock(return_value=model) + + monkeypatch.setattr(graphsanitizer, "set_trt_plugin_domain", Mock(return_value=model)) + monkeypatch.setattr(graphsanitizer, "get_custom_layers", get_custom_layers) + monkeypatch.setattr(graphsanitizer, "infer_types_shapes_tensorrt", infer_types_shapes) + + sanitizer = GraphSanitizer(model, min_opset=22, onnx_path=str(model_path)) + sanitizer.find_custom_nodes() + + get_custom_layers.assert_called_once_with(str(model_path.resolve()), []) + infer_types_shapes.assert_called_once_with(model, [], all_tensor_info=tensor_info) diff --git a/tests/unit/onnx/autocast/test_referencerunner.py b/tests/unit/onnx/autocast/test_referencerunner.py index 00075714c69..5c5c3c00ab2 100644 --- a/tests/unit/onnx/autocast/test_referencerunner.py +++ b/tests/unit/onnx/autocast/test_referencerunner.py @@ -17,9 +17,12 @@ import os import tempfile from collections import OrderedDict +from unittest.mock import Mock import numpy as np import onnx +import polygraphy.backend.onnx as polygraphy_onnx +import polygraphy.backend.onnxrt as polygraphy_onnxrt import pytest from onnx import TensorProto, helper @@ -67,6 +70,54 @@ def test_init(simple_model): assert isinstance(runner.model, onnx.ModelProto) +def test_get_ort_runner_uses_file_when_modified_byte_size_fails(monkeypatch, reference_runner): + modified_model = Mock() + modified_model.ByteSize.side_effect = ValueError("model size unavailable") + check_external_data = Mock(return_value=False) + save_onnx = Mock() + bytes_from_onnx = Mock(side_effect=AssertionError("Byte serialization should not be used")) + session_from_onnx = Mock(return_value="session") + onnxrt_runner = Mock(return_value="runner") + + monkeypatch.setattr(onnx_utils, "check_model_uses_external_data", check_external_data) + monkeypatch.setattr(onnx_utils, "save_onnx", save_onnx) + monkeypatch.setattr(polygraphy_onnx, "BytesFromOnnx", bytes_from_onnx) + monkeypatch.setattr(polygraphy_onnxrt, "SessionFromOnnx", session_from_onnx) + monkeypatch.setattr(polygraphy_onnxrt, "OnnxrtRunner", onnxrt_runner) + + runners, model_temp_dir = reference_runner._get_ort_runner(lambda: modified_model) + + assert model_temp_dir is not None + with model_temp_dir as model_temp_path: + save_onnx.assert_called_once() + model_path = save_onnx.call_args.args[1] + assert runners == ["runner"] + assert os.path.dirname(model_path) == model_temp_path + check_external_data.assert_called_once_with(modified_model) + modified_model.ByteSize.assert_called_once_with() + save_onnx.assert_called_once_with(modified_model, model_path, save_as_external_data=True) + bytes_from_onnx.assert_not_called() + session_from_onnx.assert_called_once_with(model_path, providers=reference_runner.providers) + onnxrt_runner.assert_called_once_with("session") + + +def test_run_cleans_model_tempdir_when_input_loading_fails(monkeypatch, reference_runner): + model_temp_dir = Mock() + monkeypatch.setattr( + reference_runner, "_get_ort_runner", Mock(return_value=([], model_temp_dir)) + ) + monkeypatch.setattr( + reference_runner, + "_load_inputs", + Mock(side_effect=ValueError("input loading failed")), + ) + + with pytest.raises(ValueError, match="input loading failed"): + reference_runner.run() + + model_temp_dir.cleanup.assert_called_once_with() + + def test_run_with_random_inputs(reference_runner): """Test running inference with random inputs.""" results = reference_runner.run() diff --git a/tests/unit/onnx/test_onnx_utils.py b/tests/unit/onnx/test_onnx_utils.py index 36face35b90..590a9f24186 100644 --- a/tests/unit/onnx/test_onnx_utils.py +++ b/tests/unit/onnx/test_onnx_utils.py @@ -14,6 +14,8 @@ # limitations under the License. import os +import sys +from unittest.mock import Mock import numpy as np import onnx @@ -28,6 +30,7 @@ make_tensor_value_info, ) +import modelopt.onnx.utils as onnx_utils from modelopt.onnx.trt_utils import load_onnx_model from modelopt.onnx.utils import ( clear_stale_value_info, @@ -56,6 +59,190 @@ def test_save_onnx(tmp_path): assert os.path.exists(os.path.join(tmp_path, "test.onnx")) +def _make_external_initializer( + name: str, location: str = "missing-shared-data.bin" +) -> onnx.TensorProto: + initializer = onnx.TensorProto(name=name, data_type=onnx.TensorProto.FLOAT, dims=[1]) + initializer.data_location = onnx.TensorProto.EXTERNAL + for key, value in [ + ("location", location), + ("offset", "0"), + ("length", "4"), + ]: + entry = initializer.external_data.add() + entry.key = key + entry.value = value + return initializer + + +def test_duplicate_shared_constants_preserves_and_materializes_external_data(tmp_path): + external_values = np.array([3.25], dtype=np.float32) + external_data_path = tmp_path / "shared.bin" + external_data_path.write_bytes(external_values.tobytes()) + shared = _make_external_initializer("shared", external_data_path.name) + collision = make_tensor("shared_1", onnx.TensorProto.FLOAT, [1], [0.0]) + sparse_collision = onnx.helper.make_sparse_tensor( + make_tensor("shared_2", onnx.TensorProto.FLOAT, [1], [0.0]), + make_tensor("", onnx.TensorProto.INT64, [1], [0]), + [1], + ) + nodes = [ + make_node("Add", ["input", "shared"], ["intermediate"]), + make_node("Add", ["intermediate", "shared"], ["output"]), + ] + graph = make_graph( + nodes, + "shared_external_initializer", + [ + make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1]), + make_tensor_value_info("shared", onnx.TensorProto.FLOAT, [1]), + ], + [make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1])], + initializer=[shared, collision], + ) + graph.sparse_initializer.append(sparse_collision) + model = make_model(graph) + + result, modified = onnx_utils.duplicate_shared_constants(model) + + assert result is model + assert modified + duplicated_names = ("shared_3", "shared_4") + assert [node.input[1] for node in result.graph.node] == list(duplicated_names) + initializers = {initializer.name: initializer for initializer in result.graph.initializer} + assert set(initializers) == {"shared_1", *duplicated_names} + assert result.graph.sparse_initializer[0].values.name == "shared_2" + assert {graph_input.name for graph_input in result.graph.input} == {"input"} + for name in duplicated_names: + initializer = initializers[name] + assert initializer.data_location == onnx.TensorProto.EXTERNAL + assert [(entry.key, entry.value) for entry in initializer.external_data] == [ + ("location", "shared.bin"), + ("offset", "0"), + ("length", "4"), + ] + assert not initializer.HasField("raw_data") + + onnx.external_data_helper.load_external_data_for_model(result, str(tmp_path)) + for name in duplicated_names: + np.testing.assert_array_equal( + onnx.numpy_helper.to_array(initializers[name]), external_values + ) + + +def test_duplicate_shared_constants_fast_path_returns_original_model(): + shared = _make_external_initializer("single_use") + graph = make_graph( + [make_node("Add", ["input", "single_use"], ["output"])], + "single_external_initializer", + [make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1])], + [make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1])], + initializer=[shared], + ) + model = make_model(graph) + serialized_model = model.SerializeToString() + + result, modified = onnx_utils.duplicate_shared_constants(model) + + assert result is model + assert not modified + assert result.SerializeToString() == serialized_model + + +def test_duplicate_shared_constants_retains_initializer_captured_by_subgraphs(): + def make_branch(name): + output = make_tensor_value_info("branch_output", onnx.TensorProto.FLOAT, [1]) + return make_graph([make_node("Identity", ["shared"], [output.name])], name, [], [output]) + + shared = make_tensor("shared", onnx.TensorProto.FLOAT, [1], [1.0]) + nodes = [ + make_node("Add", ["input", "shared"], ["left"]), + make_node("Add", ["input", "shared"], ["right"]), + make_node( + "If", + ["condition"], + ["branch_value"], + then_branch=make_branch("then_branch"), + else_branch=make_branch("else_branch"), + ), + make_node("Sum", ["left", "right", "branch_value"], ["output"]), + ] + graph = make_graph( + nodes, + "captured_initializer", + [ + make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1]), + make_tensor_value_info("condition", onnx.TensorProto.BOOL, []), + make_tensor_value_info("shared", onnx.TensorProto.FLOAT, [1]), + ], + [make_tensor_value_info("output", onnx.TensorProto.FLOAT, [1])], + initializer=[shared], + ) + model = make_model(graph, opset_imports=[make_opsetid("", 17)]) + + result, modified = onnx_utils.duplicate_shared_constants(model) + + assert modified + assert [node.input[1] for node in result.graph.node[:2]] == ["shared_1", "shared_2"] + assert {initializer.name for initializer in result.graph.initializer} == { + "shared", + "shared_1", + "shared_2", + } + assert "shared" in {graph_input.name for graph_input in result.graph.input} + onnx.checker.check_model(result) + + +def test_duplicate_shared_constants_retains_initializer_exposed_as_graph_output(): + shared = make_tensor("shared", onnx.TensorProto.FLOAT, [1], [1.0]) + graph = make_graph( + [ + make_node("Add", ["input", "shared"], ["left"]), + make_node("Add", ["input", "shared"], ["right"]), + ], + "initializer_graph_output", + [make_tensor_value_info("input", onnx.TensorProto.FLOAT, [1])], + [ + make_tensor_value_info("left", onnx.TensorProto.FLOAT, [1]), + make_tensor_value_info("right", onnx.TensorProto.FLOAT, [1]), + make_tensor_value_info("shared", onnx.TensorProto.FLOAT, [1]), + ], + initializer=[shared], + ) + model = make_model(graph) + + result, modified = onnx_utils.duplicate_shared_constants(model) + + assert modified + assert [node.input[1] for node in result.graph.node] == ["shared_1", "shared_2"] + assert {initializer.name for initializer in result.graph.initializer} == { + "shared", + "shared_1", + "shared_2", + } + assert [output.name for output in result.graph.output] == ["left", "right", "shared"] + onnx.checker.check_model(result) + + +@pytest.mark.parametrize( + ("model_size", "expected"), + [ + (1, False), + (onnx.checker.MAXIMUM_PROTOBUF - sys.getsizeof(b""), False), + (onnx.checker.MAXIMUM_PROTOBUF - sys.getsizeof(b"") + 1, True), + (0, True), + (ValueError("model size unavailable"), True), + ], +) +def test_is_model_too_large_for_protobuf(model_size, expected): + model = Mock() + if isinstance(model_size, Exception): + model.ByteSize.side_effect = model_size + else: + model.ByteSize.return_value = model_size + assert onnx_utils.is_model_too_large_for_protobuf(model) is expected + + def make_onnx_model_for_matmul_op(): input_left = np.array([1, 2]) input_right = np.array([1, 3]) From 78002183664b13435a18ac6b5fdda7cc0c98ee02 Mon Sep 17 00:00:00 2001 From: yeyu-nvidia Date: Thu, 10 Sep 2026 10:06:35 -0700 Subject: [PATCH 04/28] specdec: config_overrides for nested text_config checkpoints + load VLM-capable bases in merge_lora (#2289) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: New feature + bug fix Two related gaps, both hit while enabling EAGLE3 on a checkpoint whose config nests its text dims. **1. `config_overrides` for checkpoints whose `text_config` dims don't propagate.** Some multimodal checkpoints carry the real text-tower dims only under `config.text_config`, leaving the parent fields `None`. `from_pretrained` then builds a text tower with the wrong shape. `load_vlm_or_llm` gains an optional `config_overrides` dict applied to *both* the parent config and its `text_config` before instantiation, and the three entrypoints that load checkpoints — `ar_validate.py`, `export_hf_checkpoint.py`, `merge_lora.py` — get a `--config_overrides` passthrough. `main.py` threads it from `ModelArguments`. **2. `merge_lora.py` could not merge into any VLM base.** It loaded via `AutoModelForCausalLM`, which cannot load architectures absent from the CausalLM Auto map — every VLM base failed. It now goes through `load_vlm_or_llm`, which routes VLMs to `AutoModelForVision2Seq`/`AutoModelForImageTextToText` and plain LLMs to `AutoModelForCausalLM` with the same `dtype`/`device_map`, so LLM behavior is byte-for-byte unchanged. Also adds an optional `transformers_cosmos3` import so `cosmos3_omni` is registered with `AutoConfig` before use, and dispatches that `model_type` to its model class directly — that plugin registers only a *config*, never a model under `Auto*`, so `AutoModelForCausalLM` raised `KeyError('cosmos3_omni')` regardless of imports. The import is wrapped in `contextlib.suppress(ImportError)`, so it is a no-op when the plugin isn't installed. ### Usage ```bash # Checkpoint whose real dims live under config.text_config python examples/speculative_decoding/scripts/ar_validate.py \ --model_path --trust_remote_code \ --config_overrides '{"num_hidden_layers": 36, "intermediate_size": 12288, "num_key_value_heads": 8}' # Same flag on export and merge python examples/speculative_decoding/scripts/export_hf_checkpoint.py \ --model_path --export_path --config_overrides '{"num_hidden_layers": 36}' python examples/speculative_decoding/scripts/merge_lora.py \ --base_model_path --exported_lora_dir --output_path \ --config_overrides '{"num_hidden_layers": 36}' ``` ```python model = load_vlm_or_llm(path, config_overrides={"num_hidden_layers": 36}) # default None ``` ### Testing Exercised end-to-end on a Cosmos3-Nano (16B, 36-layer text tower) EAGLE3 LoRA run: - **Training** — the base loads with all 36 text layers and correct dims; two 4-epoch co-training runs completed (46,816 steps each). - **Export + merge** — produced `adapter_model.safetensors` and a merged base. Verified correct by per-layer weight diff: a `start_layer=18` run changed **exactly** layers 18-35, with layers 0-17 bit-identical to the base. - **AR validation** — `--config_overrides` loads the trained checkpoint; 80/80 MT-Bench samples, AR 3.42. - **Regression check** — `merge_lora` via `load_vlm_or_llm` produces a base loadable by `lm_eval`; ifeval/arc_challenge/winogrande all ran to completion. No local unit-test run: `nvidia-modelopt` isn't installed in my checkout, so `tests/conftest.py` fails to import. Relying on CI. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ — `config_overrides` defaults to `None`; the `merge_lora` loader swap keeps the same class, dtype and device_map for plain LLMs. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A — no new dependency; `transformers_cosmos3` is an optional import guarded by `contextlib.suppress`. - Did you write any new necessary tests?: ❌ — exercising these paths needs a checkpoint with a nested `text_config`, which the unit suite has no fixture for. Happy to add one if a reviewer can point me at a small suitable model. - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ❌ — can add a *Speculative Decoding* entry for the `merge_lora` VLM fix if you consider it changelog-worthy. - Did you get Claude approval on this PR?: ❌ — not yet run. ## Summary by CodeRabbit - **New Features** - Added JSON-based model configuration overrides across speculative decoding, training, validation, export, and LoRA workflows. - Overrides can update primary model and text configuration settings. - Expanded support for vision-language models and Cosmos3 Omni checkpoints. - **Bug Fixes** - Improved configuration handling for offline loading and checkpoint-based initialization. - Restored draft-model precision during checkpoint loading and model conversion. - Added validation for malformed, unsupported, and non-finite override values. - Standardized configuration override guidance across command-line workflows. --------- Signed-off-by: Ye Yu Co-authored-by: Claude Opus 4.8 (1M context) Signed-off-by: Chad Voegele --- examples/speculative_decoding/main.py | 6 +- .../scripts/ar_validate.py | 19 ++- .../scripts/export_hf_checkpoint.py | 19 ++- .../scripts/merge_lora.py | 40 +++++- .../speculative/plugins/hf_training_args.py | 5 + modelopt/torch/speculative/utils.py | 136 +++++++++++++++++- 6 files changed, 210 insertions(+), 15 deletions(-) diff --git a/examples/speculative_decoding/main.py b/examples/speculative_decoding/main.py index 46848203606..8e5acd91544 100644 --- a/examples/speculative_decoding/main.py +++ b/examples/speculative_decoding/main.py @@ -217,7 +217,10 @@ def train(): assert checkpoint is not None # guaranteed by checkpoint_is_hf with patch_transformers5_params_loading(): model = load_vlm_or_llm( - checkpoint, dtype="auto", trust_remote_code=recipe.model.trust_remote_code + checkpoint, + dtype="auto", + trust_remote_code=recipe.model.trust_remote_code, + config_overrides=recipe.model.config_overrides, ) tokenizer = transformers.AutoTokenizer.from_pretrained( checkpoint, trust_remote_code=recipe.model.trust_remote_code @@ -242,6 +245,7 @@ def train(): dtype="auto", device_map="cpu", trust_remote_code=recipe.model.trust_remote_code, + config_overrides=recipe.model.config_overrides, ) tokenizer = transformers.AutoTokenizer.from_pretrained( model_name_or_path, diff --git a/examples/speculative_decoding/scripts/ar_validate.py b/examples/speculative_decoding/scripts/ar_validate.py index 5699c480b7b..ccfb47dc085 100644 --- a/examples/speculative_decoding/scripts/ar_validate.py +++ b/examples/speculative_decoding/scripts/ar_validate.py @@ -28,7 +28,11 @@ import modelopt.torch.opt as mto from modelopt.torch.speculative.plugins.hf_eagle import HFARValidation -from modelopt.torch.speculative.utils import load_vlm_or_llm +from modelopt.torch.speculative.utils import ( + CONFIG_OVERRIDES_HELP, + load_vlm_or_llm, + parse_config_overrides, +) mto.enable_huggingface_checkpointing() @@ -100,11 +104,22 @@ def main(): default=None, help="Error if AR is below this threshold.", ) + parser.add_argument( + "--config_overrides", + type=str, + default=None, + help=CONFIG_OVERRIDES_HELP, + ) args = parser.parse_args() + config_overrides = parse_config_overrides(args.config_overrides) + accelerator = Accelerator() model = load_vlm_or_llm( - args.model_path, device_map="auto", trust_remote_code=args.trust_remote_code + args.model_path, + device_map="auto", + trust_remote_code=args.trust_remote_code, + config_overrides=config_overrides, ) tokenizer = AutoTokenizer.from_pretrained( args.model_path, trust_remote_code=args.trust_remote_code diff --git a/examples/speculative_decoding/scripts/export_hf_checkpoint.py b/examples/speculative_decoding/scripts/export_hf_checkpoint.py index cee2e45d0eb..2b86a47004b 100644 --- a/examples/speculative_decoding/scripts/export_hf_checkpoint.py +++ b/examples/speculative_decoding/scripts/export_hf_checkpoint.py @@ -21,7 +21,11 @@ import modelopt.torch.opt as mto from modelopt.torch.export import export_speculative_decoding -from modelopt.torch.speculative.utils import load_vlm_or_llm +from modelopt.torch.speculative.utils import ( + CONFIG_OVERRIDES_HELP, + load_vlm_or_llm, + parse_config_overrides, +) def parse_args(): @@ -33,13 +37,24 @@ def parse_args(): parser.add_argument( "--export_path", type=str, default="Destination directory for exported files." ) + parser.add_argument( + "--config_overrides", + type=str, + default=None, + help=CONFIG_OVERRIDES_HELP, + ) return parser.parse_args() mto.enable_huggingface_checkpointing() args = parse_args() -model = load_vlm_or_llm(args.model_path, dtype="auto", trust_remote_code=args.trust_remote_code) +model = load_vlm_or_llm( + args.model_path, + dtype="auto", + trust_remote_code=args.trust_remote_code, + config_overrides=parse_config_overrides(args.config_overrides), +) model.eval() with torch.inference_mode(): export_speculative_decoding( diff --git a/examples/speculative_decoding/scripts/merge_lora.py b/examples/speculative_decoding/scripts/merge_lora.py index 25d311393b5..4580e166832 100644 --- a/examples/speculative_decoding/scripts/merge_lora.py +++ b/examples/speculative_decoding/scripts/merge_lora.py @@ -31,7 +31,13 @@ from pathlib import Path from safetensors.torch import load_file -from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers import AutoTokenizer + +from modelopt.torch.speculative.utils import ( + CONFIG_OVERRIDES_HELP, + load_vlm_or_llm, + parse_config_overrides, +) def parse_args(): @@ -61,11 +67,18 @@ def parse_args(): action="store_true", help="Allow loading models that define custom code on the HF Hub. Off by default.", ) + parser.add_argument( + "--config_overrides", + type=str, + default=None, + help=CONFIG_OVERRIDES_HELP, + ) return parser.parse_args() def main(): args = parse_args() + config_overrides = parse_config_overrides(args.config_overrides) lora_dir = Path(args.exported_lora_dir) # Verify exported files exist (standard peft naming) @@ -81,13 +94,21 @@ def main(): print(f"Loaded {len(lora_sd)} LoRA tensors from {lora_dir}") print(f" Sample keys: {list(lora_sd.keys())[:4]}") - # Load the original base model + # Load the original base model. + # + # Use load_vlm_or_llm rather than AutoModelForCausalLM directly: it falls back to + # AutoModelForCausalLM for plain LLMs (same dtype/device_map, so unchanged behavior), but also + # handles VLMs and registers/loads architectures the Auto* maps don't cover. Cosmos3 is the + # motivating case -- the transformers-cosmos3 plugin registers only the `cosmos3_omni` config, + # never a model under Auto*, so AutoModelForCausalLM raises KeyError('cosmos3_omni') no matter + # what is imported. print(f"Loading base model from {args.base_model_path}...") - model = AutoModelForCausalLM.from_pretrained( + model = load_vlm_or_llm( args.base_model_path, - torch_dtype="auto", + dtype="auto", device_map="cpu", trust_remote_code=args.trust_remote_code, + config_overrides=config_overrides, ) tokenizer = AutoTokenizer.from_pretrained( args.base_model_path, trust_remote_code=args.trust_remote_code @@ -135,9 +156,18 @@ def main(): # Since LoRA only changes weights — not architecture — the original config is correct. import shutil + # ...but only when the loaded config still matches the base's. With --config_overrides the + # weights were built from corrected dims, so copying the uncorrected base config back would + # leave config.json disagreeing with model.safetensors and force every downstream reader to + # re-supply the same overrides. base_config = Path(args.base_model_path) / "config.json" output_config = Path(args.output_path) / "config.json" - if base_config.exists(): + if config_overrides: + print( + " Keeping the saved config.json (config_overrides were applied, so the original " + "base config would not match the merged weights)" + ) + elif base_config.exists(): shutil.copy2(str(base_config), str(output_config)) print(f" Restored original config.json from {base_config}") diff --git a/modelopt/torch/speculative/plugins/hf_training_args.py b/modelopt/torch/speculative/plugins/hf_training_args.py index 38d3f483e6b..b121d9f260d 100644 --- a/modelopt/torch/speculative/plugins/hf_training_args.py +++ b/modelopt/torch/speculative/plugins/hf_training_args.py @@ -44,6 +44,11 @@ class ModelArguments(BaseModel): model_name_or_path: str | None = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" use_fake_base_for_offline: bool = False trust_remote_code: bool = False + # Optional config field overrides applied to the loaded model config (and its + # text_config) before instantiation. Needed for checkpoints whose config doesn't + # round-trip cleanly through transformers (e.g. Cosmos3's Qwen3-VL text tower, + # where intermediate_size/num_key_value_heads don't propagate from text_config). + config_overrides: dict | None = None class DataArguments(BaseModel): diff --git a/modelopt/torch/speculative/utils.py b/modelopt/torch/speculative/utils.py index 9fa8fde5e15..c419d24588f 100644 --- a/modelopt/torch/speculative/utils.py +++ b/modelopt/torch/speculative/utils.py @@ -18,6 +18,7 @@ import contextlib import copy import importlib.util +import json import os import sys import warnings @@ -36,6 +37,11 @@ KIMI_K2_PACKAGE_NAME = "kimi_k2_temp" +# Attributes under which a checkpoint may nest its text-tower config. Mirrors +# modelopt.torch.speculative.plugins.modeling_fakebase._VLM_CONFIG_ATTRS. +NESTED_CONFIG_ATTRS = ["text_config", "llm_config"] + + REMOVE_THINK_CHAT_TEMPLATE = ( "{% if '' in content %}{% set content = content.split('')[-1] %}{% endif %}" ) @@ -584,6 +590,39 @@ def enable_cp_ttt_patch(cp_size: int = 1): modelopt.torch.speculative.plugins.hf_eagle.ENABLE_CP_TTT_PATCH = False +CONFIG_OVERRIDES_HELP = ( + "JSON object of config fields to override on the model config and its text_config before " + "instantiation, e.g. '{\"num_hidden_layers\": 36}'. Needed for checkpoints whose nested " + "text_config dims don't propagate to the parent config." +) + + +def parse_config_overrides(raw: str | None) -> dict | None: + """Parse a ``--config_overrides`` CLI value into a dict, or ``None`` if not supplied. + + Rejects malformed JSON and non-object payloads here, with an actionable message, rather than + letting them surface later as a raw ``JSONDecodeError`` or an ``AttributeError`` from deep + inside model loading. + """ + if not raw: + return None + try: + # Reject the JSON5-ish constants Python's json accepts by default: NaN/Infinity + # would sail through as floats and land on a config field as a dimension. + def _reject(const): + raise ValueError(f"--config_overrides contains non-finite value {const!r}") + + parsed = json.loads(raw, parse_constant=_reject) + except json.JSONDecodeError as e: + raise ValueError(f"--config_overrides is not valid JSON: {e}") from e + if not isinstance(parsed, dict): + raise ValueError( + f"--config_overrides must be a JSON object mapping field names to values, got " + f"{type(parsed).__name__}: {raw!r}" + ) + return parsed + + def load_vlm_or_llm( model_name_or_path: str, use_fake_base: bool = False, @@ -591,6 +630,7 @@ def load_vlm_or_llm( dtype: str | torch.dtype | None = None, device_map: str | None = None, trust_remote_code: bool = False, + config_overrides: dict | None = None, ): """Load a VLM or LLM. Returns the model. @@ -605,26 +645,84 @@ def load_vlm_or_llm( dtype: dtype to use when loading the model. device_map: Device map passed to ``from_pretrained``. trust_remote_code: Whether to trust remote code. + config_overrides: Optional config field overrides applied to the model config and + its ``text_config`` before instantiation (e.g. to correct dims that don't + propagate from a checkpoint's nested text config). """ + + def _warn_overrides_on_fake_base(): + # FakeBaseModel.from_source re-reads the checkpoint config itself, so overrides applied + # here never reach it -- but it is not silently wrong: from_source resolves dims from the + # nested text_config/llm_config first (modeling_fakebase._VLM_CONFIG_ATTRS), which is the + # very problem config_overrides exists to work around, so this path is already correct + # without them. Warn rather than raise: main.py forwards config_overrides unconditionally, + # so hard-failing would leave a single recipe unable to run offline at all. + if config_overrides: + warnings.warn( + "config_overrides is ignored on the FakeBaseModel path: from_source rebuilds the " + "config from the checkpoint, reading dims from the nested text_config/llm_config " + "directly, so the overrides are not needed there.", + stacklevel=2, + ) + if use_offline_training and use_fake_base: + _warn_overrides_on_fake_base() from modelopt.torch.speculative.plugins.modeling_fakebase import FakeBaseModel return FakeBaseModel.from_source(model_name_or_path, trust_remote_code=trust_remote_code) + # Import the transformers-cosmos3 plugin if available: it registers the `cosmos3_omni` + # architecture with AutoConfig on import, so the from_pretrained below recognizes it. + with contextlib.suppress(ImportError): + import transformers_cosmos3 # noqa: F401 + model_config = transformers.AutoConfig.from_pretrained( model_name_or_path, trust_remote_code=trust_remote_code, ) + # Apply caller-supplied config corrections to the parent config and every nested config + # (some checkpoints don't propagate the nested dims up to the parent). + if config_overrides: + # Cover every nested attribute VLM detection accepts, not just text_config: an + # llm_config-nesting checkpoint mirrors the fields on the parent as None, so an override + # would land on the parent, count as "applied", and never reach the real text tower. + targets = [ + cfg + for cfg in ( + model_config, + *(getattr(model_config, a, None) for a in NESTED_CONFIG_ATTRS), + ) + if cfg is not None + ] + unmatched = [] + for key, value in config_overrides.items(): + applied = False + for cfg_obj in targets: + if hasattr(cfg_obj, key): + setattr(cfg_obj, key, value) + applied = True + if not applied: + unmatched.append(key) + if unmatched: + # Silently skipping a key would hand back a wrong-shaped model while appearing to + # have applied the override -- the exact failure this option exists to correct. + raise ValueError( + f"config_overrides key(s) {sorted(unmatched)} matched no field on the model " + f"config (model_type={getattr(model_config, 'model_type', None)!r}) or its " + "text_config. Check for typos." + ) + # Detect VLMs: either "vl" in model_type (e.g. "llava") or has a nested text config # (e.g. Mistral3Config with model_type="mistral3" and text_config attribute). _is_vlm = "vl" in model_config.model_type.lower() or any( - getattr(model_config, attr, None) is not None for attr in ["text_config", "llm_config"] + getattr(model_config, attr, None) is not None for attr in NESTED_CONFIG_ATTRS ) if _is_vlm and use_offline_training: # For VLMs in offline training, FakeBaseModel loads only embed_tokens + lm_head # and auto-detects VLM weight key layouts (e.g. "language_model.model.embed_tokens"). + _warn_overrides_on_fake_base() from modelopt.torch.speculative.plugins.modeling_fakebase import FakeBaseModel return FakeBaseModel.from_source(model_name_or_path, trust_remote_code=trust_remote_code) @@ -641,13 +739,41 @@ def load_vlm_or_llm( model_cls = transformers.AutoModelForCausalLM extra = {} + + # Cosmos3 omni checkpoints: the transformers-cosmos3 plugin registers only the config + # (cosmos3_omni) with AutoConfig, never a model under the Auto* maps, so dispatch to its + # Cosmos3ForConditionalGeneration (a Qwen3-VL subclass) directly. The unused vision + # tower has mismatched dims vs the text-only use, so ignore those on load. + if getattr(model_config, "model_type", None) == "cosmos3_omni": + from transformers_cosmos3 import Cosmos3ForConditionalGeneration + + model_cls = Cosmos3ForConditionalGeneration + extra["ignore_mismatched_sizes"] = True + + # Pass our config object only when we had to modify it (overrides) or when the model class + # needs the plugin-built config; otherwise let from_pretrained build its own, exactly as before. + pass_config = bool(config_overrides) or "ignore_mismatched_sizes" in extra + + # Capture the true depth before any zeroing below, since it is restored after load. + orig_num_hidden_layers = getattr(model_config, "num_hidden_layers", None) + if use_offline_training: - extra["num_hidden_layers"] = 0 - if hasattr(model_config, "layer_types"): - extra["layer_types"] = [] + if pass_config: + # from_pretrained only forwards unrecognized kwargs into the config when it builds + # that config itself. Given a PretrainedConfig instance it deep-copies it and leaves + # the rest in model_kwargs, so num_hidden_layers=0 would never reach the config and + # the full model would be materialized. Set the fields on the config directly. + model_config.num_hidden_layers = 0 + if hasattr(model_config, "layer_types"): + model_config.layer_types = [] + else: + extra["num_hidden_layers"] = 0 + if hasattr(model_config, "layer_types"): + extra["layer_types"] = [] model = model_cls.from_pretrained( model_name_or_path, + config=model_config if pass_config else None, trust_remote_code=trust_remote_code, torch_dtype=dtype, device_map=device_map, @@ -656,7 +782,7 @@ def load_vlm_or_llm( if use_offline_training: # Preserve the original layer count since we loaded with num_hidden_layers=0 - model.config.num_orig_hidden_layers = model_config.num_hidden_layers + model.config.num_orig_hidden_layers = orig_num_hidden_layers return model From d19c7d63bc41a53da37fd3fa26e2999316cb3054 Mon Sep 17 00:00:00 2001 From: Ajinkya Rasane <131806219+ajrasane@users.noreply.github.com> Date: Thu, 10 Sep 2026 14:30:26 -0400 Subject: [PATCH 05/28] [6508436] Fix BF16 FP8 ONNX export (#2314) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix Fix FP8 ONNX export for BF16 models during real-weight compression without changing the public API or the `weights_dtype="fp32"` default. The FP8 exporter preserves BF16 initializer bits when bridging GraphSurgeon NumPy arrays to Torch, widens BF16 values exactly to FP32 for normalization, and leaves existing FP16/FP32 handling unchanged. Conv scales and dequantized outputs retain the source dtype, and scales round upward when needed so serialized values cannot cause FP8 overflow. `weights_dtype="bf16"` is accepted as a no-op only for FP8-only models whose floating parameters are all BF16. Registered buffers do not affect this weight-focused decision and may preserve higher-precision regions in the exported graph. Unsupported BF16 FP8-to-FP16 and FP32 or mixed-parameter-to-BF16 conversions are rejected with `ValueError` before temporary export paths are created. A narrow GraphSurgeon fix preserves integer BF16 value-info dtypes. ### Usage ```python onnx_bytes, metadata = get_onnx_bytes_and_metadata( quantized_fp8_model, (sample_input,), weights_dtype="bf16", onnx_opset=23, ) ``` ### Testing - Seven focused CPU regressions passed: BF16 QDQ compression and integer dtype handling, BF16-to-BF16 and FP32-to-FP16 Conv/Linear export, and four unsupported-conversion cases. - QDQ utilities: 31 passed; pytest 2.25s, wall 29.88s. - FP8 MHA exporter: 6 passed; pytest 2.05s, wall 32.71s. - Torch deploy utilities: 51 passed; pytest 8.94s, wall 25.74s. - Torch ONNX CPU export: 36 passed; pytest 4.85s, wall 32.38s. - Changed-file pre-commit hooks: all passed; wall 8.07s. - Exact-head FP8 BF16 GPU workflow at `f21d62a`: exit code 0; ONNX checker passed; 6 FP8 initializers, 3 native `DequantizeLinear` nodes, and 12 BF16 initializers. - Refreshed GitHub CI at `f21d62a`: 50 passed and 1 skipped. Unit, GPU, and regression required aggregates and Codecov passed. Two ONNX example leaves failed because the runner could not load a cuDNN sublibrary; their dependent example aggregate consequently failed. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ - 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)?: ✅ ### Additional Information - TODO: Deliver authoritative `native`/FP32/FP16/BF16 ONNX export across all quantized formats in follow-up pull requests. > 🤖 _Generated by Codex (AI agent)._ --------- Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> Co-authored-by: Codex Co-authored-by: Codex Signed-off-by: Chad Voegele --- CHANGELOG.rst | 1 + modelopt/onnx/export/fp8_exporter.py | 40 +++-- modelopt/onnx/quantization/gs_patching.py | 12 +- modelopt/torch/_deploy/utils/torch_onnx.py | 65 +++++--- .../unit/onnx/quantization/test_qdq_utils.py | 61 +++++++- .../deploy/utils/test_torch_onnx_utils.py | 142 ++++++++++++++++++ 6 files changed, 286 insertions(+), 35 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 41349796945..ceda42b6c96 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -59,6 +59,7 @@ Changelog - Fix ONNX AutoCast failing on models with external initializers larger than 2 GiB. - Avoid querying CUDA/Blackwell capability when ``NVFP4QTensor.quantize`` uses its CPU path or has the optional TensorRT-LLM fast path disabled. - Fix NVFP4 ONNX export to quantize FP4 weights with the published FP8 block scales, matching eager ModelOpt packed weights. Block scales below ``2**-9`` are now clamped to that minimum, and non-finite or negative scales raise an error. +- Fix FP8 ONNX export of BF16 models during real-weight compression. - Fix Megatron-Bridge Quantization Aware Distillation of a vision-language model silently discarding the ModelOpt state, so the distilled checkpoint restored no quantizers and exported as an unquantized model. Re-run QAD to regenerate any affected checkpoint. - Fix Megatron-Core HuggingFace export silently omitting fused (grouped GEMM) MoE experts for architectures without an ``experts.linear_fc1`` rule (e.g. ``Qwen3MoeForCausalLM``), which produced a valid-looking checkpoint containing no expert weights. The exporter now raises instead of writing that checkpoint; the scripts also avoid the situation by selecting ``SequentialMLP`` for those architectures. - Fix GatedDeltaNet (Qwen3.5) quantizer exclusions on Megatron-Core: the recipe patterns name the HuggingFace ``linear_attn`` module, so the ``conv1d`` was calibrated and the alpha / beta gate projections were exported in FP8. ``conv1d`` now has a ``self_attention`` alias in the default disabled-quantizer units, and the alpha / beta projections are exported in BF16 (they share Megatron's fused ``in_proj`` quantizer and cannot be disabled by name). diff --git a/modelopt/onnx/export/fp8_exporter.py b/modelopt/onnx/export/fp8_exporter.py index 427a7791f3b..8e94f40f332 100644 --- a/modelopt/onnx/export/fp8_exporter.py +++ b/modelopt/onnx/export/fp8_exporter.py @@ -17,6 +17,7 @@ import time +import ml_dtypes import numpy as np import onnx import onnx_graphsurgeon as gs @@ -33,6 +34,13 @@ _FP8_E4M3_SOFTMAX_SCALE = 1.0 / _FP8_E4M3_MAX +def _torch_from_numpy_for_fp8(array: np.ndarray) -> torch.Tensor: + """Convert a NumPy array to the PyTorch dtype used for FP8 normalization.""" + if array.dtype == ml_dtypes.bfloat16: + return torch.from_numpy(array.view(np.int16)).view(torch.bfloat16).float() + return torch.from_numpy(array) + + class FP8QuantExporter(ONNXQuantExporter): """Exporter for FP8 quantization.""" @@ -48,7 +56,7 @@ def compute_scales(onnx_model: onnx.ModelProto) -> onnx.ModelProto: @staticmethod def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: - """Compresses FP32/FP16 weights to FP8 by folding QDQ nodes to DQ only. + """Compresses FP32/FP16/BF16 weights to FP8 by folding QDQ nodes to DQ only. Even though modelopt supports FP8 onnx export, the weights are represented in fp32 + QDQ. The storage is therefore very bad. In this function, @@ -56,7 +64,7 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: weights in the output model. TRT custom ops are converted to native ONNX DequantizeLinear. Parameters: - onnx_model: ONNX model with FP32/FP16 weights and TRT_FP8 QDQ nodes. + onnx_model: ONNX model with FP32/FP16/BF16 weights and TRT_FP8 QDQ nodes. Returns: ONNX model with FP8 weights and native ONNX DQ nodes for weights (QDQ preserved for activations). @@ -78,8 +86,8 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: weights = node.inputs[0] scale = node.inputs[1] - torch_weights = torch.from_numpy(weights.values) - torch_scale = torch.from_numpy(scale.values) + torch_weights = _torch_from_numpy_for_fp8(weights.values) + torch_scale = _torch_from_numpy_for_fp8(scale.values) quantizer_name = scale.name.rsplit("/", 1)[0] dq_op = node.outputs[0].outputs[0] if dq_op.op != "TRT_FP8DequantizeLinear": @@ -194,14 +202,25 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: if any(out.op == "DequantizeLinear" for out in weight_input.outputs): continue - torch_weights = torch.from_numpy(weight_input.values.copy()) + torch_weights = _torch_from_numpy_for_fp8(weight_input.values.copy()) amax = torch_weights.abs().max().float() if amax == 0: continue scale_val = (amax / _FP8_E4M3_MAX).item() + scale_data = np.array(scale_val, dtype=weight_input.values.dtype) + # Round up so normalizing by the serialized scale stays within the FP8 range. + if scale_data < scale_val: + np.nextafter( + scale_data, + np.array(np.inf, dtype=scale_data.dtype), + out=scale_data, + ) + torch_scale = _torch_from_numpy_for_fp8(scale_data) # Quantize weights to FP8 (WAR: numpy doesn't support fp8) - fp8_data = (torch_weights / scale_val).to(torch.float8_e4m3fn).view(torch.uint8).numpy() + fp8_data = ( + (torch_weights / torch_scale).to(torch.float8_e4m3fn).view(torch.uint8).numpy() + ) fp8_tensor = onnx.TensorProto() fp8_tensor.data_type = onnx.TensorProto.FLOAT8E4M3FN fp8_tensor.dims.extend(fp8_data.shape) @@ -210,13 +229,16 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: node.name + "/weight_quantizer/fp8_weights", LazyValues(fp8_tensor) ) - # Scale in FP16 — DQ output type matches scale dtype, must match activation type scale_constant = gs.Constant( node.name + "/weight_quantizer/scale", - np.array(scale_val, dtype=np.float16), + scale_data, ) - dq_output = gs.Variable(node.name + "/weight_quantizer/dq_output") + dq_output = gs.Variable( + node.name + "/weight_quantizer/dq_output", + scale_data.dtype, + weight_input.values.shape, + ) dq_node = gs.Node( op="DequantizeLinear", name=node.name + "/weight_quantizer/DequantizeLinear", diff --git a/modelopt/onnx/quantization/gs_patching.py b/modelopt/onnx/quantization/gs_patching.py index a0eea84951e..503a91f9721 100644 --- a/modelopt/onnx/quantization/gs_patching.py +++ b/modelopt/onnx/quantization/gs_patching.py @@ -101,9 +101,15 @@ def _export_value_info_proto(tensor: gs.Variable, do_type_check: bool) -> onnx.V ) if tensor.dtype is not None: - dtype = getattr( - tensor, "explicit_dtype", onnx.helper.np_dtype_to_tensor_dtype(np.dtype(tensor.dtype)) - ) + dtype = getattr(tensor, "explicit_dtype", None) + if dtype is None: + dtype = tensor.dtype + if isinstance(dtype, (int, np.integer)): + dtype = int(dtype) + if dtype not in onnx.TensorProto.DataType.values(): + raise ValueError(f"Unknown ONNX tensor dtype for {tensor.name}: {dtype}") + else: + dtype = onnx.helper.np_dtype_to_tensor_dtype(np.dtype(dtype)) onnx_tensor = onnx.helper.make_tensor_value_info(tensor.name, dtype, tensor.shape) else: onnx_tensor = onnx.helper.make_empty_tensor_value_info(tensor.name) diff --git a/modelopt/torch/_deploy/utils/torch_onnx.py b/modelopt/torch/_deploy/utils/torch_onnx.py index 01fb754bbae..8e16813bd3a 100644 --- a/modelopt/torch/_deploy/utils/torch_onnx.py +++ b/modelopt/torch/_deploy/utils/torch_onnx.py @@ -507,14 +507,17 @@ def get_onnx_bytes_and_metadata( `torch.onnx.export `_. onnx_opset: The onnx opset version to use for exporting the model. dq_only: If True, the exported onnx model is converted to a dq_only model. - weights_dtype: The dtype of the weights in the onnx model. + weights_dtype: Requested high-precision dtype for exported weights. For an FP8 model, + ``"bf16"`` is accepted only when every floating parameter is already BF16. This is + a weight-focused no-op, not a graph-wide conversion: floating buffers are not + considered for eligibility and may preserve higher-precision regions. Returns: bytes: Onnx model in bytes. ModelMetadata: The model's meta data. Raises: - ValueError: If nn.Module is not passed as model. + ValueError: If model is not an nn.Module or the requested precision conversion is unsupported. """ if not isinstance(model, nn.Module): raise ValueError("Only PyTorch model compilation is supported.") @@ -527,6 +530,22 @@ def get_onnx_bytes_and_metadata( if isinstance(model, (DataParallel, DistributedDataParallel)): model = model.module + source_parameter_dtypes = { + parameter.dtype for parameter in model.parameters() if parameter.is_floating_point() + } + source_parameter_dtype_names = ", ".join(sorted(map(str, source_parameter_dtypes))) or "none" + uses_fp4 = is_fp4_quantized(model) + uses_mxfp8 = is_mxfp8_quantized(model) + uses_fp8 = is_fp8_quantized(model) + uses_int8 = is_int8_quantized(model) + uses_other_unsupported_quantizer = is_int4_quantized(model) or uses_mxfp8 or uses_int8 + is_bf16_fp8_noop = ( + weights_dtype == "bf16" + and source_parameter_dtypes == {torch.bfloat16} + and uses_fp8 + and not (uses_fp4 or uses_other_unsupported_quantizer) + ) + # Standardize model args and also tensorize them so they also appear in the onnx graph! # Floats/ints are tensorized when they are provided, but not tensorized when they are not # provided which is somewhat inconsistent (we always tensorize them!) @@ -549,11 +568,7 @@ def get_onnx_bytes_and_metadata( input_none_names = list(set(tree_spec_input.names) - set(input_names)) use_torch_autocast = not ( - is_fp4_quantized(model) - or is_mxfp8_quantized(model) - or is_fp8_quantized(model) - or is_int8_quantized(model) - or weights_dtype == "fp32" + uses_fp4 or uses_mxfp8 or uses_fp8 or uses_int8 or weights_dtype == "fp32" ) autocast = torch.autocast("cuda") if use_torch_autocast else nullcontext() @@ -575,6 +590,22 @@ def get_onnx_bytes_and_metadata( ) return onnx_model.to_bytes(), model_metadata + if weights_dtype == "fp16" and uses_fp8 and torch.bfloat16 in source_parameter_dtypes: + raise ValueError( + "Converting a BF16 FP8 ONNX graph to FP16 is not supported yet " + f"(source parameter dtypes: {source_parameter_dtype_names})" + ) + + if ( + weights_dtype == "bf16" + and (uses_fp8 or uses_other_unsupported_quantizer) + and not is_bf16_fp8_noop + ): + raise ValueError( + "Converting a quantized ONNX graph to BF16 is not supported yet " + f"(source parameter dtypes: {source_parameter_dtype_names})" + ) + # Export onnx model from pytorch model # As the maximum size of protobuf is 2GB, we cannot use io.BytesIO() buffer during export. model_name = model_name or model.__class__.__name__ @@ -583,16 +614,12 @@ def get_onnx_bytes_and_metadata( # Configure quantizers if the model is quantized in NVFP4 or MXFP8 mode quantizer_context = ( - configure_linear_module_onnx_quantizers(model) - if is_fp4_quantized(model) or is_mxfp8_quantized(model) - else nullcontext() + configure_linear_module_onnx_quantizers(model) if uses_fp4 or uses_mxfp8 else nullcontext() ) # Disable FP8 Conv weight quantizers: TorchScript custom ops produce outputs with # unknown shapes, causing _convolution symbolic to fail. Conv weights are quantized # to FP8 in post-processing by FP8QuantExporter instead. - conv_wq_context = ( - _disable_fp8_conv_weight_quantizers(model) if is_fp8_quantized(model) else nullcontext() - ) + conv_wq_context = _disable_fp8_conv_weight_quantizers(model) if uses_fp8 else nullcontext() with torch.inference_mode(), autocast, quantizer_context, conv_wq_context: additional_kwargs = {} if not dynamo_export: @@ -634,14 +661,8 @@ def get_onnx_bytes_and_metadata( if dq_only: onnx_opt_graph = qdq_to_dq(onnx_opt_graph) - if weights_dtype in ["fp16", "bf16"]: - if ( - is_int4_quantized(model) - or is_mxfp8_quantized(model) - or is_fp8_quantized(model) - or is_int8_quantized(model) - ): - assert weights_dtype == "fp16", "BF16 + MXFP8/INT4 mixed precision is not supported yet" + if weights_dtype in ["fp16", "bf16"] and not is_bf16_fp8_noop: + if uses_other_unsupported_quantizer or uses_fp8: onnx_opt_graph = convert_float_to_float16( onnx_opt_graph, keep_io_types=False, @@ -665,7 +686,7 @@ def get_onnx_bytes_and_metadata( onnx_opt_graph = remove_redundant_casts(onnx_opt_graph) # Remove Cast nodes around Q/DQ for optimal TRT fusion - if is_fp8_quantized(model): + if uses_fp8: onnx_opt_graph = fold_q_fp16_to_fp32_casts(onnx_opt_graph) onnx_opt_graph = fold_dq_fp32_to_fp16_casts(onnx_opt_graph) diff --git a/tests/unit/onnx/quantization/test_qdq_utils.py b/tests/unit/onnx/quantization/test_qdq_utils.py index 4b1e69ec538..5cdaf9203d0 100644 --- a/tests/unit/onnx/quantization/test_qdq_utils.py +++ b/tests/unit/onnx/quantization/test_qdq_utils.py @@ -15,14 +15,22 @@ import warnings +import ml_dtypes import numpy as np +import onnx import onnx_graphsurgeon as gs import onnxruntime as ort import pytest from onnx import TensorProto, helper, numpy_helper -from modelopt.onnx.export import INT4QuantExporter, MXFP8QuantExporter, NVFP4QuantExporter +from modelopt.onnx.export import ( + FP8QuantExporter, + INT4QuantExporter, + MXFP8QuantExporter, + NVFP4QuantExporter, +) from modelopt.onnx.export.nvfp4_exporter import _cast_fp4 +from modelopt.onnx.quantization.gs_patching import _export_value_info_proto from modelopt.onnx.quantization.qdq_utils import ( _cast_fp8, apply_column_major_transformation, @@ -484,6 +492,57 @@ def test_cast_fp4(self, input_array, expected_array): assert np.all(result == expected_array) +class TestFP8QuantExporter: + """Test suite for FP8QuantExporter.""" + + def test_bf16_weights_and_scale_are_compressed(self): + weight_data = np.array([0.001312255859375], dtype=ml_dtypes.bfloat16) + scale_data = np.array(0.00099945068359375, dtype=ml_dtypes.bfloat16) + weight = gs.Constant("weight", weight_data) + scale = gs.Constant("linear/weight_quantizer/scale", scale_data) + quantized = gs.Variable("quantized", dtype=np.uint8, shape=weight_data.shape) + dequantized = gs.Variable( + "dequantized", dtype=TensorProto.BFLOAT16, shape=weight_data.shape + ) + value_info = _export_value_info_proto(dequantized, do_type_check=True) + assert value_info.type.tensor_type.elem_type == TensorProto.BFLOAT16 + graph = gs.Graph( + nodes=[ + gs.Node( + op="TRT_FP8QuantizeLinear", + inputs=[weight, scale], + outputs=[quantized], + ), + gs.Node( + op="TRT_FP8DequantizeLinear", + inputs=[quantized, scale], + outputs=[dequantized], + ), + ], + outputs=[dequantized], + opset=23, + ) + + converted_model = FP8QuantExporter.compress_weights(gs.export_onnx(graph)) + + onnx.checker.check_model(converted_model) + assert [node.op_type for node in converted_model.graph.node] == ["DequantizeLinear"] + assert converted_model.graph.output[0].type.tensor_type.elem_type == TensorProto.BFLOAT16 + fp8_weight = next( + initializer + for initializer in converted_model.graph.initializer + if initializer.name == "linear/weight_quantizer/fp8_weights" + ) + assert fp8_weight.data_type == TensorProto.FLOAT8E4M3FN + assert fp8_weight.raw_data == b"\x3b" + output_scale = next( + initializer + for initializer in converted_model.graph.initializer + if initializer.name == scale.name + ) + assert output_scale.data_type == TensorProto.BFLOAT16 + + class TestMXFP8QuantExporter: """Test suite for MXFP8QuantExporter.""" diff --git a/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py b/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py index 399085ad699..2fd1a1c8504 100644 --- a/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py +++ b/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py @@ -13,7 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy import json +import tempfile from contextlib import nullcontext import numpy as np @@ -24,6 +26,7 @@ import torch.nn as nn from _test_utils.torch.deploy.lib_test_models import BaseDeployModel, get_deploy_models +import modelopt.torch.quantization as mtq from modelopt.onnx.utils import get_batch_size_from_bytes, validate_batch_size from modelopt.torch._deploy.utils import ( OnnxBytes, @@ -57,6 +60,64 @@ } +class _FP8ModelWithBuffer(nn.Sequential): + fp32_buffer: torch.Tensor + + def forward(self, inputs): + return super().forward(inputs) + self.fp32_buffer + + +def _make_fp8_model(source_dtype, kind="fp8"): + if kind == "format": + model = nn.Sequential(*(nn.Linear(128, 128, bias=False) for _ in range(2))) + sample_input = torch.ones(1, 128, dtype=source_dtype) + config = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + config["quant_cfg"].extend( + [ + { + "quantizer_name": "1.weight_quantizer", + "cfg": {"num_bits": 4, "block_sizes": {-1: 128, "type": "static"}}, + }, + {"quantizer_name": "1.input_quantizer", "enable": False}, + ] + ) + else: + model = _FP8ModelWithBuffer( + nn.Conv2d(1, 1, 1, bias=False), + nn.Flatten(), + nn.Linear(4, 4, bias=False), + ) + sample_input = torch.ones(1, 1, 2, 2, dtype=source_dtype) + config = mtq.FP8_DEFAULT_CFG + model = model.eval().to(source_dtype) + if kind != "format": + model.register_buffer("fp32_buffer", torch.ones(4)) + if kind == "parameters": + model.register_parameter("unused_fp32_parameter", nn.Parameter(torch.ones(1))) + quantized_model = mtq.quantize( + model, + config, + forward_loop=lambda quantized_model: quantized_model(sample_input), + ) + if kind != "format": + # Keep calibration nonzero while exercising Conv scale underflow during export. + with torch.no_grad(): + quantized_model[0].weight.fill_(1e-38 if source_dtype == torch.bfloat16 else 1.0) + return quantized_model, sample_input + + +def _export_fp8_model(source_dtype, weights_dtype): + model, sample_input = _make_fp8_model(source_dtype) + onnx_bytes, _ = get_onnx_bytes_and_metadata( + model, + (sample_input,), + weights_dtype=weights_dtype, + onnx_opset=23, + ) + onnx_bytes_obj = OnnxBytes.from_bytes(onnx_bytes) + return onnx.load_model_from_string(onnx_bytes_obj.get_onnx_model_file_bytes()) + + @pytest.mark.parametrize( "model", deploy_benchmark_dynamo.values(), ids=deploy_benchmark_dynamo.keys() ) @@ -157,6 +218,87 @@ def test_onnx_export_and_inputs(model: BaseDeployModel): ) +@pytest.mark.parametrize( + ("source_dtype", "weights_dtype", "expected_onnx_dtype"), + [ + (torch.bfloat16, "bf16", onnx.TensorProto.BFLOAT16), + (torch.float32, "fp16", onnx.TensorProto.FLOAT16), + ], + ids=["bf16-weight-focused-buffer", "fp32-to-fp16"], +) +def test_fp8_export_with_supported_weights_dtype(source_dtype, weights_dtype, expected_onnx_dtype): + exported_model = _export_fp8_model(source_dtype, weights_dtype) + + onnx.checker.check_model(exported_model, full_check=True) + assert {"TRT_FP8QuantizeLinear", "TRT_FP8DequantizeLinear"}.isdisjoint( + node.op_type for node in exported_model.graph.node + ) + initializer_by_name = { + initializer.name: initializer for initializer in exported_model.graph.initializer + } + fp8_weight_dq_nodes = [ + node + for node in exported_model.graph.node + if node.op_type == "DequantizeLinear" + and node.input[0] in initializer_by_name + and initializer_by_name[node.input[0]].data_type == onnx.TensorProto.FLOAT8E4M3FN + ] + assert len(fp8_weight_dq_nodes) == 2 + for node in fp8_weight_dq_nodes: + assert initializer_by_name[node.input[1]].data_type == expected_onnx_dtype + assert {0x7F, 0xFF}.isdisjoint(initializer_by_name[node.input[0]].raw_data) + assert all( + value.type.tensor_type.elem_type == expected_onnx_dtype + for value in exported_model.graph.input + ) + expected_output_dtype = ( + onnx.TensorProto.FLOAT if weights_dtype == "bf16" else expected_onnx_dtype + ) + assert all( + value.type.tensor_type.elem_type == expected_output_dtype + for value in exported_model.graph.output + ) + if weights_dtype == "bf16": + fp32_buffer = next( + initializer + for initializer in exported_model.graph.initializer + if initializer.name.endswith("fp32_buffer") + ) + assert fp32_buffer.data_type == onnx.TensorProto.FLOAT + assert any( + node.op_type == "Add" and fp32_buffer.name in node.input + for node in exported_model.graph.node + ) + + +@pytest.mark.parametrize( + ("kind", "source_dtype", "weights_dtype", "error"), + [ + ("fp8", torch.float32, "bf16", "torch.float32"), + ("fp8", torch.bfloat16, "fp16", "torch.bfloat16"), + ("parameters", torch.bfloat16, "bf16", "torch.bfloat16, torch.float32"), + ("format", torch.bfloat16, "bf16", "torch.bfloat16"), + ], + ids=["fp32-to-bf16", "bf16-to-fp16", "mixed-parameters", "mixed-format"], +) +def test_fp8_export_rejects_unsupported_dtype_conversion( + kind, source_dtype, weights_dtype, error, monkeypatch, tmp_path +): + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + model, sample_input = _make_fp8_model(source_dtype, kind) + with pytest.raises( + ValueError, + match=rf"Converting .* to {weights_dtype.upper()}.*source parameter dtypes: {error}", + ): + get_onnx_bytes_and_metadata( + model, + (sample_input,), + weights_dtype=weights_dtype, + onnx_opset=23, + ) + assert not any(tmp_path.iterdir()) + + class SingleArgModel(nn.Module): def forward(self, x: torch.Tensor): return torch.add(x, x) - x From 2fbcca252fbbf451e128ba53275be20c38c7b2e1 Mon Sep 17 00:00:00 2001 From: noeyy-mino <174223378+noeyy-mino@users.noreply.github.com> Date: Fri, 11 Sep 2026 19:21:23 +0800 Subject: [PATCH 06/28] deprecate trtllm-build in weight_sparsity (#2371) Type of change: export PTS/finetuned model to Hugging Face checkpoint, then replace trtllm-build with trtllm-serve Renamed export_trtllm_ckpt.py to export_hf_ckpt.py. Replaced the legacy export_tensorrt_llm_checkpoint() flow with export_hf_checkpoint(). Fix bug: 5823190 ``` python examples/llm_sparsity/weight_sparsity/hf_pts.py --model_name_or_path Llama-3.1-8B-Instruct --device cuda --model_max_length 1024 --dtype fp16 --sparsity_fmt sparsegpt --calib_size 128 --output_dir Llama-3.1-8B-Instruct_pts python examples/llm_sparsity/weight_sparsity/export_hf_ckpt.py --model_name_or_path Llama-3.1-8B-Instruct --model_max_length 1024 --dtype fp16 --modelopt_restore_path Llama-3.1-8B-Instruct_pts/pts_modelopt_state.pth --output_dir Llama-3.1-8B-Instruct_pts/trtllm/ckpt_pts trtllm-serve Llama-3.1-8B-Instruct_pts/trtllm/ckpt_pts \ --tp_size 1 \ --pp_size 1 \ --host 0.0.0.0 \ --port 8000 ``` PTS and SAT tested Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: N/A - 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?: N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - Did you get Claude approval on this PR?: N/A N/A * **Documentation** * Updated sparsity example instructions to export Hugging Face checkpoints and serve models with `trtllm-serve`. * Documented tensor and pipeline parallelism, host and port settings, and the OpenAI-compatible chat completions endpoint. * Corrected the PTS model restoration path. * **Bug Fixes** * Model export now saves the tokenizer alongside the checkpoint. * Model length configuration is interpreted as an integer. --------- Signed-off-by: Noey Yang <174223378+noeyy-mino@users.noreply.github.com> Signed-off-by: Chad Voegele --- .../llm_sparsity/weight_sparsity/README.md | 31 +++++++++++++------ ...xport_trtllm_ckpt.py => export_hf_ckpt.py} | 29 +++-------------- 2 files changed, 26 insertions(+), 34 deletions(-) rename examples/llm_sparsity/weight_sparsity/{export_trtllm_ckpt.py => export_hf_ckpt.py} (84%) diff --git a/examples/llm_sparsity/weight_sparsity/README.md b/examples/llm_sparsity/weight_sparsity/README.md index 97563aff007..42bac030ce9 100644 --- a/examples/llm_sparsity/weight_sparsity/README.md +++ b/examples/llm_sparsity/weight_sparsity/README.md @@ -135,28 +135,39 @@ ROUGE scores: {'rouge1': 42.174, 'rouge2': 19.2724, 'rougeL': 28.6989, 'rougeLsu Please refer to [link](../llm_eval/README.md#Evaluation-scripts-for-LLM-tasks) for more details of how to evaluate the sparsified models on other benchmarks, such as MMLU and HumanEval. -## Export TensorRT-LLM Checkpoint +## Export and Serve a Hugging Face Checkpoint -To export the PTS pytorch model to a TensorRT-LLM checkpoint, run the following command: +Export the PTS model to a Hugging Face checkpoint: ```sh -python export_trtllm_ckpt.py --model_name_or_path meta-llama/Llama-2-7b-hf \ +python export_hf_ckpt.py --model_name_or_path meta-llama/Llama-2-7b-hf \ --model_max_length 1024 \ --dtype fp16 \ - --modelopt_restore_path saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/pts_modelopt_state.pth \ - --output_dir saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/trtllm/ckpt_pts + --modelopt_restore_path saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/pts/pts_modelopt_state.pth \ + --output_dir saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/hf/pts ``` -To export the finetuned pytorch model to a TensorRT-LLM checkpoint, run the following command: +Export the finetuned model in the same way: ```sh -python export_trtllm_ckpt.py --model_name_or_path meta-llama/Llama-2-7b-hf \ +python export_hf_ckpt.py --model_name_or_path meta-llama/Llama-2-7b-hf \ --model_max_length 1024 \ --dtype fp16 \ --modelopt_restore_path saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/finetuned/finetuned_modelopt_state.pth \ - --output_dir saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/trtllm/ckpt_finetuned + --output_dir saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/hf/finetuned ``` -## Build TensorRT-LLM Engine +The output is a standard Hugging Face checkpoint, including the tokenizer and sparse model weights. +TensorRT-LLM loads it directly; no checkpoint conversion or engine build is required. -For guidance on how to build TensorRT-LLM engines, please refer to [link](https://nvidia.github.io/TensorRT-LLM/commands/trtllm-build.html#trtllm-build) and use the `--weight_sparsity` flag. +Serve the exported checkpoint and set the deployment parallelism at startup: + +```sh +trtllm-serve saved_models_Llama-2-7b-hf_sparsegpt_tp1_pp1/hf/pts \ + --tp_size 1 \ + --pp_size 1 \ + --host 0.0.0.0 \ + --port 8000 +``` + +After the server starts, send requests to the OpenAI-compatible `/v1/chat/completions` endpoint. diff --git a/examples/llm_sparsity/weight_sparsity/export_trtllm_ckpt.py b/examples/llm_sparsity/weight_sparsity/export_hf_ckpt.py similarity index 84% rename from examples/llm_sparsity/weight_sparsity/export_trtllm_ckpt.py rename to examples/llm_sparsity/weight_sparsity/export_hf_ckpt.py index 2cf7ca3a7ac..7ef22c52eca 100644 --- a/examples/llm_sparsity/weight_sparsity/export_trtllm_ckpt.py +++ b/examples/llm_sparsity/weight_sparsity/export_hf_ckpt.py @@ -23,7 +23,7 @@ import modelopt.torch.opt as mto import modelopt.torch.sparsity as mts -from modelopt.torch.export import export_tensorrt_llm_checkpoint, get_model_type +from modelopt.torch.export import export_hf_checkpoint DEFAULT_PAD_TOKEN = "[PAD]" @@ -103,8 +103,6 @@ def main(args): model=model, ) - # Export the sparse model to trt-llm checkpoint - model_type = get_model_type(model) if args.modelopt_restore_path: print(f"Loading sparsity state from {args.modelopt_restore_path}") if not os.path.isfile(args.modelopt_restore_path): @@ -112,17 +110,11 @@ def main(args): mto.restore(model, args.modelopt_restore_path) - print(f"Exporting trt-llm checkpoint to {args.output_dir}") + print(f"Exporting Hugging Face checkpoint to {args.output_dir}") with torch.inference_mode(): model = mts.export(model) - export_tensorrt_llm_checkpoint( - model, - model_type, - torch.float16, - export_dir=args.output_dir, - inference_tensor_parallel=args.inference_tensor_parallel, - inference_pipeline_parallel=args.inference_pipeline_parallel, - ) + export_hf_checkpoint(model, export_dir=args.output_dir) + tokenizer.save_pretrained(args.output_dir) if __name__ == "__main__": @@ -140,22 +132,11 @@ def main(args): parser.add_argument("--dtype", help="Model data type.", default="fp16") parser.add_argument( "--model_max_length", + type=int, default=2048, help="Maximum sequence length. Sequences will be right padded (and possibly truncated).", ) parser.add_argument("--output_dir", default="output_dir") - parser.add_argument( - "--inference_tensor_parallel", - help="Number of tensor parallel groups for inference.", - type=int, - default=1, - ) - parser.add_argument( - "--inference_pipeline_parallel", - help="Number of pipeline parallel groups for inference.", - type=int, - default=1, - ) parser.add_argument( "--trust_remote_code", help="Set trust_remote_code for Huggingface models and tokenizers", From 68db2c963436cfcc761405b9abdab247657c356a Mon Sep 17 00:00:00 2001 From: noeyy-mino <174223378+noeyy-mino@users.noreply.github.com> Date: Fri, 11 Sep 2026 22:46:10 +0800 Subject: [PATCH 07/28] Fix protobuf size-check failures in ONNX deployment (#2403) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix：6701737 The ONNX deployment path assumed that ModelProto.ByteSize() would always return a valid size. With newer protobuf versions, querying the size of a model exceeding the protobuf serialization limit can itself raise EncodeError: Failed to serialize proto. Replaced both direct size checks with the existing is_model_too_large_for_protobuf() helper. This helper handles size-query failures conservatively and checks the protobuf size limit: Shape inference now selects the external-data/file-based path when ByteSize() fails or the model is too large. Metadata creation uses the same safe check instead of raising another serialization error. The unused TWO_GB constant was also removed. ### Usage ``` python examples/diffusers/quantization/diffusion_trt.py --model flux-dev --benchmark --skip-image ``` ### Testing the above test case pass on B100 ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: N/A - 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?: N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - Did you get Claude approval on this PR?: N/A ### Additional Information N/A ## Summary by CodeRabbit * **Bug Fixes** * Improved ONNX model size detection during shape inference and export processing. * Ensured large models consistently use the appropriate external-data handling path. Signed-off-by: Noey Yang <174223378+noeyy-mino@users.noreply.github.com> Signed-off-by: Chad Voegele --- modelopt/torch/_deploy/utils/onnx_optimizer.py | 4 +++- modelopt/torch/_deploy/utils/torch_onnx.py | 4 ++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/_deploy/utils/onnx_optimizer.py b/modelopt/torch/_deploy/utils/onnx_optimizer.py index 08c4c40ca97..c520983f7c0 100644 --- a/modelopt/torch/_deploy/utils/onnx_optimizer.py +++ b/modelopt/torch/_deploy/utils/onnx_optimizer.py @@ -21,6 +21,8 @@ import onnx import onnx_graphsurgeon as gs +from modelopt.onnx.utils import is_model_too_large_for_protobuf + class Optimizer: """Optimizer for onnx graphs.""" @@ -62,7 +64,7 @@ def select_outputs(self, keep, names=None): def infer_shapes(self, return_onnx=False): """Infers shapes of the onnx graph.""" onnx_graph = gs.export_onnx(self.graph) - if onnx_graph.ByteSize() > (2 * (1024**3)): # 2GB limit + if is_model_too_large_for_protobuf(onnx_graph): # 2GB limit temp_dir = tempfile.TemporaryDirectory().name os.makedirs(temp_dir, exist_ok=True) onnx_orig_path = os.path.join(temp_dir, "model.onnx") diff --git a/modelopt/torch/_deploy/utils/torch_onnx.py b/modelopt/torch/_deploy/utils/torch_onnx.py index 8e16813bd3a..b217f1188b3 100644 --- a/modelopt/torch/_deploy/utils/torch_onnx.py +++ b/modelopt/torch/_deploy/utils/torch_onnx.py @@ -56,6 +56,7 @@ get_output_names, get_output_shapes, infer_shapes, + is_model_too_large_for_protobuf, remove_node_training_mode, remove_redundant_casts, ) @@ -96,7 +97,6 @@ def _patched_remove_unnecessary_cast_node(graph): # a few constants... DEFAULT_ONNX_OPSET = 20 ONNX_EXPORT_OUT_PREFIX = "out" -TWO_GB = 2 * 1024 * 1024 * 1024 class OnnxBytes: @@ -765,6 +765,6 @@ def create_model_metadata( "output_onnx_names": get_output_names(onnx_graph), "signature": inspect.signature(model.forward), "onnx_node_names": get_node_names(onnx_graph), - "is_bytes_pickled": onnx_graph.ByteSize() > TWO_GB, + "is_bytes_pickled": is_model_too_large_for_protobuf(onnx_graph), "config": model.config if hasattr(model, "config") else None, } From aca3f134a0a0ad55934d2a4057af75d2143e4430 Mon Sep 17 00:00:00 2001 From: realAsma <86726418+realAsma@users.noreply.github.com> Date: Fri, 11 Sep 2026 11:18:23 -0700 Subject: [PATCH 08/28] Document the MSE calibration API (#2405) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: documentation Expose `mse_calibrate` through `model_calib.__all__` so Sphinx autosummary includes the existing MSE calibration API on the generated `model_calib` reference page. The documentation configuration honors each module's curated `__all__` surface. Although `mse_calibrate` was implemented and used by the calibration dispatcher, it was missing from that surface and was therefore filtered out during API generation. ### Usage ```python from modelopt.torch.quantization.model_calib import mse_calibrate ``` ### Testing - `pre-commit run --files modelopt/torch/quantization/model_calib.py` - `git diff --check -- modelopt/torch/quantization/model_calib.py` - Generated the recursive autosummary API tree using the repository Sphinx configuration and module template. - Verified the generated RST contains `mse_calibrate`. - Rendered the focused module page and verified its function-table link, anchor, signature, and docstring. The canonical full build was unavailable in the active environment because the configured `shibuya` theme is not installed. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - 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?: N/A — the existing documentation build directly exercises this declarative autosummary contract. - Did you update Changelog?: N/A — this is a documentation-visibility repair for an existing API. - Did you get Claude approval on this PR?: N/A ### Additional Information No source implementation or calibration behavior changed. ## Summary by CodeRabbit - **New Features** - Made the MSE calibration capability publicly available for quantization workflows. - **Chores** - Increased documentation build time limits to improve reliability for longer-running builds. - Increased multi-version test time limits to better accommodate extended test runs. --------- Signed-off-by: realAsma Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Signed-off-by: Chad Voegele --- modelopt/torch/quantization/model_calib.py | 1 + 1 file changed, 1 insertion(+) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index c85e97a104d..2a504597c3f 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -73,6 +73,7 @@ "local_hessian_calibrate", "lsq", "max_calibrate", + "mse_calibrate", "nvfp4_act_headroom_calibrate", "smoothquant", "svdquant", From c51568cfd01ce2eb2cd6213213ae640310c2eb8f Mon Sep 17 00:00:00 2001 From: Jenny Chen Date: Fri, 11 Sep 2026 15:28:43 -0400 Subject: [PATCH 09/28] Fix TEGroupedMLP quantizer checkpoint resharding (#2319) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix for https://github.com/NVIDIA/Model-Optimizer/issues/2209 Fix TEGroupedMLP per-expert weight quantizer checkpoint resharding. `TEGroupedMLP` now saves its per-expert quantizer state as singleton local shards, allowing the distributed checkpoint format to retain each expert's global identity. Restore also initializes scalar `_amax` placeholders after ModelOpt extra-state restoration so distributed checkpoint loading can populate quantizer state for experts that move between ranks. This fixes restoring quantized TEGroupedMLP checkpoints across expert-parallel and tensor-parallel topology changes. Previously there was a bug that had two parts 1. TEGroupedMLP did not mark its per-expert quantizer state as singleton_local_shards. That meant the scalar weight_quantizer.._amax state was not saved with the same globally unique expert identity as the grouped-expert weights, so DCP could not reliably redistribute it across EP layouts. 2. During restore, ModelOpt’s extra-state restoration can leave _amax absent for experts that were not local on the checkpoint’s saving rank. The subsequent distributed checkpoint load then had no destination tensor to populate. ### Usage ```python # Add a code snippet demonstrating how to use this ``` ### Testing - `ruff format`, `ruff check`, `mypy`, `bandit`, and repository pre-commit hooks - Focused GPU regression: ```bash python3 -m pytest tests/gpu_megatron/torch/quantization/plugins/test_megatron.py \ -k te_grouped_sharded_state_dict_reshard -v Replaced the prior metadata-only TEGroupedMLP sharded-state test with an end-to-end distributed-checkpoint save/restore regression test. The new test: - Quantizes a TEGroupedMLP with per-expert NVFP4 weight quantizers. - Assigns each local expert a distinct, deterministic `_amax` based on its global expert index. - Saves both the model distributed checkpoint and sharded ModelOpt state. - Rebuilds the model under a different TP/EP topology. - Restores ModelOpt state, loads the distributed checkpoint, and verifies each target-local expert received the expected global-expert `_amax`. The parameterized test covers: - EP=2 -> EP=1 - EP=1 -> EP=2 - TP=1 -> TP=2 - TP=2 -> TP=1 ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ / ❌ / N/A - 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?: ✅ / ❌ / N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ / ❌ / N/A - Did you get Claude approval on this PR?: ✅ / ❌ / N/A ### Additional Information ## Summary by CodeRabbit - **Bug Fixes** - Improved checkpoint restoration for grouped quantizers by initializing missing quantization statistics with compatible shapes. - Improved restoration across supported grouped quantizer configurations, including sequential groups and parallel checkpoint layouts. - Extra module state is now finalized through supported post-load callbacks when available. - Preserved populated quantized output-layer state during checkpoint operations while removing empty placeholders. - **Tests** - Expanded checkpoint resharding coverage across tensor- and expert-parallel configurations. - Added coverage for disabled, dynamic, and other grouped quantizer scenarios. --------- Signed-off-by: Jennifer Chen Signed-off-by: Jenny Chen Signed-off-by: Chad Voegele --- .../opt/plugins/mcore_dist_checkpointing.py | 4 + .../torch/quantization/plugins/megatron.py | 36 +++ .../quantization/plugins/test_megatron.py | 262 +++++++++++++----- 3 files changed, 228 insertions(+), 74 deletions(-) diff --git a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py index 2bae53fa6f1..a4796162f6d 100644 --- a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py +++ b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py @@ -181,6 +181,10 @@ def _load_extra_state_from_sharded_checkpoint( module, "modelopt_set_extra_state_callbacks" ): module.set_extra_state(extra_state_dict_no_prefix[key]) + for module in model.modules(): + post_load_extra_state = getattr(module, "modelopt_post_load_extra_state", None) + if callable(post_load_extra_state): + post_load_extra_state() def restore_sharded_modelopt_state( diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 8c2352b5265..ec1958649dd 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -115,6 +115,39 @@ def _check_nvfp4_static_tp_supported(model: torch.nn.Module) -> None: ) +def _initialize_grouped_weight_quantizer_state(module: torch.nn.Module) -> None: + """Create per-expert buffers as destinations for the subsequent checkpoint load.""" + grouped_leaves = [ + quantizer if isinstance(quantizer, torch.nn.Sequential) else [quantizer] + for quantizer in [module.weight_quantizer[idx] for idx in range(module.num_gemms)] + ] + for sibling_leaves in zip(*grouped_leaves): + eligible_leaves = [ + quantizer + for quantizer in sibling_leaves + if ( + quantizer.is_enabled + and not quantizer.is_mx_format + and not getattr(quantizer, "_dynamic", False) + and not getattr(quantizer, "_lsq", False) + ) + ] + for state_name in ("_amax", "_global_amax"): + reference = next( + ( + state + for quantizer in eligible_leaves + if (state := getattr(quantizer, state_name, None)) is not None + ), + None, + ) + if reference is None: + continue + for quantizer in eligible_leaves: + if getattr(quantizer, state_name, None) is None: + quantizer.register_buffer(state_name, torch.zeros_like(reference)) + + def real_quant_module_get_extra_state(self) -> dict: """Populating real_quantizer_state and q_tensor_state.""" extra_state = {} @@ -818,6 +851,9 @@ class _QuantTELayerNormColumnParallelLinear( # Quantized subclasses to support TEGroupedLinear quantization class _QuantMegatronTEGroupedLinear(_QuantTEGroupedLinear, _MegatronParallelLinear): + def modelopt_post_load_extra_state(self): + _initialize_grouped_weight_quantizer_state(self) + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): # _sharded_state_dict_grouped adds _extra_state{gemm_idx} for gemm_idx:[1, num_gemms] in # sharded_state_dict which is same as _extra_state. The _extra_state{gemm_idx} is used for diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index bec404f9c24..041a0bd3fca 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -15,7 +15,6 @@ import copy import math -import re import sys import types from contextlib import nullcontext @@ -38,7 +37,9 @@ get_batch, get_forward, initialize_for_megatron, + load_distributed_checkpoint, run_mcore_inference, + save_distributed_checkpoint, sharded_state_dict_test_helper, ) from _test_utils.torch.misc import set_seed @@ -54,6 +55,7 @@ from megatron.core.parallel_state import ( destroy_model_parallel, get_data_parallel_group, + get_expert_model_parallel_rank, get_tensor_model_parallel_group, ) from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear @@ -64,10 +66,15 @@ import modelopt import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq +from modelopt.torch.opt.plugins.mcore_dist_checkpointing import ( + restore_sharded_modelopt_state, + save_sharded_modelopt_state, +) from modelopt.torch.quantization.algorithms import QuantRecipe, _AutoQuantizeBaseSearcher from modelopt.torch.quantization.nn import QuantModuleRegistry, SequentialQuantizer from modelopt.torch.quantization.nn.modules.quant_linear import RealQuantLinear from modelopt.torch.quantization.plugins.megatron import ( + _initialize_grouped_weight_quantizer_state, _output_layer_extra_state_has_data, _output_layer_untied, _QuantMegatronTEGroupedLinear, @@ -1094,100 +1101,207 @@ def test_te_grouped_vs_sequential_default_amax(dist_workers_size_1, quant_cfg): ) -def _te_grouped_expert_identity_from_sharded_state(module): - """Return {local_key: (global_expert_idx, num_global_experts)} for per-expert amax shards. +def _set_te_grouped_weight_quantizer_state(model, ep_rank, num_local_experts): + """Give every local expert distinct quantizer state derived from its global index.""" + for linear in model.modules(): + if not isinstance(linear, _QuantMegatronTEGroupedLinear): + continue + for local_expert_idx in range(linear.num_gemms): + quantizer = linear.weight_quantizer[local_expert_idx] + leaves = list(quantizer) if isinstance(quantizer, SequentialQuantizer) else [quantizer] + for leaf in leaves: + amax = getattr(leaf, "_amax", None) + if amax is not None: + amax.fill_(1.0 + ep_rank * num_local_experts + local_expert_idx) + global_amax = getattr(leaf, "_global_amax", None) + if global_amax is not None: + global_amax.fill_(1.0 + ep_rank * num_local_experts + local_expert_idx) + + +def _assert_te_grouped_weight_quantizer_state(model, expected_amax, expect_global_amax): + checked = 0 + for linear in model.modules(): + if not isinstance(linear, _QuantMegatronTEGroupedLinear): + continue + for local_expert_idx in range(linear.num_gemms): + quantizer = linear.weight_quantizer[local_expert_idx] + leaves = list(quantizer) if isinstance(quantizer, SequentialQuantizer) else [quantizer] + for leaf in leaves: + amax = getattr(leaf, "_amax", None) + assert amax is not None, ( + "TEGrouped per-expert weight quantizer amax was not restored" + ) + checked += 1 + assert torch.equal(amax, torch.full_like(amax, expected_amax[local_expert_idx])) + global_amax = getattr(leaf, "_global_amax", None) + if expect_global_amax: + assert global_amax is not None, ( + "TEGrouped per-expert weight quantizer global_amax was not restored" + ) + assert torch.equal( + global_amax, + torch.full_like(global_amax, expected_amax[local_expert_idx]), + ) + assert checked > 0, "no TEGrouped per-expert weight quantizer amax was checked" - The grouped linear must give each fused expert the same global identity the weights use: - the dict key keeps the local expert index (maps to the local buffer on restore) while the - ShardedTensor carries the global expert offset. Called with sharded_offsets=() so the expert - axis is the (only) prepended axis at index 0. - """ - sharded_sd = module.sharded_state_dict(prefix="", sharded_offsets=(), metadata=None) - identity = {} - for key, sh_ten in sharded_sd.items(): - if re.match(r"weight_quantizer\.\d+\..*_amax$", key): - assert sh_ten.prepend_axis_num >= 1, f"{key}: expected a prepended expert axis" - identity[key] = (int(sh_ten.global_offset[0]), int(sh_ten.global_shape[0])) - return identity +def test_initialize_grouped_weight_quantizer_state_for_restore(): + """Missing grouped state inherits the shape and dtype of a populated sibling.""" + source = mtq.nn.StaticBlockScaleQuantizer.from_tensor_quantizer( + mtq.nn.TensorQuantizer(amax=torch.tensor([1.0, 2.0])), global_amax=torch.tensor(2.0) + ) + target = mtq.nn.StaticBlockScaleQuantizer.from_tensor_quantizer(mtq.nn.TensorQuantizer()) + disabled = mtq.nn.StaticBlockScaleQuantizer.from_tensor_quantizer(mtq.nn.TensorQuantizer()) + disabled.disable() + mx = mtq.nn.TensorQuantizer( + mtq.config.QuantizerAttributeConfig( + num_bits=(4, 3), block_sizes={-1: 32, "type": "dynamic", "scale_bits": (8, 0)} + ) + ) -def _test_te_grouped_sharded_state_dict_global_expert_identity_helper( - tp_size, ep_size, quant_cfg, rank, size -): - """Per-expert quantizer amax must persist all num_global_experts across EP. + model = torch.nn.Module() + model.weight = torch.nn.Parameter(torch.empty(1)) + model.num_gemms = 4 + model.weight_quantizer = torch.nn.ModuleList([target, source, disabled, mx]) - With EP>1 the base linear emitted ``weight_quantizer.{local_i}._amax`` at the local index with - no expert offset, so every rank wrote identical keys and torch_dist dedup collapsed them to a - single rank's experts. Assert each rank's fused experts now carry distinct global identities so - the union across ranks covers every global expert. - """ + _initialize_grouped_weight_quantizer_state(model) + + assert torch.equal(target.amax, torch.zeros_like(source.amax)) + assert torch.equal(target.global_amax, torch.zeros_like(source.global_amax)) + assert disabled.amax is None + assert disabled.global_amax is None + assert mx.amax is None + assert not hasattr(mx, "_amax") + + +def _test_te_grouped_sharded_state_dict_reshard_helper( + save_tp_size, + save_ep_size, + load_tp_size, + load_ep_size, + quant_cfg, + expect_global_amax, + checkpoint_path, + rank, + size, +): + """Round-trip TEGroupedMLP amax through a topology change.""" + num_experts = 4 + save_num_local_experts = num_experts // save_ep_size initialize_for_megatron( - tensor_model_parallel_size=tp_size, - expert_model_parallel_size=ep_size, + tensor_model_parallel_size=save_tp_size, + expert_model_parallel_size=save_ep_size, seed=SEED, ) - num_experts = 4 - num_local = num_experts // ep_size - te_grouped = _gpt_model_provider( - tp_size=tp_size, - ep_size=ep_size, + source = _gpt_model_provider( + tp_size=save_tp_size, + ep_size=save_ep_size, hidden_size=32, moe_grouped_gemm=True, transformer_impl="transformer_engine", num_moe_experts=num_experts, ) - forward = get_forward(te_grouped, batch_size=8) - for module in te_grouped.modules(): + forward = get_forward(source, batch_size=8) + for module in source.modules(): if isinstance(module, TopKRouter): module.topk = module.num_experts - mtq.quantize(te_grouped, quant_cfg, forward) + mtq.quantize(source, copy.deepcopy(quant_cfg), forward) + _set_te_grouped_weight_quantizer_state( + source, get_expert_model_parallel_rank(), save_num_local_experts + ) + save_distributed_checkpoint(checkpoint_path, source) + save_sharded_modelopt_state([source], checkpoint_path) + torch.distributed.barrier() + del source + destroy_model_parallel() - grouped_linears = [ - m for m in te_grouped.modules() if isinstance(m, _QuantMegatronTEGroupedLinear) - ] - assert grouped_linears, "No grouped quant linears found" - - expected_global = {rank * num_local + i for i in range(num_local)} - for linear in grouped_linears: - # Give each expert a distinct amax so a value mix-up would also be observable. - for i in range(linear.num_gemms): - wq = linear.weight_quantizer[i] - leaves = list(wq) if isinstance(wq, SequentialQuantizer) else [wq] - for leaf in leaves: - if hasattr(leaf, "_amax") and leaf._amax is not None: - leaf._amax.fill_(1.0 + rank * num_local + i) - - identity = _te_grouped_expert_identity_from_sharded_state(linear) - # One entry per local expert per amax buffer; dict keys keep the LOCAL index. - local_keys = {int(re.search(r"weight_quantizer\.(\d+)\.", k).group(1)) for k in identity} - assert local_keys == set(range(num_local)), ( - f"Expected local expert keys {set(range(num_local))}, got {local_keys}" - ) - # ShardedTensor global identity: this rank owns experts {rank*num_local + i}. - local_global = {gidx for gidx, _ in identity.values()} - assert local_global == expected_global, ( - f"rank {rank}: expected global experts {expected_global}, got {local_global}" - ) - assert all(total == num_experts for _, total in identity.values()), ( - f"num_global_experts should be {num_experts}, got {identity}" + initialize_for_megatron( + tensor_model_parallel_size=load_tp_size, + expert_model_parallel_size=load_ep_size, + seed=SEED, + ) + target = _gpt_model_provider( + tp_size=load_tp_size, + ep_size=load_ep_size, + hidden_size=32, + moe_grouped_gemm=True, + transformer_impl="transformer_engine", + num_moe_experts=num_experts, + ) + target_models = [target] + restore_sharded_modelopt_state(target_models, checkpoint_path) + target = target_models[0] + load_distributed_checkpoint(checkpoint_path, target) + load_num_local_experts = num_experts // load_ep_size + expected_amax = tuple( + range( + 1 + get_expert_model_parallel_rank() * load_num_local_experts, + 1 + (get_expert_model_parallel_rank() + 1) * load_num_local_experts, ) - - # Gather the global expert indices across all EP ranks: the union must cover every expert. - gathered = [None] * size - torch.distributed.all_gather_object(gathered, sorted(expected_global)) - union = set() - for part in gathered: - union.update(part) - assert union == set(range(num_experts)), ( - f"Union of global experts across EP ranks should be {set(range(num_experts))}, got {union}" ) + _assert_te_grouped_weight_quantizer_state(target, expected_amax, expect_global_amax) -@pytest.mark.parametrize("quant_cfg", [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG]) -def test_te_grouped_sharded_state_dict_global_expert_identity(dist_workers_size_2, quant_cfg): +@pytest.mark.parametrize( + ( + "quant_cfg", + "expect_global_amax", + "save_tp_size", + "save_ep_size", + "load_tp_size", + "load_ep_size", + ), + [ + pytest.param(mtq.FP8_DEFAULT_CFG, False, 1, 2, 1, 1, id="fp8-ep-downsize"), + pytest.param(mtq.FP8_DEFAULT_CFG, False, 1, 1, 1, 2, id="fp8-ep-upsize"), + pytest.param(mtq.FP8_DEFAULT_CFG, False, 1, 1, 2, 1, id="fp8-tp-upsize"), + pytest.param(mtq.FP8_DEFAULT_CFG, False, 2, 1, 1, 1, id="fp8-tp-downsize"), + pytest.param(mtq.NVFP4_DEFAULT_CFG, False, 1, 2, 1, 1, id="nvfp4-ep-downsize"), + pytest.param(mtq.NVFP4_DEFAULT_CFG, False, 1, 1, 1, 2, id="nvfp4-ep-upsize"), + pytest.param(mtq.NVFP4_DEFAULT_CFG, False, 1, 1, 2, 1, id="nvfp4-tp-upsize"), + pytest.param(mtq.NVFP4_DEFAULT_CFG, False, 2, 1, 1, 1, id="nvfp4-tp-downsize"), + pytest.param( + mtq.NVFP4_W4A4_WEIGHT_MSE_FP8_SWEEP_CFG, + True, + 1, + 2, + 1, + 1, + id="nvfp4-mse-ep-downsize", + ), + pytest.param( + mtq.NVFP4_W4A4_WEIGHT_MSE_FP8_SWEEP_CFG, + True, + 1, + 1, + 1, + 2, + id="nvfp4-mse-ep-upsize", + ), + ], +) +def test_te_grouped_sharded_state_dict_reshard( + dist_workers_size_2, + tmp_path, + save_tp_size, + save_ep_size, + load_tp_size, + load_ep_size, + quant_cfg, + expect_global_amax, +): dist_workers_size_2.run( - partial(_test_te_grouped_sharded_state_dict_global_expert_identity_helper, 1, 2, quant_cfg) + partial( + _test_te_grouped_sharded_state_dict_reshard_helper, + save_tp_size, + save_ep_size, + load_tp_size, + load_ep_size, + quant_cfg, + expect_global_amax, + tmp_path, + ) ) From e75d6c1c7ee53de82beb6e6d4acda42a311ea87c Mon Sep 17 00:00:00 2001 From: Zhiyu Date: Fri, 11 Sep 2026 13:26:35 -0700 Subject: [PATCH 10/28] feat(quantization): fail fast when a quant config matches no weight quantizer (#2203) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: New feature (fail-fast guard; behavior change on a previously silent path) A `quant_cfg` whose module patterns don't match the model is not an error to `set_quantizer_by_cfg` — every pattern simply matches nothing. The run then calibrates, exports, and hands back a checkpoint that is silently unquantized: ```json {"quantization": {"quant_algo": null, "kv_cache_quant_algo": "FP8", "quantized_layers": {}}} ``` Nothing in the run says so. It has only ever been caught by someone reading the exported `hf_quant_config.json` afterwards — most recently on Step-3.7 ([NVBug 6518665](https://nvbugspro.nvidia.com/bug/6518665), after a full 8×B200 calibration), and before that on MiniMax-M3, where fused-expert detection skipped the experts and an experts-only recipe matched nothing. `mtq.quantize` now compares the config's intent against the outcome and raises **before calibration**: ``` RuntimeError: The quantization config asks for weight quantization but no weight quantizer was enabled, so nothing would be quantized (3 quantizer(s) inserted). These patterns matched no weight quantizer: *.experts.*weight_quantizer Either the patterns do not match this architecture's module names (check the model-specific recipes under modelopt_recipes/huggingface//), or the modules holding the weights were never converted to quantized modules (an unsupported custom module, e.g. a trust_remote_code MoE layout). ``` Scoped to avoid false positives: - **Only configs that ask for weight quantization** are checked (an entry with `enable` and `weight_quantizer` in its pattern), so activation-only and KV-cache-only configs are unaffected. - **Intent is read from each pattern's final entry**, since `quant_cfg` entries apply in order: a pattern that is enabled and then disabled later asks for nothing by the end. - **Configs refining an already-quantized model** (weight quantizers enabled by an earlier `mtq.quantize`) are left alone. Matching goes through `conversion._match_quantizer` — the same matcher `set_quantizer_by_cfg` used to apply the config — so "did this pattern match anything?" is answered exactly as the applying code would. A local `fnmatch` diverges on the two cases that matcher handles: `SequentialQuantizer` modules (W4A8-style list-valued `cfg`) and fused-experts names (`..._weight_quantizers.0` normalizing to `..._weight_quantizer`). ### Usage No API change. A config that would previously have produced an unquantized checkpoint now raises: ```python mtq.quantize(model, {"quant_cfg": [ {"quantizer_name": "*", "enable": False}, {"quantizer_name": "*.experts.*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, ]}, forward_loop) # RuntimeError if the model has no `experts` modules ``` ### Testing Seven tests in `tests/unit/torch/quantization/test_quantize_cpu.py`, one per branch of the guard: patterns matching nothing raise; an activation-only config still runs; weight patterns disabled by a later entry still run; enabled-then-retracted patterns still run; `SequentialQuantizer` (list-valued `cfg`) and fused-experts quantizer names count as matched; and the already-quantized refinement path is exercised. Each was checked to be non-vacuous by removing the corresponding branch and confirming exactly that test fails. **One existing test changed.** `tests/gpu/torch/export/test_fsdp2_export.py` parametrized over `NVFP4_MLP_ONLY_CFG`, but its `SmallQKVModel` has no MLP — so that case ran the FSDP2 paths against an *unquantized* model, and the new guard reported it (4 GPU failures on the first CI run, all `quant_config6`; `NVFP4_OMLP_ONLY_CFG` passed because that model does have `o_proj`). The parametrization is dropped with a comment; `NVFP4_OMLP_ONLY_CFG` keeps the scoped-recipe coverage. **If reviewers would rather not change that test's meaning, the alternative is to downgrade the guard to a warning — flagging it explicitly as a decision.** I also swept every shipped `mtq.*_CFG` against `SmallQKVModel`: only the four MLP/experts-scoped configs raise, and the other three (`NVFP4_EXPERTS_ONLY_CFG`, `MXFP4_MLP_WEIGHT_ONLY_CFG`, `NVFP4_MLP_WEIGHT_ONLY_CFG`) are used elsewhere only against real MoE models (Qwen3-MoE, gpt-oss), so no other test is affected. Ran locally after rebasing onto current `main` (torch 2.11, transformers 5.5.4): `tests/unit/torch/quantization` + `tests/unit/recipe` — 1271 passed, 7 skipped. Full `tests/unit` (minus onnx, and puzzletron which needs `hydra`): 2676 passed, with 4 pre-existing `test_quant_aware_conversion.py` failures that reproduce unchanged on clean `main`. GPU tests were not run locally (no suitable GPU); the FSDP2 change above is reasoned from the CI failure, not re-run. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ❌ — deliberately. A config that previously produced a `quant_algo: null` checkpoint now raises. Any such run was already not doing what it claimed; the three scoping rules above keep intentional non-weight quantization working. - 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)?: ✅ (Backward Breaking Changes) - Did you get Claude approval on this PR?: ❌ ### Additional Information Pairs with #2202 (PTQ support for Step-3.7 MoE checkpoints), which fixes the specific model that motivated this. Independent branches; either can merge first. 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit - **Bug Fixes** - Quantization now detects enabled weight-quantization patterns that do not apply to any model weights and reports a clear validation error before calibration. - Broad wildcard patterns and nested quantizers are now handled correctly. - Overlapping patterns respect the final matching setting, including later disabling rules. - Existing quantized models can be refined using the parsed configuration. - Activation-only and explicitly disabled weight-quantization configurations remain supported. - Pipeline-parallel stages without targeted weights can bypass this validation when configured to do so. - **Documentation** - Documented the process-wide override for bypassing unmatched weight-quantizer validation. --------- Signed-off-by: Zhiyu Cheng Co-authored-by: Claude Opus 5 (1M context) Signed-off-by: Chad Voegele --- CHANGELOG.rst | 1 + modelopt/torch/quantization/model_quant.py | 85 +++++++- tests/gpu/torch/export/test_fsdp2_export.py | 8 +- .../torch/quantization/test_quantize_cpu.py | 184 ++++++++++++++++++ 4 files changed, 274 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ceda42b6c96..85047aee189 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -47,6 +47,7 @@ Changelog - Remove in-trainer quantization via ``QuantizationArguments.quant_cfg`` / ``--quant_cfg`` (deprecated in 0.45); use ``--recipe``. New recipes ``general/ptq/mxfp4_mlp_weight_only`` and ``general/ptq/nvfp4_mlp_weight_only`` replace ``MXFP4_MLP_WEIGHT_ONLY_CFG`` / ``NVFP4_MLP_WEIGHT_ONLY_CFG`` in the ``examples/gpt-oss`` QAT flow. - Remove the ``QuantizationArgumentsWithConfig`` alias in ``modelopt.torch.quantization.plugins.transformers_trainer`` (deprecated in 0.45). Use ``QuantizationArguments``. - Transformer Engine ``TEGroupedLinear`` (fused MoE experts) now uses **per-expert** weight quantization (one ``amax`` per expert) instead of a single shared ``amax``, so ModelOpt checkpoints containing quantized ``TEGroupedLinear`` modules saved before 0.47 are **not compatible** with 0.47. Re-run PTQ to regenerate compatible checkpoints. +- ``mtq.quantize`` now raises when a config asks for weight quantization but none of its weight-quantizer patterns match the model, instead of calibrating and exporting a silently unquantized checkpoint (``"quant_algo": null``). Configs that quantize activations or the KV cache only are unaffected, as are patterns that match and are then disabled by a later entry. If this fires, use the recipe for that architecture under ``modelopt_recipes/huggingface//`` or fix the module patterns. Set ``MODELOPT_SKIP_WEIGHT_QUANT_CHECK=1`` to disable the check process-wide, e.g. for a pipeline-parallel rank whose local stage legitimately has none of the targeted modules. **Deprecations** diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 3f6040fd4ef..4588c6ac54b 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -144,6 +144,85 @@ def postprocess_amax(model: nn.Module, key: str, post_process_fn) -> nn.Module: return model +_SKIP_WEIGHT_QUANT_CHECK_ENV = "MODELOPT_SKIP_WEIGHT_QUANT_CHECK" + + +def _check_weight_quantization_took_effect(model: nn.Module, config: QuantizeConfig) -> None: + """Raise when a config asks for weight quantization but no weight quantizer is enabled. + + A config whose module patterns do not match the model is not an error to + :func:`set_quantizer_by_cfg` — every pattern simply matches nothing — so the run + proceeds through calibration and export and produces a checkpoint that is silently + unquantized (``"quant_algo": null`` with an empty ``quantized_layers``). That has + bitten several MoE architectures whose module naming differs from the wildcards in + the general recipes, and it is only noticed when someone reads the exported config. + + By the time this runs, :func:`set_quantizer_by_cfg` (or the ``apply_mode`` conversion + that calls it) has already applied ``config`` to ``model``, so each quantizer's + ``is_enabled`` *is* the true outcome of that application — checking it directly cannot + diverge from what the config actually did. An earlier version of this check instead + re-derived "did this pattern match anything?" via a separate matcher call, which missed + the case of two *different* overlapping patterns (e.g. ``*weight_quantizer`` enabling + something a later, broader ``*`` then disables): the narrower pattern registered as + "matched" even though the quantizer it matched ended up disabled. + + A config that never asks for weight quantization (activation-only or KV-cache-only) + must not raise, so the check first looks at the config's own intent — via each + pattern's *final* entry, since entries apply in order and the last one for a pattern + wins — before looking at the model at all. + """ + if os.environ.get(_SKIP_WEIGHT_QUANT_CHECK_ENV) == "1": + return + + # Later entries override earlier ones, so only each pattern's final state states intent. + # A pattern naming ``weight_quantizer`` explicitly (the common case, e.g. + # ``*weight_quantizer``, ``*.experts.*weight_quantizer``) is caught by the substring + # check. A broad wildcard that never mentions "weight" -- a bare ``"*"`` catch-all, or + # ``"*_quantizer"`` -- can still match weight quantizers at runtime, so it must count + # too, or a config built only from patterns like that would never trip the guard + # regardless of what the model contains. ``fnmatch`` against the literal probe string + # ``"weight_quantizer"`` catches those (a pattern matching that bare name is, by + # construction, asking for one) without replacing the substring check: the probe alone + # would miss ``*.experts.*weight_quantizer`` (there is no ``.experts.`` in the probe + # string), which is what recognizes model-scoped patterns like the Step / MoE recipes use. + last_entry_per_pattern = {entry.quantizer_name: entry for entry in config.quant_cfg} + weight_patterns = [ + pattern + for pattern, entry in last_entry_per_pattern.items() + if entry.enable + and ("weight_quantizer" in pattern or fnmatch.fnmatch("weight_quantizer", pattern)) + ] + if not weight_patterns: + return + + # `SequentialQuantizer.is_enabled` delegates to its first member, so a list-valued `cfg`'s + # quantizers are already covered here without naming `SequentialQuantizer` explicitly: + # `named_modules()` recurses into the container and yields those children too, individually, + # named `...weight_quantizer.0` / `.1` (the substring match below still applies to them). + if any( + module.is_enabled + for name, module in model.named_modules() + if isinstance(module, TensorQuantizer) and "weight_quantizer" in name + ): + return + + patterns = "\n ".join(sorted(weight_patterns)) + raise RuntimeError( + "The quantization config asks for weight quantization but no weight quantizer is " + f"enabled, so nothing would be quantized. These patterns asked for it:\n {patterns}\n" + "Either the patterns do not match this architecture's module names (check the " + "model-specific recipes under modelopt_recipes/huggingface//), or the " + "modules holding the weights were never converted to quantized modules (an " + "unsupported custom module, e.g. a trust_remote_code MoE layout).\n" + "Under pipeline parallelism, a rank whose local stage genuinely has none of the " + "targeted modules (e.g. a pure-attention stage under an experts-only recipe) hits " + "this too, while other ranks proceed into calibration -- a collective hang, not " + f"just a wrong per-rank verdict. Set {_SKIP_WEIGHT_QUANT_CHECK_ENV}=1 to bypass this " + "check in that situation -- note this is process-global, so it silences the check " + "on every rank, not only the one with the legitimately empty stage." + ) + + def quantize( model: nn.Module, config: dict[str, Any | QuantizeConfig], @@ -241,12 +320,14 @@ def forward_loop(model) -> None: Returns: A pytorch model which has been quantized and calibrated. """ + quantize_config = QuantizeConfig(**dict(config)) if not is_quantized(model): model = apply_mode(model, mode=[("quantize", dict(config))], registry=QuantizeModeRegistry) else: # Already quantized, so lets apply the quant_cfg from the config - quant_cfg = QuantizeConfig(**dict(config)).quant_cfg - set_quantizer_by_cfg(model, quant_cfg) + set_quantizer_by_cfg(model, quantize_config.quant_cfg) + # Fail before calibration rather than after exporting an unquantized checkpoint. + _check_weight_quantization_took_effect(model, quantize_config) return calibrate(model, config.get("algorithm"), forward_loop=forward_loop) diff --git a/tests/gpu/torch/export/test_fsdp2_export.py b/tests/gpu/torch/export/test_fsdp2_export.py index 65f33abd41c..dc35974c047 100644 --- a/tests/gpu/torch/export/test_fsdp2_export.py +++ b/tests/gpu/torch/export/test_fsdp2_export.py @@ -261,7 +261,9 @@ def test_fsdp2_weight_update_context_for_export(dist_workers): # mtq.W4A8_AWQ_BETA_CFG, #TODO: Fix unit test for this case # mtq.FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG, #TODO: Fix unit test for this case mtq.W4A8_MXFP4_FP8_CFG, - mtq.NVFP4_MLP_ONLY_CFG, + # NVFP4_MLP_ONLY_CFG is omitted: SmallQKVModel has no MLP, so that config matched + # nothing and the parametrization exercised an unquantized model. NVFP4_OMLP_ONLY_CFG + # covers the same scoped-recipe shape here because it also selects `o_proj`. mtq.NVFP4_OMLP_ONLY_CFG, ], ) @@ -281,7 +283,9 @@ def test_fsdp2_weight_update_context_for_fuse_layers(dist_workers, quant_config, # mtq.W4A8_AWQ_BETA_CFG, #TODO: Fix unit test for this case # mtq.FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG, #TODO: Fix unit test for this case mtq.W4A8_MXFP4_FP8_CFG, - mtq.NVFP4_MLP_ONLY_CFG, + # NVFP4_MLP_ONLY_CFG is omitted: SmallQKVModel has no MLP, so that config matched + # nothing and the parametrization exercised an unquantized model. NVFP4_OMLP_ONLY_CFG + # covers the same scoped-recipe shape here because it also selects `o_proj`. mtq.NVFP4_OMLP_ONLY_CFG, ], ) diff --git a/tests/unit/torch/quantization/test_quantize_cpu.py b/tests/unit/torch/quantization/test_quantize_cpu.py index 3e4925e7b63..eb62e2fbbcb 100644 --- a/tests/unit/torch/quantization/test_quantize_cpu.py +++ b/tests/unit/torch/quantization/test_quantize_cpu.py @@ -16,6 +16,8 @@ """High-level tests for quantization.""" import copy +import os +from unittest import mock import pytest import torch @@ -441,6 +443,188 @@ def test_enable_only_entry_preserves_attributes(): assert module.axis == 0, "axis should be preserved by enable-only entry" +def test_weight_patterns_matching_nothing_raise(): + """A config whose weight patterns match no module must fail, not quantize nothing. + + Otherwise calibration and export run to completion and produce a checkpoint that is + silently unquantized (``"quant_algo": null``). + """ + model = SimpleLinear() + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + # No module in this model is named `experts`. + {"quantizer_name": "*.experts.*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + ], + "algorithm": "max", + } + with pytest.raises(RuntimeError, match="no weight quantizer is enabled"): + mtq.quantize(model, config, lambda m: m(m.get_input())) + + +def test_config_without_weight_quantization_is_allowed(): + """Activation-only configs quantize no weight on purpose and must still run.""" + model = SimpleLinear() + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*input_quantizer", "cfg": {"num_bits": 8, "axis": None}}, + ], + "algorithm": "max", + } + model = mtq.quantize(model, config, lambda m: m(m.get_input())) + + for name, module in model.named_modules(): + if name.endswith("weight_quantizer"): + assert not module.is_enabled + + +def test_weight_quantizers_disabled_by_a_later_entry_are_allowed(): + """Patterns that match and are then switched off are a choice, not a mismatch.""" + model = SimpleLinear() + config = { + "quant_cfg": [ + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 4, "axis": 0}}, + {"quantizer_name": "*weight_quantizer", "enable": False}, + ], + "algorithm": "max", + } + model = mtq.quantize(model, config, lambda m: m(m.get_input())) + + for name, module in model.named_modules(): + if name.endswith("weight_quantizer"): + assert not module.is_enabled + + +def test_sequential_weight_quantizers_do_not_trip_the_guard(): + """List-valued `cfg` builds `SequentialQuantizer`s, but the guard only ever looks at + `TensorQuantizer` instances -- this pins that it still doesn't wrongly raise for them. + + `SequentialQuantizer` (itself an `nn.Sequential`) is not special-cased: its `TensorQuantizer` + children are reachable directly via `named_modules()`, individually named + `...weight_quantizer.0` / `.1`, so the substring match already finds them. + """ + model = SimpleLinear() + calib_data = [model.get_input() for _ in range(2)] + quantize_model_and_forward(model, copy.deepcopy(WINT4INT8_CFG), calib_data) + + for name, module in model.named_modules(): + if name.endswith("weight_quantizer"): + assert isinstance(module, SequentialQuantizer) + + +def test_fused_experts_quantizer_names_do_not_trip_the_guard(): + """Fused-experts quantizers are named `..._weight_quantizers.N`, plural and indexed, but + still contain `weight_quantizer` as a substring and so are still read by the guard. + """ + model = SimpleLinear() + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda m: m(m.get_input())) + + # Rename as the fused-experts path does; the config's `*weight_quantizer` must still match. + linear = model.net[0] + linear.add_module("gate_up_proj_weight_quantizers", torch.nn.ModuleList([TensorQuantizer()])) + + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*gate_up_proj_weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + ], + "algorithm": "max", + } + mtq.quantize(model, config, lambda m: m(m.get_input())) + + +def test_refining_an_already_quantized_model_does_not_raise(): + """A second config whose own patterns match nothing still sees the earlier weight + quantizers as enabled, so this is refining an already-quantized model, not a no-op run. + """ + model = SimpleLinear() + model = mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda m: m(m.get_input())) + assert any(m.is_enabled for n, m in model.named_modules() if n.endswith("weight_quantizer")) + + refinement = { + "quant_cfg": [ + {"quantizer_name": "*.experts.*weight_quantizer", "cfg": {"num_bits": 4, "axis": 0}}, + ], + "algorithm": None, + } + mtq.quantize(model, refinement) + + # The earlier weight quantization is untouched. + assert any(m.is_enabled for n, m in model.named_modules() if n.endswith("weight_quantizer")) + + +def test_weight_patterns_enabled_then_retracted_do_not_raise(): + """An unmatched pattern that a later entry disables asks for nothing by the end.""" + model = SimpleLinear() + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*.missing.*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + {"quantizer_name": "*.missing.*weight_quantizer", "enable": False}, + ], + "algorithm": "max", + } + model = mtq.quantize(model, config, lambda m: m(m.get_input())) + + for name, module in model.named_modules(): + if name.endswith("weight_quantizer"): + assert not module.is_enabled + + +def test_bare_wildcard_pattern_matching_nothing_raises(): + """A bare `"*"` (or another pattern never mentioning "weight") still expresses weight + intent if it would match a `weight_quantizer` name -- and must still raise if nothing in + the model actually has one, exactly like an explicit `*weight_quantizer` pattern would. + + A model with zero quantizable modules (e.g. `nn.Module()`, no Linear/Conv anywhere) is + the degenerate case where this matters: nothing in the config's own text says "weight", + so a substring-only check would silently return without ever looking at the model. + """ + model = torch.nn.Module() # no quantizable submodules at all + config = {"quant_cfg": [{"quantizer_name": "*", "cfg": {"num_bits": 8, "axis": 0}}]} + with pytest.raises(RuntimeError, match="no weight quantizer is enabled"): + mtq.quantize(model, config) + + +def test_overlapping_patterns_disabled_by_a_broader_later_one_still_raise(): + """A narrower pattern "matching" is not enough -- the final enabled state is what counts. + + `*weight_quantizer` enables real quantizers, but the later, broader `*` disables + everything again; the config's net effect is still "nothing quantized" and must raise. + A check that asked "did any weight pattern match something" instead of "is anything + actually enabled" would miss this, since the narrower pattern did match. + """ + model = SimpleLinear() + config = { + "quant_cfg": [ + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + {"quantizer_name": "*", "enable": False}, + ], + "algorithm": "max", + } + with pytest.raises(RuntimeError, match="no weight quantizer is enabled"): + mtq.quantize(model, config, lambda m: m(m.get_input())) + + +def test_skip_weight_quant_check_env_var_bypasses_the_guard(): + """Documented escape hatch for pipeline-parallel ranks whose local stage legitimately + has none of the targeted modules (e.g. a pure-attention stage under an experts-only + recipe): raising there while other ranks proceed into calibration is a collective hang, + not just a wrong per-rank verdict. + """ + model = SimpleLinear() + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*.experts.*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + ], + "algorithm": "max", + } + with mock.patch.dict(os.environ, {"MODELOPT_SKIP_WEIGHT_QUANT_CHECK": "1"}): + mtq.quantize(model, config, lambda m: m(m.get_input())) + + def test_atomicity_later_cfg_entry_does_not_inherit_earlier(): """When two cfg-bearing entries match the same quantizer, the second fully replaces the first.""" model = SimpleLinear() From 83876db16caeb970a0dbc41aceae9399a360e37a Mon Sep 17 00:00:00 2001 From: Ajinkya Rasane <131806219+ajrasane@users.noreply.github.com> Date: Fri, 11 Sep 2026 18:04:19 -0400 Subject: [PATCH 11/28] [6701308][OMNIML-5805] Correct ONNX PTQ documentation contracts (#2413) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: documentation Align the ONNX PTQ README, guide, and executable example with the implemented contracts: - use the canonical `--calibration_data_path` CLI option; - load `.npy` calibration data before passing it to the Python API; - document the supported Autotune modes and calibration methods; - correct the minimum opsets to INT8 19, FP8 19, and INT4 21; and - describe the no-data fallback as random calibration inputs. This also removes an inaccurate source comment without changing runtime behavior. ### Usage ```bash python -m modelopt.onnx.quantization \ --onnx_path=model.onnx \ --quantize_mode=int8 \ --calibration_data_path=calib.npy \ --output_path=model.quant.onnx ``` ### Testing - `pre-commit run --files docs/source/guides/_onnx_quantization.rst examples/onnx_ptq/README.md modelopt/onnx/quantization/quantize.py tests/examples/test_onnx_ptq.sh` - `bash -n tests/examples/test_onnx_ptq.sh` - `CUDA_VISIBLE_DEVICES="" python -m pytest -o addopts="" -p no:cacheprovider --confcutdir=tests/unit/onnx/quantization -q tests/unit/onnx/quantization/test_autotune_quantization_integration.py` (4 passed) - `nox -s docs` (passed; Sphinx built 881 HTML files) - Focused before/after contract probe covering the documented CLI option, API data type, Autotune modes and methods, opset minimums, and random-input wording ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: N/A - 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?: N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - Did you get Claude approval on this PR?: N/A ### Additional Information Tracking: [6701308] > 🤖 _Generated by Codex (AI agent)._ ## Summary by CodeRabbit - **Documentation** - Clarified that random calibration inputs are used when no calibration dataset is provided. - Updated ONNX post-training quantization examples with minimum opset requirements and the `calibration_data_path` argument. - Clarified Autotune support for FP8 and INT8 calibration methods using `max` or `entropy`. - **Tests** - Updated quantization command examples to use the current calibration data path option. Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> Co-authored-by: Codex Signed-off-by: Chad Voegele --- docs/source/guides/_onnx_quantization.rst | 2 +- examples/onnx_ptq/README.md | 16 +++++++++------- modelopt/onnx/quantization/quantize.py | 1 - tests/examples/test_onnx_ptq.sh | 2 +- 4 files changed, 11 insertions(+), 10 deletions(-) diff --git a/docs/source/guides/_onnx_quantization.rst b/docs/source/guides/_onnx_quantization.rst index e4d0c2d93d6..22fd9721f44 100644 --- a/docs/source/guides/_onnx_quantization.rst +++ b/docs/source/guides/_onnx_quantization.rst @@ -37,7 +37,7 @@ Requirements Apply Post Training Quantization (PTQ) ====================================== -PTQ should be done with a calibration dataset. If calibration dataset is not provided, ModelOpt will use random scales for the QDQ nodes. +PTQ should be done with a calibration dataset. Random calibration inputs are used when no calibration dataset is provided. Prepare calibration dataset --------------------------- diff --git a/examples/onnx_ptq/README.md b/examples/onnx_ptq/README.md index de983b409bd..65c62f0655a 100644 --- a/examples/onnx_ptq/README.md +++ b/examples/onnx_ptq/README.md @@ -111,7 +111,7 @@ The model can be quantized as an FP8, INT8 or INT4 model using either the CLI or > *For NVFP4 and MXFP8 ONNX, see the [PyTorch to ONNX example](../torch_onnx/).* -> *Minimum opset requirements: int8 (13+), fp8 (21+), int4 (21+). ModelOpt will automatically upgrade lower opset versions to meet these requirements.* +> *Minimum opset requirements: int8 (19+), fp8 (19+), int4 (21+). ModelOpt will automatically upgrade lower opset versions to meet these requirements.* #### Option 1: Command-line interface @@ -119,7 +119,7 @@ The model can be quantized as an FP8, INT8 or INT4 model using either the CLI or python -m modelopt.onnx.quantization \ --onnx_path=vit_base_patch16_224.onnx \ --quantize_mode= \ - --calibration_data=calib.npy \ + --calibration_data_path=calib.npy \ --calibration_method= \ --output_path=vit_base_patch16_224.quant.onnx ``` @@ -127,12 +127,14 @@ python -m modelopt.onnx.quantization \ #### Option 2: Python API ```python +import numpy as np + from modelopt.onnx.quantization import quantize quantize( onnx_path="vit_base_patch16_224.onnx", quantize_mode="int8", # fp8, int8, int4 etc. - calibration_data="calib.npy", + calibration_data=np.load("calib.npy"), calibration_method="max", # max, entropy, awq_clip, rtn_dq etc. output_path="vit_base_patch16_224.quant.onnx", ) @@ -198,7 +200,7 @@ To enable per node calibration, add the `--calibrate_per_node` flag to your quan python -m modelopt.onnx.quantization \ --onnx_path=vit_base_patch16_224.onnx \ --quantize_mode= \ - --calibration_data=calib.npy \ + --calibration_data_path=calib.npy \ --calibrate_per_node \ --output_path=vit_base_patch16_224.quant.onnx ``` @@ -253,9 +255,9 @@ To access this feature in the ONNX quantization workflow, simply add `--autotune ```bash python -m modelopt.onnx.quantization \ --onnx_path=vit_base_patch16_224.onnx \ - --quantize_mode= \ - --calibration_data=calib.npy \ - --calibration_method= \ + --quantize_mode= \ + --calibration_data_path=calib.npy \ + --calibration_method= \ --output_path=vit_base_patch16_224.quant.onnx \ --autotune= ``` diff --git a/modelopt/onnx/quantization/quantize.py b/modelopt/onnx/quantization/quantize.py index fb6b662916d..ef16390b90e 100755 --- a/modelopt/onnx/quantization/quantize.py +++ b/modelopt/onnx/quantization/quantize.py @@ -779,7 +779,6 @@ def quantize( ) if calibration_data_reader is None: - # Use random scales if calibration data is not supplied if calibration_data is None: calibration_data_reader = RandomDataProvider(onnx_path, calibration_shapes) else: diff --git a/tests/examples/test_onnx_ptq.sh b/tests/examples/test_onnx_ptq.sh index 2e4805abfe9..f0a10b8eb1c 100755 --- a/tests/examples/test_onnx_ptq.sh +++ b/tests/examples/test_onnx_ptq.sh @@ -152,7 +152,7 @@ for model_path in "${model_paths[@]}"; do python -m modelopt.onnx.quantization \ --onnx_path=$model_dir/fp16/model.onnx \ --quantize_mode=$quant_mode \ - --calibration_data=$calib_data_path \ + --calibration_data_path=$calib_data_path \ --output_path=$model_dir/$quant_mode/model.quant.onnx \ --calibration_eps=cuda done From 5ecede066abe33ce50ba9aad96cfb6512eda0f2c Mon Sep 17 00:00:00 2001 From: Ajinkya Rasane <131806219+ajrasane@users.noreply.github.com> Date: Fri, 11 Sep 2026 23:49:01 -0400 Subject: [PATCH 12/28] [6463897] Fix narrow FP16 histogram calibration (#2412) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix FP16 entropy calibration can fail for sufficiently narrow activation ranges because NumPy may construct the histogram bin edges at FP16 precision. NumPy 2.2 and later reject the resulting collapsed bin spacing, while earlier versions can silently return invalid, non-monotonic edges. The same precision issue can recur when ONNX Runtime merges later calibration batches. Losslessly widen FP16 activation values to FP32 while calculating and merging histograms in both ONNX entropy calibration paths, then restore the source dtype at the calibration-to-quantization boundary. This keeps the histogram bins stable without changing FP16 Q/DQ scale or graph dtype semantics. FP32 inputs and public APIs are unchanged. For full-range FP16 activations, ONNX Runtime can overflow while subtracting FP16 calibration endpoints before it widens the result. Retry only a non-finite FP16 scale calculation with FP32 endpoints, then cast the finite scale back to FP16. Existing finite FP16 calculations and all non-FP16 calculations continue to use ONNX Runtime's original result. The fallback intentionally patches only the `qdq_quantizer` binding used for calibrated activation ranges. Initializer and weight quantization continue to use ONNX Runtime's existing `quant_utils` path unchanged; full-range FP16 weight scaling is outside this calibration fix. The regression tests exercise both collectors across initial collection, an equal-range merge, and an expanding-range merge. They also verify the internal FP32 histogram and external FP16 calibration-range contract, including finite saturation when restoring sanitized values. A real entropy calibration test covers full-range FP16 values and verifies finite FP16 Q/DQ scales and a loadable ONNX Runtime graph. The AutoCast integration verifies the same FP16 scale-type contract through the public quantization path. N/A — no API or usage change. All tests ran with CUDA hidden. - NumPy 1.26.4: affected histogram, AutoCast, and ORT-patching modules: `39 passed`. - NumPy 2.2.3: affected histogram, AutoCast, and ORT-patching modules: `39 passed`. - NumPy 2.3.5: affected histogram, AutoCast, and ORT-patching modules: `39 passed`. - Public INT8 entropy quantization with full-range FP16 calibration data: finite FP16 Q/DQ scales, full ONNX check passed, and the CPU ONNX Runtime session loaded. - Changed-file pre-commit hooks: passed. Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ - 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?: ❌ Follow-up to #1558. > 🤖 _Generated by Codex (AI agent)._ --------- Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> Co-authored-by: Codex Signed-off-by: Chad Voegele --- CHANGELOG.rst | 1 + modelopt/onnx/quantization/ort_patching.py | 61 +++++- tests/gpu/onnx/test_ort_patching.py | 15 -- .../test_ort_patching_histogram.py | 182 ++++++++++++++++++ tests/unit/onnx/test_autocast_quantize.py | 13 +- 5 files changed, 249 insertions(+), 23 deletions(-) create mode 100644 tests/unit/onnx/quantization/test_ort_patching_histogram.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 85047aee189..1dd7936723f 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -56,6 +56,7 @@ Changelog **Bug Fixes** +- Fix ONNX INT8 entropy calibration failing or producing invalid quantization parameters for FP16 activations. - Speed up ``mtq.quantize`` on FSDP2-sharded fused-MoE models. Promoting static-block weight quantizers gathered each expert's slice of the fused weight across ranks even though only quantizer state is read, adding a collective per expert to calibration. - Fix ONNX AutoCast failing on models with external initializers larger than 2 GiB. - Avoid querying CUDA/Blackwell capability when ``NVFP4QTensor.quantize`` uses its CPU path or has the optional TensorRT-LLM fast path disabled. diff --git a/modelopt/onnx/quantization/ort_patching.py b/modelopt/onnx/quantization/ort_patching.py index f10b97d13d1..d2957ec183f 100755 --- a/modelopt/onnx/quantization/ort_patching.py +++ b/modelopt/onnx/quantization/ort_patching.py @@ -52,7 +52,7 @@ import onnxruntime as ort import pynvml from onnx import onnx_pb -from onnxruntime.quantization import calibrate +from onnxruntime.quantization import calibrate, qdq_quantizer from onnxruntime.quantization.base_quantizer import BaseQuantizer from onnxruntime.quantization.calibrate import ( CalibraterBase, @@ -74,6 +74,7 @@ QuantType, add_infer_metadata, ) +from onnxruntime.quantization.quant_utils import compute_scale_zp as _ort_compute_scale_zp from onnxruntime.quantization.quantize import check_static_quant_arguments from onnxruntime.quantization.registry import QDQRegistry, QLinearOpsRegistry from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference @@ -94,6 +95,53 @@ def load_model_with_shape_infer(model_path: Path) -> onnx.ModelProto: return model +def _compute_scale_zp(rmin, rmax, qmin, qmax, symmetric=False, min_real_range=None): + """Retry FP16 scale calculation in FP32 when range subtraction overflows.""" + range_dtype = np.asarray(rmax).dtype + if range_dtype != np.float16: + return _ort_compute_scale_zp(rmin, rmax, qmin, qmax, symmetric, min_real_range) + + with np.errstate(over="ignore", invalid="ignore"): + zero_point, scale = _ort_compute_scale_zp(rmin, rmax, qmin, qmax, symmetric, min_real_range) + if np.all(np.isfinite(scale)): + return zero_point, scale + + zero_point, scale = _ort_compute_scale_zp( + np.asarray(rmin, dtype=np.float32), + np.asarray(rmax, dtype=np.float32), + qmin, + qmax, + symmetric, + min_real_range, + ) + return zero_point, np.asarray(scale, dtype=range_dtype) + + +def _prepare_histogram_data(histogram_collector, tensor, data_arr): + """Use FP32 for histogram math while remembering the source dtype.""" + if data_arr.dtype != np.float16: + return data_arr + + original_dtypes = getattr(histogram_collector, "_modelopt_original_dtypes", {}) + original_dtypes[tensor] = data_arr.dtype + histogram_collector._modelopt_original_dtypes = original_dtypes + return data_arr.astype(np.float32) + + +def _restore_histogram_calibration_dtypes(histogram_collector, tensors_range): + """Restore source dtypes at the calibration-to-quantization boundary.""" + original_dtypes = getattr(histogram_collector, "_modelopt_original_dtypes", {}) + for tensor, dtype in original_dtypes.items(): + if tensor not in tensors_range: + continue + tensor_data = tensors_range[tensor] + dtype_limits = np.finfo(dtype) + for attribute in ("lowest", "highest", "avg", "std"): + if hasattr(tensor_data, attribute): + value = np.clip(getattr(tensor_data, attribute), dtype_limits.min, dtype_limits.max) + setattr(tensor_data, attribute, np.asarray(value, dtype=dtype)) + + def _collect_value(histogram_collector, name_to_arr): """Collect histogram on real value.""" for tensor, data_arr in tqdm(name_to_arr.items()): @@ -105,6 +153,7 @@ def _collect_value(histogram_collector, name_to_arr): curr_data_arr = curr_data_arr.flatten() concat_data_arr = np.concatenate((concat_data_arr, curr_data_arr)) + concat_data_arr = _prepare_histogram_data(histogram_collector, tensor, concat_data_arr) data_arr = concat_data_arr # ========================================================== if data_arr.size > 0: @@ -130,9 +179,6 @@ def _collect_value(histogram_collector, name_to_arr): old_histogram, data_arr, min_value, max_value, threshold ) else: - # Cast range endpoints to Python float so numpy computes bin edges in - # float64. A fp16 threshold here can underflow the 128-bin linspace - # and trip "Too many bins for data range" on numpy >= 2.0. range_max = float(threshold) hist, hist_edges = np.histogram( data_arr, histogram_collector.num_bins, range=(-range_max, range_max) @@ -1126,6 +1172,7 @@ def _collect_value_histogram_collector_single_node_calibration(histogram_collect """Collect histogram on real value.""" for tensor, data_arr in name_to_arr.items(): data_arr = np.asarray(data_arr).flatten() + data_arr = _prepare_histogram_data(histogram_collector, tensor, data_arr) min_value, max_value = (np.min(data_arr), np.max(data_arr)) if data_arr.size > 0 else (0, 0) # Replace inf/nan with float32 min/max @@ -1147,9 +1194,6 @@ def _collect_value_histogram_collector_single_node_calibration(histogram_collect threshold, ) else: - # Cast range endpoints to Python float so numpy computes bin edges in - # float64. A fp16 threshold here can underflow the 128-bin linspace - # and trip "Too many bins for data range" on numpy >= 2.0. range_max = float(threshold) hist, hist_edges = np.histogram( data_arr, histogram_collector.num_bins, range=(-range_max, range_max) @@ -1685,6 +1729,8 @@ def _quantize_static( raise TypeError( f"Unexpected type {type(tensors_range)} for tensors_range and calibrator={type(calibrator)}." ) + if isinstance(calibrator, HistogramCalibrater): + _restore_histogram_calibration_dtypes(calibrator.collector, tensors_range) del calibrator check_static_quant_arguments(quant_format, activation_type, weight_type) @@ -1795,4 +1841,5 @@ def patch_ort_modules(calibrate_per_node: bool = False): CalibraterBase.select_tensors_to_calibrate = _select_tensors_to_calibrate QDQQuantizer.check_opset_version = _check_opset_version BaseQuantizer.adjust_tensor_ranges = _adjust_tensor_ranges + qdq_quantizer.compute_scale_zp = _compute_scale_zp CalibraterBase.__init__ = _init_calibrater_base diff --git a/tests/gpu/onnx/test_ort_patching.py b/tests/gpu/onnx/test_ort_patching.py index 26ca49f39b4..84224dcffa0 100644 --- a/tests/gpu/onnx/test_ort_patching.py +++ b/tests/gpu/onnx/test_ort_patching.py @@ -153,21 +153,6 @@ def test_collect_value(self, mock_histogram_collector, sample_tensor_data): assert "tensor1" in mock_histogram_collector.histogram_dict assert "tensor2" in mock_histogram_collector.histogram_dict - def test_collect_value_fp16_narrow_range(self, mock_histogram_collector): - # fp16 activations with a small range (threshold ~1e-5) used to raise - # "Too many bins for data range" on numpy >= 2.0, because the fp16 range - # produced a fp16 linspace where consecutive bin edges rounded together. - activations = np.zeros(1000, dtype=np.float16) - activations[0] = np.float16(1e-5) - name_to_arr = {"narrow_fp16_tensor": [activations]} - - _collect_value(mock_histogram_collector, name_to_arr) - - hist, edges, _, _, _ = mock_histogram_collector.histogram_dict["narrow_fp16_tensor"] - assert hist.sum() == activations.size - assert len(edges) == mock_histogram_collector.num_bins + 1 - assert not np.any(np.diff(edges) == 0), "fp16 bin edges collapsed" - def test_collect_absolute_value(self, mock_histogram_collector, sample_tensor_data): """Test _collect_absolute_value function.""" # Convert to float32 to avoid the float64 assertion error diff --git a/tests/unit/onnx/quantization/test_ort_patching_histogram.py b/tests/unit/onnx/quantization/test_ort_patching_histogram.py new file mode 100644 index 00000000000..e986454bb57 --- /dev/null +++ b/tests/unit/onnx/quantization/test_ort_patching_histogram.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. + +"""Tests for ONNX Runtime histogram quantization patches.""" + +import numpy as np +import onnx +import onnxruntime as ort +import pytest +from onnx import TensorProto, helper, numpy_helper +from onnxruntime.quantization.calibrate import ( + CalibrationDataReader, + CalibrationMethod, + HistogramCollector, + TensorData, + TensorsData, +) + +from modelopt.onnx.quantization.ort_patching import ( + _collect_value, + _collect_value_histogram_collector_single_node_calibration, + _compute_scale_zp, + _prepare_histogram_data, + _quantize_static, + _restore_histogram_calibration_dtypes, + patch_ort_modules, +) + + +def test_compute_scale_zp_fp16_overflow_fallback(): + zero_point, scale = _compute_scale_zp( + np.array(-65504, dtype=np.float16), + np.array(65504, dtype=np.float16), + np.array(-128, dtype=np.int8), + np.array(127, dtype=np.int8), + symmetric=True, + ) + + assert zero_point.dtype == np.int8 + assert zero_point == 0 + assert scale.dtype == np.float16 + assert scale == np.float16(514) + + +def test_quantize_static_fp16_high_range_scale(tmp_path): + class HighRangeDataReader(CalibrationDataReader): + def __init__(self): + self.rewind() + + def get_next(self): + return next(self.data, None) + + def rewind(self): + values = np.array([[-65504, 65504, -32752, 32752]], dtype=np.float16) + self.data = iter([{"input": values}]) + + model_path = tmp_path / "model.onnx" + output_path = tmp_path / "model.quant.onnx" + graph = helper.make_graph( + [helper.make_node("MatMul", ["input", "weight"], ["output"], name="matmul")], + "fp16_high_range", + [helper.make_tensor_value_info("input", TensorProto.FLOAT16, [1, 4])], + [helper.make_tensor_value_info("output", TensorProto.FLOAT16, [1, 4])], + [numpy_helper.from_array(np.eye(4, dtype=np.float16), name="weight")], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21)]) + model.ir_version = min(model.ir_version, 10) + onnx.save(model, model_path) + + patch_ort_modules(False) + _quantize_static( + model_path, + output_path, + HighRangeDataReader(), + nodes_to_quantize=["matmul"], + op_types_to_quantize=["MatMul"], + calibrate_method=CalibrationMethod.Entropy, + extra_options={ + "ExecutionProviders": ["CPUExecutionProvider"], + "ActivationSymmetric": True, + "AddQDQPairToWeight": True, + }, + ) + + quantized_model = onnx.load(output_path) + initializers = { + initializer.name: initializer for initializer in quantized_model.graph.initializer + } + scale_initializers = [ + initializers[node.input[1]] + for node in quantized_model.graph.node + if node.op_type in {"QuantizeLinear", "DequantizeLinear"} and node.input[1] in initializers + ] + assert scale_initializers + assert {initializer.data_type for initializer in scale_initializers} == {TensorProto.FLOAT16} + assert all( + np.isfinite(numpy_helper.to_array(initializer)).all() for initializer in scale_initializers + ) + ort.InferenceSession(output_path, providers=["CPUExecutionProvider"]) + + +@pytest.mark.parametrize( + "collect_value", + [ + _collect_value, + _collect_value_histogram_collector_single_node_calibration, + ], +) +def test_collect_value_fp16_narrow_range(collect_value): + collector = HistogramCollector( + method="entropy", + symmetric=False, + num_bins=128, + num_quantized_bins=128, + percentile=None, + scenario="same", + ) + activations = np.zeros(1000, dtype=np.float16) + for activation_max in (1e-6, 1e-6, 2e-6): + activations[0] = np.float16(activation_max) + collect_value(collector, {"narrow_fp16_tensor": [activations]}) + + hist, edges, _, _, threshold = collector.histogram_dict["narrow_fp16_tensor"] + assert hist.sum() == 3 * activations.size + assert len(hist) > collector.num_bins + assert edges.dtype == np.float32 + assert np.all(np.diff(edges) > 0), "fp16 bin edges are not strictly increasing" + assert np.asarray(threshold).dtype.itemsize >= np.dtype(np.float32).itemsize + + tensors_range = TensorsData(CalibrationMethod.Entropy, collector.compute_collection_result()) + _restore_histogram_calibration_dtypes(collector, tensors_range) + tensor_range = tensors_range["narrow_fp16_tensor"] + assert tensor_range.lowest.dtype == np.float16 + assert tensor_range.highest.dtype == np.float16 + assert tensor_range.bins.dtype == np.float32 + + +def test_restore_histogram_calibration_dtypes_clamps_fp16(): + collector = HistogramCollector( + method="distribution", + symmetric=False, + num_bins=512, + num_quantized_bins=128, + percentile=None, + scenario="same", + ) + _prepare_histogram_data(collector, "tensor", np.array([], dtype=np.float16)) + _prepare_histogram_data(collector, "missing_tensor", np.array([], dtype=np.float16)) + + fp32_max = np.finfo(np.float32).max + tensor_data = TensorData( + lowest=np.float32(-fp32_max), + highest=np.float32(fp32_max), + avg=np.float32(fp32_max), + std=np.float32(fp32_max), + hist=np.array([1]), + hist_edges=np.array([-1, 1], dtype=np.float32), + ) + tensors_range = TensorsData(CalibrationMethod.Distribution, {"tensor": tensor_data}) + + _restore_histogram_calibration_dtypes(collector, tensors_range) + + fp16_limits = np.finfo(np.float16) + tensor_range = tensors_range["tensor"] + assert tensor_range.lowest == fp16_limits.min + assert tensor_range.highest == fp16_limits.max + assert tensor_range.avg == fp16_limits.max + assert tensor_range.std == fp16_limits.max + assert tensor_range.hist_edges.dtype == np.float32 + assert "missing_tensor" not in tensors_range diff --git a/tests/unit/onnx/test_autocast_quantize.py b/tests/unit/onnx/test_autocast_quantize.py index bc123aad9dc..930187980f8 100644 --- a/tests/unit/onnx/test_autocast_quantize.py +++ b/tests/unit/onnx/test_autocast_quantize.py @@ -55,7 +55,18 @@ def test_autocast_quantize_int8(tmp_path, keep_io_types, bias_add): assert os.path.isfile(output_onnx_path) # Load the output model and check QDQ node placements - graph = gs.import_onnx(onnx.load(output_onnx_path)) + quantized_model = onnx.load(output_onnx_path) + graph = gs.import_onnx(quantized_model) + + activation_scale_names = { + node.input[1] for node in quantized_model.graph.node if node.op_type == "QuantizeLinear" + } + activation_scale_types = { + initializer.data_type + for initializer in quantized_model.graph.initializer + if initializer.name in activation_scale_names + } + assert activation_scale_types == {onnx.TensorProto.FLOAT16} # Check that all MatMul nodes are quantized mm_nodes = [n for n in graph.nodes if n.op == "MatMul"] From 72a53b0da3afa0144ad0dee4604fd3320297f845 Mon Sep 17 00:00:00 2001 From: Joshua <97406952+joshua-hill@users.noreply.github.com> Date: Sun, 13 Sep 2026 10:55:32 -0400 Subject: [PATCH 13/28] Fix distributed AutoQuantize scoring and share backward setup (#2231) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix AutoQuantize can measure a group of quantized expert layers at their enclosing MLP output. That enclosing module is often a plain PyTorch container and does not carry distributed-group information, so its sensitivity score was not combined across data- or expert-parallel workers. This PR obtains the distributed groups from the quantized layers when the scoring module does not provide them. It also preserves construction order for quantized modules, scoring modules, and their registered hyperparameters so every worker accumulates scores in the same order. The temporary state needed by backward-based scoring is now managed by one shared session. The session installs and removes forward patches and invocation-specific output-gradient hooks, controls parameter gradients, and restores the active quantization recipes even when scoring raises an exception. Scoring methods remain responsible for their own score calculation. ### Usage N/A — this fixes existing AutoQuantize behavior and does not add an API or flag. ### Testing - `pre-commit run --files modelopt/torch/quantization/algorithms.py tests/unit/torch/quantization/test_autoquant.py` - `pytest -q tests/unit/torch/quantization/test_autoquant.py` — 102 passed - Added a real two-rank gradient AutoQuantize test covering MoE experts scored at an enclosing MLP. - Added regressions for deterministic hyperparameter registration, per-invocation replay for reused score modules, exact `forward`-attribute restoration, partial setup rollback, and cleanup after a scoring failure. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors). - Is this change backward compatible?: ✅ — no API or checkpoint format changes; distributed sensitivity values now include the missing reduction. - 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)?: N/A — no new feature, deprecation, breaking change, or critical release-note item. - Did you get Claude approval on this PR?: ❌ — pending review. ## Summary by CodeRabbit * **Bug Fixes** * Improved quantization scoring consistency through deterministic ordering and invocation handling. * Added more reliable distributed score aggregation, including support for mixture-of-experts models. * Improved gradient-based scoring for repeated evaluations, tuple outputs, and checkpoint-compatible workflows. * Ensured model behavior and scoring state are restored after successful or failed evaluations. * Avoided unnecessary output replay when gradients are not required. --------- Signed-off-by: Joshua Hill Co-authored-by: Cursor Signed-off-by: Chad Voegele --- modelopt/torch/quantization/algorithms.py | 447 +++++++++++------- .../unit/torch/quantization/test_autoquant.py | 277 +++++++++++ 2 files changed, 550 insertions(+), 174 deletions(-) diff --git a/modelopt/torch/quantization/algorithms.py b/modelopt/torch/quantization/algorithms.py index 7beeef6ad7f..133f278d24a 100644 --- a/modelopt/torch/quantization/algorithms.py +++ b/modelopt/torch/quantization/algorithms.py @@ -23,7 +23,7 @@ from abc import ABC, abstractmethod from collections import defaultdict from collections.abc import Callable, Sequence -from contextlib import nullcontext +from contextlib import ExitStack, nullcontext from typing import Any import regex as re @@ -413,8 +413,10 @@ def __init__( self.allow_no_quant = allow_no_quant self.is_fixed = fixed_recipe is not None - self.quant_modules = list(set(quant_modules or [])) - self.score_modules = list(set(score_modules or self.quant_modules)) + # Module hashes depend on object identity, so sets can produce different orders per rank. + self.quant_modules = list(dict.fromkeys(quant_modules or [])) + self.score_modules = list(dict.fromkeys(score_modules or self.quant_modules)) + self._warned_parallel_state_fallbacks: set[nn.Module] = set() fixed_quantizers = ( { @@ -467,11 +469,11 @@ def __init__( quant_recipe: dict.fromkeys(self.score_modules) for quant_recipe in self.choices } - # Attach this hparam to each score_module's set of hparams it scores + # Registration order follows the rank-stable runtime-group construction order. for score_module in self.score_modules: if not hasattr(score_module, "_hparams_for_scoring"): - score_module._hparams_for_scoring = set() - score_module._hparams_for_scoring.add(self) + score_module._hparams_for_scoring = [] + score_module._hparams_for_scoring.append(self) @property def active(self) -> HPType: @@ -532,6 +534,31 @@ def get_score(self, recipe: QuantRecipe) -> float: continue parallel_state = getattr(score_module, "parallel_state", None) + if parallel_state is None: + # TODO: Prefer parallel_state owned by the score module; this temporary fallback + # inherits the first quantized child's state and assumes all grouped quant modules + # share the same parallel groups. + parallel_state_source = next( + ( + (module, state) + for module in self.quant_modules + if (state := getattr(module, "parallel_state", None)) is not None + ), + None, + ) + if parallel_state_source is not None: + quant_module, parallel_state = parallel_state_source + if ( + torch.distributed.is_initialized() + and score_module not in self._warned_parallel_state_fallbacks + ): + warnings.warn( + "Distributed training is initialized but no parallel_state is set for " + f"score module {type(score_module)}. Using parallel_state from its first " + f"quantized child {type(quant_module)}. All grouped quant modules must " + "share the same parallel groups." + ) + self._warned_parallel_state_fallbacks.add(score_module) if parallel_state is None: total_score += importance.cpu().item() @@ -1438,7 +1465,222 @@ def _add_auto_quantize_score(grad_output, output_diff, score_tensor): score_tensor += _get_auto_quantize_score(grad_output, output_diff) -class AutoQuantizeGradientSearcher(_AutoQuantizeBaseSearcher): +class _AutoQuantizeBackwardScoringSession(ABC): + """Manage temporary model state used by activation-backward scoring.""" + + def __init__( + self, + model: nn.Module, + score_modules: Sequence[nn.Module], + is_param_grad_enabled: Callable, + verbose: bool = False, + ) -> None: + self.model = model + self.score_modules = tuple(score_modules) + self.is_param_grad_enabled = is_param_grad_enabled + self.verbose = verbose + self._stack = ExitStack() + self._original_forwards: dict[nn.Module, Callable] = {} + self._output_grad_hook_handles: set[Any] = set() + self._grad_accumulators: list[Any] = [] + + def __enter__(self): + """Install scoring hooks and parameter settings.""" + try: + hparams = list( + dict.fromkeys( + hparam + for module in self.score_modules + for hparam in module._hparams_for_scoring + ) + ) + for hparam in hparams: + self._stack.callback(setattr, hparam, "active", hparam.active) + + def patched_forward(module, *args, **kwargs): + return self.forward(module, *args, **kwargs) + + for module in self.score_modules: + original_forward = module.forward + self._original_forwards[module] = original_forward + had_instance_forward = "forward" in module.__dict__ + instance_forward = module.__dict__.get("forward") + module.forward = types.MethodType(patched_forward, module) + if had_instance_forward: + self._stack.callback(setattr, module, "forward", instance_forward) + else: + self._stack.callback(module.__dict__.pop, "forward", None) + + for name, param in self.model.named_parameters(): + requires_grad = param.requires_grad + enable_grad = self.is_param_grad_enabled(name, self.model) + param.requires_grad = enable_grad + self._stack.callback(setattr, param, "requires_grad", requires_grad) + if not enable_grad: + continue + if self.verbose: + print_rank_0(f"AutoQuantize: Enabling gradient for param {name}.") + accumulator, hook = create_param_grad_clear_hook(param) + self._grad_accumulators.append(accumulator) + self._stack.callback(hook.remove) + except Exception: + self._stack.close() + raise + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + """Restore all model state changed for scoring.""" + self._clear_output_grad_hooks() + self._stack.close() + self._original_forwards.clear() + self._grad_accumulators.clear() + + def original_forward(self, module: nn.Module) -> Callable: + """Return the forward method saved before scoring.""" + return self._original_forwards[module] + + def _clear_output_grad_hooks(self) -> None: + """Remove output hooks whose backward pass has not run.""" + for handle in self._output_grad_hook_handles: + handle.remove() + self._output_grad_hook_handles.clear() + + def _register_output_grad_hook(self, output: torch.Tensor, hook: Callable) -> None: + """Attach an invocation-specific output-gradient hook for this session.""" + + def run_once(grad): + try: + return hook(grad) + finally: + handle.remove() + self._output_grad_hook_handles.discard(handle) + + handle = output.register_hook(run_once) + self._output_grad_hook_handles.add(handle) + + @abstractmethod + def forward(self, module: nn.Module, *args, **kwargs): + """Run a score module forward pass and collect method-specific state.""" + + +class _AutoQuantizeGradientScoringSession(_AutoQuantizeBackwardScoringSession): + """Collect gradient-based scores while candidate recipes are replayed.""" + + def forward(self, module: nn.Module, *args, **kwargs): + """Run the reference forward and cache each recipe's output perturbation.""" + no_quant_recipe = QuantRecipe(quant_cfg=None) + for hparam in module._hparams_for_scoring: + if hparam.is_configurable: + hparam.active = no_quant_recipe + + output = self.original_forward(module)(*args, **kwargs) + + # Checkpointed modules recompute with gradients enabled during backward. + base_output = output[0] if isinstance(output, tuple) else output + if not torch.is_grad_enabled() or not base_output.requires_grad: + return output + + output_diffs = {hparam: {} for hparam in module._hparams_for_scoring} + with torch.no_grad(): + for hparam in module._hparams_for_scoring: + if not hparam.is_configurable: + continue + for recipe in hparam.choices: + if recipe == no_quant_recipe: + continue + hparam.active = recipe + replay = self.original_forward(module)(*args, **kwargs) + output_diff = ( + replay[0] - output[0] if isinstance(replay, tuple) else replay - output + ) + output_diffs[hparam][recipe] = output_diff.detach() + hparam.active = no_quant_recipe + + self._register_output_grad_hook( + base_output, + lambda grad_output: self._accumulate_scores(module, output_diffs, grad_output), + ) + return output + + def _accumulate_scores(self, module, invocation_diffs, grad_output) -> None: + """Accumulate scores for the invocation that produced ``grad_output``.""" + for hparam, output_diffs in invocation_diffs.items(): + for recipe, output_diff in output_diffs.items(): + importance = hparam._importance_dict[recipe][module] + if importance is None: + hparam._importance_dict[recipe][module] = _get_auto_quantize_score( + grad_output, output_diff + ) + else: + _add_auto_quantize_score(grad_output, output_diff, importance) + + +class _AutoQuantizeBackwardScoringSearcher(_AutoQuantizeBaseSearcher): + """Share orchestration used by activation-backward scoring methods.""" + + score_module_rules = [ + # Score MoE projections together at their enclosing MLP or mixer output. + r"^(.*?\.mlp)\.experts\.\d+\.(gate_proj|up_proj|down_proj)$", + r"^(.*?\.mixer)\.experts\.\d+\.(up_proj|down_proj)$", + r"^(.*?)\.(\d+\.(w1|w2|w3))$", + r"^(.*?)\.((w1_linear|w2_linear|w3_linear)\.\d+)$", + ] + + _custom_support: list[tuple[Callable, Callable, Callable]] = [] + + @classmethod + def register_custom_support( + cls, + is_supported_checker: Callable, + grad_ckpt_context: Callable, + is_param_grad_enabled: Callable, + ) -> None: + """Register optional hooks for memory-efficient backward scoring. + + `is_supported_checker` selects models that use these hooks. + `grad_ckpt_context` enables their gradient-checkpointing context, and + `is_param_grad_enabled` selects the minimum parameters needed to propagate + activation gradients. + """ + cls._custom_support.append((is_supported_checker, grad_ckpt_context, is_param_grad_enabled)) + + def _configurable_score_modules(self) -> list[nn.Module]: + return [ + module + for module in self.model.modules() + if hasattr(module, "_hparams_for_scoring") + and any(hparam.is_configurable for hparam in module._hparams_for_scoring) + ] + + @abstractmethod + def _estimate_auto_quantize_scores(self, is_param_grad_enabled: Callable) -> None: + """Estimate scores while activation gradients are enabled.""" + + def estimate_sensitivity_scores(self) -> None: + """Run backward scoring with the first matching model-specific support hook.""" + self.model.eval() + + def default_is_param_grad_enabled(_name, _model): + return True + + grad_checkpointing_context = None + is_param_grad_enabled = default_is_param_grad_enabled + for is_supported, context_candidate, grad_candidate in self._custom_support: + if is_supported(self.model): + grad_checkpointing_context = context_candidate + is_param_grad_enabled = grad_candidate + break + + context = ( + grad_checkpointing_context(self.model) + if grad_checkpointing_context is not None + else nullcontext() + ) + with context: + self._estimate_auto_quantize_scores(is_param_grad_enabled) + + +class AutoQuantizeGradientSearcher(_AutoQuantizeBackwardScoringSearcher): """A searcher for AutoQuantize algorithm that uses gradient based score estimation. In AutoQuantize, we search for the best per-layer quantization configuration that minimizes the sum of per-layer @@ -1472,17 +1714,6 @@ class AutoQuantizeGradientSearcher(_AutoQuantizeBaseSearcher): method_name = "gradient" - score_module_rules = [ - # Use MLP layer output for gate_proj, up_proj, down_proj for Qwen3 like MoE models (local and shared experts) - r"^(.*?\.mlp)\.experts\.\d+\.(gate_proj|up_proj|down_proj)$", - r"^(.*?\.mixer)\.experts\.\d+\.(up_proj|down_proj)$", # NemotronH MoE experts - r"^(.*?)\.(\d+\.(w1|w2|w3))$", # mixtral experts - r"^(.*?)\.((w1_linear|w2_linear|w3_linear)\.\d+)$", # dbrx experts - ] - - # See `register_custom_support` for details - _custom_support: list[tuple[Callable, Callable, Callable]] = [] - @property def default_search_config(self): """Get the default config for the searcher.""" @@ -1511,30 +1742,6 @@ def sanitize_search_config(self, config: SearchConfig | None) -> SearchConfig: return config - @classmethod - def register_custom_support( - cls, - is_supported_checker: Callable, - grad_ckpt_context: Callable, - is_param_grad_enabled: Callable, - ) -> None: - """(Optional) Register custom support for `AutoQuantize` score estimation. - - This custom support is used to enable memory/compute efficient backward gradient propagation. This involves: - - - `grad_ckpt_context`: backward pass with gradient checkpointing enabled - - `is_param_grad_enabled`: AutoQuantize only needs activation gradients to be computed (not weight - gradients). `is_param_grad_enabled` is used to select which parameters should have gradients enabled, - limiting gradient computation to only what's needed for activation gradients. For LLMs, to trigger all - activation gradient computation, just enabling the embedding layer weight gradient is sufficient. This will - enable gradient computation for all the activation gradients downstream. - - If the `is_supported_checker(model)` returns True, the `grad_ckpt_context(model)` will be - used to enable gradient checkpointing and `is_param_grad_enabled(pname, model)` - will be used to select which parameters have gradients enabled to minimize gradient computation. - """ - cls._custom_support.append((is_supported_checker, grad_ckpt_context, is_param_grad_enabled)) - def _get_default_forward_backward_step(self): def forward_backward_step(model, data): output = self.config["forward_step"](model, data) @@ -1552,143 +1759,35 @@ def forward_backward_step(model, data): @torch.enable_grad() def _estimate_auto_quantize_scores(self, is_param_grad_enabled): - # TODO: remove the no-quant recipe - def auto_quantize_score_estimate_forward(module, input, *args, **kwargs): - for hparam in module._hparams_for_scoring: - if hparam.is_configurable: - hparam.active = QuantRecipe(quant_cfg=None) - - output = module._forward_original(input, *args, **kwargs) - - # If gradient checkpointing is enabled, gradient will not be enabled in the global forward pass. - # With gradient checkpointing, gradients are computed in the local forward pass during backward pass - - # Lets compute the output_diff and save it in memory only if gradient is enabled to be memory efficient - if not torch.is_grad_enabled(): - return output - - module.output_diff_dict = {hparam: {} for hparam in module._hparams_for_scoring} - with torch.no_grad(): - for hparam in module._hparams_for_scoring: - if not hparam.is_configurable: - continue - for recipe in hparam.choices: - if recipe == QuantRecipe(quant_cfg=None): - continue - hparam.active = recipe - output_diff = module._forward_original(input, *args, **kwargs) - - if isinstance(output_diff, tuple): - output_diff = output_diff[0] - output[0] - else: - output_diff -= output - module.output_diff_dict[hparam][recipe] = output_diff.detach() - - # Disable the configurable hparam now that we have computed the diff - hparam.active = QuantRecipe(quant_cfg=None) - - return output - - def backward_hook(module, grad_input, grad_output): - for hparam, output_diff_dict in module.output_diff_dict.items(): - for recipe, output_diff in output_diff_dict.items(): - if hparam._importance_dict[recipe][module] is None: - hparam._importance_dict[recipe][module] = _get_auto_quantize_score( - grad_output[0], output_diff - ) - else: - _add_auto_quantize_score( - grad_output[0], output_diff, hparam._importance_dict[recipe][module] - ) - - def setup_params_for_score_estimation(name, param, params_metadata, enable_grad=True): - # Let us delete the gradient as soon as they are computed to save memory - params_metadata[name] = {"requires_grad": param.requires_grad} - param.requires_grad = enable_grad - if not enable_grad: - return - if self.config.get("verbose", False): - print_rank_0(f"AutoQuantize: Enabling gradient for param {name}.") - accum_grad, handle = create_param_grad_clear_hook(param) - params_metadata[name]["accum_grad"] = accum_grad # We need to keep the accum_grad alive - params_metadata[name]["handle"] = handle - - def setup_module_for_score_estimation(module): - module._forward_original = module.forward - module.forward = types.MethodType(auto_quantize_score_estimate_forward, module) - module._backward_hook_handle = module.register_full_backward_hook(backward_hook) - - def cleanup_module_after_score_estimation(module): - module.forward = module._forward_original - del module._forward_original - - module._backward_hook_handle.remove() - - def cleanup_params_after_score_estimation(name, param, params_metadata): - param.requires_grad = params_metadata[name]["requires_grad"] - handle = params_metadata[name].get("handle", None) - if handle is not None: - handle.remove() - - score_modules = set() - for name, module in self.model.named_modules(): - if ( - hasattr(module, "_hparams_for_scoring") - and any(hparam.is_configurable for hparam in module._hparams_for_scoring) - and module not in score_modules - ): - # Monkey patch the forward methods to cache (Q(Y) - Y) - setup_module_for_score_estimation(module) - score_modules.add(module) - - params_metadata = {} - for name, param in self.model.named_parameters(): - setup_params_for_score_estimation( - name, param, params_metadata, is_param_grad_enabled(name, self.model) + score_modules = self._configurable_score_modules() + with _AutoQuantizeGradientScoringSession( + self.model, + score_modules, + is_param_grad_enabled, + verbose=self.config.get("verbose", False), + ) as scoring_session: + gc.collect() + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + report_memory("AutoQuantize: starting score estimation, ") + + def score_step(model, data): + try: + return self.config["forward_backward_step"](model, data) + finally: + scoring_session._clear_output_grad_hooks() + + self._run_func( + score_step, + num_iters=self.config["num_score_steps"], + desc="Estimating auto_quantize scores", ) - gc.collect() - if torch.cuda.is_available(): - torch.cuda.reset_peak_memory_stats() - report_memory("AutoQuantize: starting score estimation, ") - - self._run_func( - self.config["forward_backward_step"], - num_iters=self.config["num_score_steps"], - desc="Estimating auto_quantize scores", - ) - - if torch.cuda.is_available(): - report_memory("AutoQuantize: After score estimation") + if torch.cuda.is_available(): + report_memory("AutoQuantize: After score estimation") - for module in score_modules: - cleanup_module_after_score_estimation(module) - - for name, param in self.model.named_parameters(): - cleanup_params_after_score_estimation(name, param, params_metadata) - - # Delete the params_metadata - del params_metadata gc.collect() - def estimate_sensitivity_scores(self) -> None: - """Estimate sensitivity scores using hessian approximation.""" - self.model.eval() - - def _default_is_param_grad_enabled(pname, model): - return True - - grad_checkpointing_ctxt = None - is_param_grad_enabled = _default_is_param_grad_enabled - for is_supported_checker, ctxt_candidate, grad_enabled_candidate in self._custom_support: - if is_supported_checker(self.model): - grad_checkpointing_ctxt = ctxt_candidate - is_param_grad_enabled = grad_enabled_candidate - break - - with grad_checkpointing_ctxt(self.model) if grad_checkpointing_ctxt else nullcontext(): - self._estimate_auto_quantize_scores(is_param_grad_enabled) - def run_search_with_stats(self, max_weight_size, verbose=False): """Linear Programming Solve for gradient based auto_quantize. diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py index e83f7fa0a70..0013f2da538 100644 --- a/tests/unit/torch/quantization/test_autoquant.py +++ b/tests/unit/torch/quantization/test_autoquant.py @@ -15,6 +15,7 @@ import copy import io +import warnings from types import SimpleNamespace import pytest @@ -36,6 +37,7 @@ QuantRecipe, QuantRecipeHparam, _AutoQuantizeBaseSearcher, + _AutoQuantizeGradientScoringSession, _module_search_space_signature, estimate_quant_compression, ) @@ -98,6 +100,47 @@ def get_input(self): return torch.randn(1, 4, 32) +class _ScoredMoeExpert(torch.nn.Module): + def __init__(self): + super().__init__() + self.gate_proj = torch.nn.Linear(8, 8) + self.up_proj = torch.nn.Linear(8, 8) + self.down_proj = torch.nn.Linear(8, 8) + + def forward(self, x): + return self.down_proj(self.gate_proj(x) + self.up_proj(x)) + + +class _ScoredMoeMlp(torch.nn.Module): + def __init__(self): + super().__init__() + self.experts = torch.nn.ModuleList([_ScoredMoeExpert(), _ScoredMoeExpert()]) + + def forward(self, x): + output = torch.zeros_like(x) + for expert in self.experts: + output = output + expert(x) + return output + + +class _ScoredMoeLayer(torch.nn.Module): + def __init__(self): + super().__init__() + self.mlp = _ScoredMoeMlp() + + def forward(self, x): + return self.mlp(x) + + +class _ScoredMoeModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.layers = torch.nn.ModuleList([_ScoredMoeLayer()]) + + def forward(self, x): + return self.layers[0](x) + + @pytest.mark.parametrize( ("quant_cfg", "other_quant_cfg", "is_less_than"), [ @@ -905,6 +948,240 @@ def test_data_parallel_auto_quantize(skip_on_windows): spawn_multiprocess_job(2, _test_data_parallel_auto_quantize, backend="gloo") +def _test_data_parallel_moe_score_module(rank, size): + torch.manual_seed(1234) + model = _ScoredMoeModel() + data_loader = [torch.randn(2, 3, 8) for _ in range(2)] + with warnings.catch_warnings(record=True) as caught_warnings: + warnings.simplefilter("always") + model, search_history = mtq.auto_quantize( + model, + constraints={"effective_bits": 12.0}, + quantization_formats=[mtq.INT8_DEFAULT_CFG], + data_loader=data_loader, + forward_step=lambda model, batch: model(batch), + loss_func=lambda output, data: output.square().mean(), + num_calib_steps=2, + num_score_steps=2, + ) + + hparam = model.layers[0].mlp.experts[0].gate_proj.get_hparam("quant_recipe") + assert hparam.score_modules == [model.layers[0].mlp] + assert isinstance(model.layers[0].mlp._hparams_for_scoring, list) + + fallback_warnings = [ + str(warning.message) + for warning in caught_warnings + if "no parallel_state is set for score module" in str(warning.message) + ] + assert len(fallback_warnings) == len(model.layers[0].mlp._hparams_for_scoring) + assert all(str(type(model.layers[0].mlp)) in message for message in fallback_warnings) + assert all("first quantized child" in message for message in fallback_warnings) + + recipe = QuantRecipe(mtq.INT8_DEFAULT_CFG) + local_score = sum(hparam._importance_dict[recipe][m] for m in hparam.score_modules) + candidate = next( + candidate + for candidate in search_history["candidate_stats"].values() + if "layers.0.mlp.experts.0.gate_proj" in candidate["module_names"] + ) + recipe_idx = candidate["formats"].index(recipe) + torch.testing.assert_close( + local_score * size, + torch.tensor( + candidate["scores"][recipe_idx], + device=local_score.device, + dtype=local_score.dtype, + ), + ) + + scores = { + name: candidate["scores"] for name, candidate in search_history["candidate_stats"].items() + } + rank_zero_scores = DistributedProcessGroup.get_dist_syncd_obj( + scores if rank == 0 else None, + DistributedProcessGroup(None), + lambda values: values[0], + ) + assert scores == rank_zero_scores + + +def test_data_parallel_moe_score_module(skip_on_windows): + spawn_multiprocess_job(2, _test_data_parallel_moe_score_module, backend="gloo") + + +def test_score_hparam_registration_preserves_order(): + quant_modules = [ + mtq.quantize(torch.nn.Linear(4, 4), mtq.INT8_DEFAULT_CFG), + mtq.quantize(torch.nn.Linear(4, 4), mtq.INT8_DEFAULT_CFG), + ] + score_module = torch.nn.Identity() + recipe = QuantRecipe(mtq.INT8_DEFAULT_CFG) + + first = QuantRecipeHparam( + [recipe], + quant_modules=[quant_modules[0], quant_modules[1], quant_modules[0]], + score_modules=[score_module, score_module], + ) + second = QuantRecipeHparam( + [recipe], + quant_modules=[quant_modules[1]], + score_modules=[score_module], + ) + + assert first.quant_modules == quant_modules + assert first.score_modules == [score_module] + assert score_module._hparams_for_scoring == [first, second] + + +def test_gradient_scoring_tracks_reused_module_invocations(): + """Each autograd use of a shared score module retains its own replay difference.""" + no_quant_recipe = QuantRecipe(quant_cfg=None) + quant_recipe = QuantRecipe(mtq.INT8_DEFAULT_CFG) + + class TestHparam: + is_configurable = True + choices = [no_quant_recipe, quant_recipe] + active = no_quant_recipe + + class ScoreModule(torch.nn.Module): + def forward(self, x): + scale = 1.0 if hparam.active == no_quant_recipe else 2.0 + return scale * x + + class ReusedScoreModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.shared = ScoreModule() + + def forward(self, x): + first = self.shared(x) + second = self.shared(3.0 * x) + # This use does not participate in autograd. + self.shared(5.0 * x.detach()) + # This output requires grad but is intentionally unused by the loss. + self.shared(7.0 * x) + return first.sum() + 2.0 * second.sum() + + model = ReusedScoreModule() + score_module = model.shared + hparam = TestHparam() + hparam._importance_dict = {recipe: {score_module: None} for recipe in hparam.choices} + score_module._hparams_for_scoring = [hparam] + inputs = torch.tensor([[1.0, 2.0]], requires_grad=True) + + hparam._importance_dict[quant_recipe][score_module] = None + delayed_session = _AutoQuantizeGradientScoringSession(model, [score_module], lambda *_: True) + with delayed_session: + delayed_loss = model(inputs) + assert delayed_session._output_grad_hook_handles + + assert not delayed_session._output_grad_hook_handles + delayed_loss.backward() + assert hparam._importance_dict[quant_recipe][score_module] is None + + with _AutoQuantizeGradientScoringSession(model, [score_module], lambda *_: True): + model(inputs).backward() + + # First use: sum(x**2) = 5. Second use: sum((2 * 3x)**2) = 180. + importance = hparam._importance_dict[quant_recipe][score_module] + torch.testing.assert_close( + importance, + torch.tensor(185.0, device=importance.device), + rtol=0, + atol=1e-6, + ) + + +def test_gradient_scoring_restores_model_after_failure(): + model = SimpleLinear() + patched_modules = [] + + def fail_during_scoring(model, data): + model(data) + patched_modules.extend( + module + for module in model.modules() + if getattr(module.forward, "__name__", None) == "patched_forward" + ) + raise RuntimeError("stop after scoring forward") + + with pytest.raises(RuntimeError, match="stop after scoring forward"): + mtq.auto_quantize( + model, + constraints={"effective_bits": 12.0}, + quantization_formats=[mtq.INT8_DEFAULT_CFG], + data_loader=[model.get_input()], + forward_step=lambda model, batch: model(batch), + forward_backward_step=fail_during_scoring, + num_calib_steps=1, + num_score_steps=1, + ) + + assert patched_modules + assert all( + getattr(module.forward, "__name__", None) != "patched_forward" for module in patched_modules + ) + assert all("forward" not in module.__dict__ for module in patched_modules) + assert all(param.requires_grad for param in model.parameters()) + for module in model.modules(): + for hparam in getattr(module, "_hparams_for_scoring", []): + assert hparam.active == hparam.original + + +def test_backward_scoring_session_restores_partial_setup(): + model = torch.nn.Sequential(torch.nn.Linear(4, 4)) + score_module = model[0] + score_module._hparams_for_scoring = [] + score_module.weight.requires_grad = False + original_requires_grad = {name: param.requires_grad for name, param in model.named_parameters()} + + def fail_on_second_parameter(name, _model): + if name.endswith("bias"): + raise RuntimeError("stop during scoring setup") + return True + + session = _AutoQuantizeGradientScoringSession( + model, + [score_module], + fail_on_second_parameter, + ) + with pytest.raises(RuntimeError, match="stop during scoring setup"), session: + pytest.fail("scoring setup should not complete") + + assert "forward" not in score_module.__dict__ + assert { + name: param.requires_grad for name, param in model.named_parameters() + } == original_requires_grad + + +@pytest.mark.parametrize("instance_override", [False, True]) +def test_gradient_scoring_restores_forward_attribute_layout(instance_override): + module = torch.nn.Identity() + module._hparams_for_scoring = [] + + def original_forward(x): + return x + 1 + + original_override = original_forward + if instance_override: + module.forward = original_override + + session = _AutoQuantizeGradientScoringSession( + module, + [module], + lambda _name, _model: False, + ) + with pytest.raises(RuntimeError, match="stop during scoring"), session: + assert module.__dict__["forward"] is not original_override + raise RuntimeError("stop during scoring") + + if instance_override: + assert module.__dict__["forward"] is original_override + else: + assert "forward" not in module.__dict__ + + def test_auto_quantize_budget_uses_no_quant_candidate_cost(monkeypatch): class _BudgetCaptureSearcher(AutoQuantizeGradientSearcher): def run_search_with_stats(self, max_weight_size, verbose=False): From 60638ec0bd3422298560bf8a35c2d1aa16afc71c Mon Sep 17 00:00:00 2001 From: Chenjie Luo <108829653+cjluo-nv@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:56:23 -0700 Subject: [PATCH 14/28] Document the nvfp4_act_headroom calibration variant in ptq.md (#2439) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: documentation `general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml` appears in the shipped-recipes table in `modelopt_recipes/ptq.md`, but the **Calibration variants** section — which documents `max`, `mse`, `input_scale1`, `gptq`, and the `layerwise` variants — had no entry for it. Someone scanning that section for "which calibration do I pick when NVFP4 W4A4 regresses?" only found `mse`, which searches **weight** scales and so cannot help when the loss comes from activation clipping. This adds the missing entry: the scale formula (`amax = max(rho * anchor, upper)`) and its defaults, the fact that it costs one calibration pass and exports a standard NVFP4 checkpoint with coverage identical to `nvfp4_default-kv_fp8_cast`, and the symptoms that should route you here rather than to a weight-side calibration — an A16 ablation clears the regression while `mse` does not, the symptom is behavioral (verbose or runaway generations, hitting the generation cap) rather than a flat score drop, inference contexts run longer than the calibration set, or a few rare blocks dominate the activation error. MoE experts-only scopes are called out as the common case. It also extends step 3 of **Choosing a general recipe** so the escalation path reads `mse` first, then `nvfp4_act_headroom` when the evidence points at activations rather than weights. **Evidence.** The guidance comes from a GLM-5.3-Flash NVFP4 experts-only W4A4 root-cause study on SciCode (temperature 1.0), which established causally that activation quantization at the routed-expert `down_proj` input drove a large generation-length blow-up. Swapping `max` for `nvfp4_act_headroom` cut the median generation-length regression versus source from +38% to +19% and the mean from +19% to +4%, with no capped generations. The entry states plainly that this was the best strict-W4A4 result in that study but still missed the p50/p75 near-lossless gate, so headroom is presented as a strong first lever for activation-driven regressions rather than a guaranteed fix, with a note that `rho` should be swept. ### Usage No API or recipe change; the recipe already ships. This PR only documents when to select it over plain `max`: ```python from modelopt.recipe import load_recipe cfg = load_recipe("general/ptq/nvfp4_act_headroom-kv_fp8_cast") ``` ### Testing Docs-only change; no code paths touched, so no new or updated tests. - `pre-commit run --files modelopt_recipes/ptq.md` — all applicable hooks pass, including `markdownlint-cli2` and `check-modelopt-recipes`. - Re-read the rendered section to confirm the new bullet nests correctly in the existing `Calibration variants` list and that surrounding entries are unchanged. - Cross-checked every claim against the implementation (`modelopt/torch/quantization/calib/nvfp4_act_headroom.py`), the recipe YAML, and the existing `CHANGELOG.rst` entry, so the documented defaults (`anchor_percentile=1`, `upper_percentile=99.99`, `rho=16384`) and the NVFP4-input-quantizer-only scope match the code. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - 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?: N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - Did you get Claude approval on this PR?: ❌ ### Additional Information The entry deliberately does not sell this on accuracy: in that study the quantized subtask accuracy (56.80%) was *above* the source checkpoint (51.18%), so what headroom recovered was generation-length behavior, not accuracy. 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **New Features** * Added an NVFP4 activation headroom calibration option with configurable percentile-based scaling. * Supports standard NVFP4 checkpoint export with a single calibration pass. * Applies to dynamic-block NVFP4 activation quantizers while keeping weight-scale configuration independent. * Added guidance for addressing activation-related W4A4 accuracy regressions, including calibration coverage and recipe selection. Signed-off-by: Chenjie Luo Co-authored-by: Claude Opus 5 (1M context) Signed-off-by: Chad Voegele --- modelopt_recipes/ptq.md | 28 +++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 8cb2b718bbc..b12c3404d67 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -170,6 +170,31 @@ How the quantization scales are searched. The default (no suffix) is `max`. (amax) calibrated as in the default recipes. Costs more calibration time but recovers accuracy NVFP4 W4A4 can lose under plain max. Reach for it when a `max` recipe regresses. +- **`nvfp4_act_headroom`** (`nvfp4_act_headroom-kv_fp8_cast`) — leaves headroom + on the NVFP4 **activation** global scale. Plain `max` anchors that scale to the + largest per-block amax seen during calibration, so any larger activation at + inference saturates. This variant instead sets `amax = max(rho * anchor, upper)` + from the per-block amaxes at `anchor_percentile` (default 1) and + `upper_percentile` (default 99.99), with headroom factor `rho` (default 16384): + calibrated blocks sit low in the FP8 block-scale range, leaving the rest for + unseen outliers. Costs one calibration pass and exports a standard NVFP4 + checkpoint; coverage matches `nvfp4_default-kv_fp8_cast`, and the nested + `weight_scale_algorithm` still accepts `mse` / `local_hessian`. Affects NVFP4 + **input** quantizers only — a no-op for FP8 and weight-only recipes. + + Reach for it when a W4A4 recipe regresses and the **activations**, not the + weights, are to blame: an A16 ablation of the same scope clears it while `mse` + does not; the symptom is behavioral (verbose or runaway generations, hitting + the generation cap) rather than a flat score drop; inference contexts run + longer than the calibration set; or a few rare blocks dominate the activation + error. MoE experts-only scopes are the common case — one observer covers every + expert in a layer, and `down_proj` inputs clip first. On a GLM-5.3-Flash + experts-only W4A4 study (SciCode, temperature 1.0) it cut the median + generation-length regression from +38% to +19% and the mean from +19% to +4% + with no capped generations: the best strict-W4A4 result there, but still short + of the p50/p75 gate. A strong first lever, not a guaranteed fix — and sweep + `rho`, since headroom above the calibrated range costs resolution inside it. + - **`input_scale1`** (`nvfp4_experts_only_input_scale1-kv_fp8_cast`) — pins the expert **activation** per-tensor amax to a constant `2688.0` (= E2M1_MAX × E4M3_MAX = 6 × 448) via `constant_amax`, so the exported NVFP4 @@ -215,7 +240,8 @@ These can also be **stacked** when a single method isn't enough — e.g. `mse` + target requires, checking accuracy as you go. 3. **Recover accuracy via calibration before backing off the scope.** If a wider-scope recipe regresses, switch its `max` to the `mse` variant before - retreating to a narrower scope. + retreating to a narrower scope. If `mse` doesn't clear it but an A16 ablation + does, the **activations** are the problem: try `nvfp4_act_headroom` next. 4. **Pick KV by deployment.** `kv_fp8_cast` is the safe default (usually as accurate as calibrated `kv_fp8`); use `kv_nvfp4_cast` for maximum KV compression. From 9880ff96d36ce8553728977b9dc9942e41393283 Mon Sep 17 00:00:00 2001 From: Chad Voegele Date: Wed, 16 Sep 2026 13:43:37 -0500 Subject: [PATCH 15/28] docs: replace legacy TensorRT-LLM engine deployment guidance (#2436) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Chad's Agent Type of change: documentation. Replace legacy TensorRT checkpoint export, support matrix, and engine-build instructions with `export_hf_checkpoint` and TensorRT-LLM's PyTorch backend. Preserve the existing 0.48.0 deprecation / 0.49.0 removal notice and page URL. Update the customized-model guide and deployment skill to match. Follow the linked unified HF export guide. No API changes. - `git diff --check` passed. - `uvx pre-commit run --files docs/source/deployment/1_tensorrt_llm.rst docs/source/guides/_customized_model_quantization.rst plugins/modelopt/skills/deployment/references/trtllm.md` passed. - Full Sphinx build delegated to the Docs workflow; preview expected after deployment. - Is this change backward compatible?: ✅ Documentation only; page URL retained. - 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?: N/A — documentation only. - Did you update Changelog?: N/A — guidance correction; no new API deprecation. - Did you get Claude approval on this PR?: ❌ Not requested yet. Removes instructions for the TensorRT backend that current TensorRT-LLM releases no longer support. - **Documentation** - Updated TensorRT-LLM deployment guidance to use `export_hf_checkpoint` with the PyTorch backend. - Clarified that this workflow does not require TensorRT engine construction. - Updated DBRX customization instructions for exporting and deploying quantized models. - Added TensorRT-LLM version requirements and links to unified Hugging Face deployment guidance. - Removed guidance for the legacy TensorRT-LLM checkpoint exporter and outdated troubleshooting steps. --------- Signed-off-by: Chad Voegele --- docs/source/deployment/1_tensorrt_llm.rst | 159 ++---------------- .../guides/_customized_model_quantization.rst | 2 +- .../skills/deployment/references/trtllm.md | 6 +- 3 files changed, 16 insertions(+), 151 deletions(-) diff --git a/docs/source/deployment/1_tensorrt_llm.rst b/docs/source/deployment/1_tensorrt_llm.rst index 43195e3f627..77e4dab649c 100644 --- a/docs/source/deployment/1_tensorrt_llm.rst +++ b/docs/source/deployment/1_tensorrt_llm.rst @@ -2,149 +2,16 @@ TensorRT-LLM ========================== -**Deprecation Notice**: The export_tensorrt_llm_checkpoint API will be deprecated in future releases. Users are encouraged to transition to the :doc:`unified HF export API <3_unified_hf>`, which provides enhanced functionality and flexibility for exporting models to multiple inference frameworks including TensorRT-LLM, vLLM, and SGLang. - -.. note:: - - Please read the `TensorRT-LLM checkpoint workflow `_ - first before going through this section. - - - -ModelOpt toolkit supports automatic conversion of ModelOpt exported LLM to the TensorRT-LLM checkpoint and the engines for accelerated inferencing. - -This conversion is achieved by: - -#. Converting Huggingface, Megatron-Bridge and ModelOpt exported checkpoints to the TensorRT-LLM checkpoint. -#. Building TensorRT-LLM engine from the TensorRT-LLM checkpoint. - - -Export Quantized Model -====================== - -After the model is quantized, the quantized model can be exported to the TensorRT-LLM checkpoint format stored as - -#. A single JSON file recording the model structure and metadata (config.json) -#. A group of safetensors files, each recording the local calibrated model on a single GPU rank (model weights, scaling factors per GPU). - -The export API (:meth:`export_tensorrt_llm_checkpoint `) can be used as follows: - -.. code-block:: python - - from modelopt.torch.export import export_tensorrt_llm_checkpoint - - with torch.inference_mode(): - export_tensorrt_llm_checkpoint( - model, # The quantized model. - decoder_type, # The type of the model as str, e.g gpt, gptj, llama. - dtype, # the weights data type to export the unquantized layers. - export_dir, # The directory where the exported files will be stored. - inference_tensor_parallel, # The number of GPUs used in the inference time for tensor parallelism. - inference_pipeline_parallel, # The number of GPUs used in the inference time for pipeline parallelism. - ) - -If the :meth:`export_tensorrt_llm_checkpoint ` call is successful, the TensorRT-LLM checkpoint will be saved. Otherwise, e.g. the ``decoder_type`` is not supported, a torch state_dict checkpoint will be saved instead. - -.. list-table:: Model support matrix for the TensorRT-LLM checkpoint export - :header-rows: 1 - - * - Model / Quantization - - FP16 / BF16 - - FP8 - - INT8_SQ - - INT4_AWQ - * - GPT2 - - Yes - - Yes - - Yes - - No - * - GPTJ - - Yes - - Yes - - Yes - - Yes - * - LLAMA 2 - - Yes - - Yes - - Yes - - Yes - * - LLAMA 3 - - Yes - - Yes - - No - - Yes - * - Mistral - - Yes - - Yes - - Yes - - Yes - * - Mixtral 8x7B - - Yes - - Yes - - No - - Yes - * - Falcon 40B, 180B - - Yes - - Yes - - Yes - - Yes - * - Falcon 7B - - Yes - - Yes - - Yes - - No - * - MPT 7B, 30B - - Yes - - Yes - - Yes - - Yes - * - Baichuan 1, 2 - - Yes - - Yes - - Yes - - Yes - * - ChatGLM2, 3 6B - - Yes - - No - - No - - Yes - * - Bloom - - Yes - - Yes - - Yes - - Yes - * - Phi-1, 2, 3 - - Yes - - Yes - - Yes - - Yes - * - Nemotron 8 - - Yes - - Yes - - No - - Yes - * - Gemma 2B, 7B - - Yes - - Yes - - No - - Yes - * - Recurrent Gemma - - Yes - - Yes - - Yes - - Yes - * - StarCoder 2 - - Yes - - Yes - - Yes - - Yes - * - Qwen-1, 1.5 - - Yes - - Yes - - Yes - - Yes - -Convert to TensorRT-LLM -======================= - -Once the TensorRT-LLM checkpoint is available, please follow the `TensorRT-LLM build API `_ to build and deploy the quantized LLM. +For current TensorRT-LLM deployments, export quantized models with +:meth:`export_hf_checkpoint ` +and load the exported Hugging Face checkpoint with TensorRT-LLM's PyTorch backend. +See the :doc:`unified HF export guide <3_unified_hf>` for export and deployment +examples, supported models, and quantization formats. This workflow does not require +building a TensorRT engine. + +.. warning:: + + The ``export_tensorrt_llm_checkpoint`` API exports checkpoints for the legacy + TensorRT backend, which current TensorRT-LLM releases no longer support. + The API will be deprecated in a future release. + Use ``export_hf_checkpoint`` instead. diff --git a/docs/source/guides/_customized_model_quantization.rst b/docs/source/guides/_customized_model_quantization.rst index c8078678ec4..b2fa879ecd4 100644 --- a/docs/source/guides/_customized_model_quantization.rst +++ b/docs/source/guides/_customized_model_quantization.rst @@ -16,7 +16,7 @@ As ModelOpt cannot detect these linear ops out-of-the-box, a HugggingFace plugin #. Rewrite the linear ops (w1, v1 and v2) as a standard ``nn.Linear`` op, and re-implement the ``forward`` method. #. Register the new dynamic ``_QuantDbrxExperts`` to replace the ``DbrxExperts`` from the modeling_dbrx.py in the ``transformers`` library #. Try quantize the DBRX model after the plugin is implemented, feel free to follow the `hf_ptq example `_. -#. TensorRT-LLM is open-sourced. If this customized model is not supported by TensorRT-LLM yet, please modify :meth:`export_tensorrt_llm_checkpoint ` or :meth:`export_hf_checkpoint ` to export the quantized model for deployment with a customized TensorRT-LLM modeling implementation. Feel free to :doc:`contact us <../support/1_contact>` if further support is needed. +#. Export the quantized model with :meth:`export_hf_checkpoint `. If the customized model is not supported by TensorRT-LLM, add support in its PyTorch backend and adapt the HF exporter if needed. See the :doc:`unified HF export guide <../deployment/3_unified_hf>` or :doc:`contact us <../support/1_contact>` for help. The following code snippet is excerpted from ``modelopt/torch/quantization/plugins/huggingface.py`` diff --git a/plugins/modelopt/skills/deployment/references/trtllm.md b/plugins/modelopt/skills/deployment/references/trtllm.md index 27185793307..a24835cdcf8 100644 --- a/plugins/modelopt/skills/deployment/references/trtllm.md +++ b/plugins/modelopt/skills/deployment/references/trtllm.md @@ -2,7 +2,7 @@ ## Requirements -- TensorRT-LLM >= 0.17.0 +- TensorRT-LLM >= 1.2.0 (see `docs/source/deployment/3_unified_hf.rst`) - Typically installed via NVIDIA container: `nvcr.io/nvidia/tensorrt-llm/release:` - Or: `pip install tensorrt-llm` @@ -65,7 +65,7 @@ directly together with a ModelOpt-quantized checkpoint. The legacy export path using `export_tensorrt_llm_checkpoint()` is deprecated. Use the unified HF checkpoint format with `export_hf_checkpoint()` instead. -If you encounter a legacy checkpoint (no `hf_quant_config.json`, has `rank*.safetensors` pattern), it needs the TRT-LLM build API to create an engine before deployment. See `docs/source/deployment/1_tensorrt_llm.rst`. +Current TensorRT-LLM releases no longer support the legacy TensorRT backend. Re-export the quantized source model with `export_hf_checkpoint()` and deploy using the PyTorch backend. See `docs/source/deployment/3_unified_hf.rst`. ## Evaluation with TRT-LLM @@ -91,5 +91,3 @@ explicitly for few-shot tasks. |-------|-----| | `No module named tensorrt_llm` | Install via container or pip | | NVFP4 inference fails on Hopper | NVFP4 requires Blackwell GPUs for inference | -| Slow first inference | Engine compilation happens on first run; subsequent runs are cached | -| OOM during engine build | Reduce `--max_batch_size` or increase TP | From 6ea33c5eb3b5c67217316d32a8109f1eead9f333 Mon Sep 17 00:00:00 2001 From: Wei-Ming Chen <17592131+meenchen@users.noreply.github.com> Date: Thu, 17 Sep 2026 10:52:46 -0700 Subject: [PATCH 16/28] Fail fast on non-finite AutoQuantize output gradients (#2432) Type of change: Bug fix Fail fast when AutoQuantize receives non-finite output gradients, before accumulating sensitivity scores. The error names the affected module and suggests checking the model, data, and loss; it also identifies cuDNN SDPA backward on fully masked rows as one possible cause and gives an explicit retry workaround. Unlike the earlier revision, this does not disable cuDNN or change any attention backend settings. Invalid gradients are not zeroed or ignored, and there is no automatic retry. No API or recipe changes. For the reproduced cuDNN failure, the caller can explicitly set `torch.backends.cuda.enable_cudnn_sdp(False)` before a fresh AutoQuantize run. - AutoQuantize unit suite: **110 passed**. Coverage includes NaN and positive/negative infinity, module diagnostics, preventing invalid score accumulation, model-state cleanup, and unchanged SDPA backend settings. - Real-model E2E on **four GB300 GPUs**, Qwen/Qwen3.6-35B-A3B, main `8025a3dc5481129aa21fef99cb13a879e1b5847e` plus this patch, batch size 8, 512 calibration samples, and `w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml`: - Default backend: the new diagnostic fired at `model.language_model.layers.39.self_attn.q_proj`; cuDNN remained enabled and no quantized model was exported. Expected-error check passed. - Explicit cuDNN-disabled fresh run: both 64-batch calibration passes, all 16 scoring batches, optimization at **5.99 effective bits**, and checkpoint export completed with exit code 0. Verified all three indexed safetensors shards and quantization configuration; the index contains 93,563 tensors. - All applicable pre-commit checks and `git diff --check` passed. Contributor guidelines and security guidance followed; commits are signed and signed off. - Is this change backward compatible?: Yes; finite-gradient behavior and backend settings are unchanged. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A; neither added. - Did you write any new necessary tests?: Yes. - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: Yes, 0.48.0 bug fixes. - Did you get Claude approval on this PR?: No; awaiting review. This improves error reporting rather than fixing the upstream cuDNN kernel. Blackwell-specificity is not established. Checkpoint deployment/reload was not tested. A calibration-only checkpoint-resume attempt completed scoring but encountered a separate `candidate_stats` KeyError. That issue is outside this patch; the successful export validation above used a fresh run without search-state resume. * **Bug Fixes** * AutoQuantize now fails fast when output gradients contain non-finite values, with an actionable error identifying the affected module. * Attention backend settings are preserved and restored after successful runs and failures. * Model state is restored when sensitivity scoring encounters an error during setup or execution. --------- Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com> Signed-off-by: Chad Voegele --- CHANGELOG.rst | 1 + modelopt/torch/quantization/algorithms.py | 10 +++ .../unit/torch/quantization/test_autoquant.py | 70 +++++++++++++++++++ 3 files changed, 81 insertions(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 1dd7936723f..5af8f8588b9 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -56,6 +56,7 @@ Changelog **Bug Fixes** +- Fail fast on non-finite AutoQuantize output gradients with an actionable error before accumulating sensitivity scores, without changing attention backend settings. - Fix ONNX INT8 entropy calibration failing or producing invalid quantization parameters for FP16 activations. - Speed up ``mtq.quantize`` on FSDP2-sharded fused-MoE models. Promoting static-block weight quantizers gathered each expert's slice of the fused weight across ranks even though only quantizer state is read, adding a collective per expert to calibration. - Fix ONNX AutoCast failing on models with external initializers larger than 2 GiB. diff --git a/modelopt/torch/quantization/algorithms.py b/modelopt/torch/quantization/algorithms.py index 133f278d24a..889290231e3 100644 --- a/modelopt/torch/quantization/algorithms.py +++ b/modelopt/torch/quantization/algorithms.py @@ -1604,6 +1604,16 @@ def forward(self, module: nn.Module, *args, **kwargs): def _accumulate_scores(self, module, invocation_diffs, grad_output) -> None: """Accumulate scores for the invocation that produced ``grad_output``.""" + if not torch.isfinite(grad_output).all(): + module_name = next( + name for name, child in self.model.named_modules() if child is module + ) + raise RuntimeError( + f"AutoQuantize: Non-finite output gradients in module '{module_name or ''}'. " + "Cannot compute reliable sensitivity scores. Check the model, data, and loss. " + "cuDNN SDPA backward on fully masked attention rows is one possible cause; " + "try torch.backends.cuda.enable_cudnn_sdp(False) before rerunning auto_quantize." + ) for hparam, output_diffs in invocation_diffs.items(): for recipe, output_diff in output_diffs.items(): importance = hparam._importance_dict[recipe][module] diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py index 0013f2da538..58ba4c4f627 100644 --- a/tests/unit/torch/quantization/test_autoquant.py +++ b/tests/unit/torch/quantization/test_autoquant.py @@ -16,6 +16,7 @@ import copy import io import warnings +from contextlib import nullcontext from types import SimpleNamespace import pytest @@ -1129,6 +1130,75 @@ def fail_during_scoring(model, data): assert hparam.active == hparam.original +@pytest.mark.parametrize("cudnn_enabled", [False, True]) +@pytest.mark.parametrize("fail", [False, True]) +def test_backward_scoring_session_preserves_sdpa_backends(cudnn_enabled, fail): + original = torch.backends.cuda.cudnn_sdp_enabled() + others = ( + torch.backends.cuda.flash_sdp_enabled(), + torch.backends.cuda.mem_efficient_sdp_enabled(), + torch.backends.cuda.math_sdp_enabled(), + ) + try: + torch.backends.cuda.enable_cudnn_sdp(cudnn_enabled) + with ( + pytest.raises(RuntimeError, match="scoring failed") if fail else nullcontext(), + _AutoQuantizeGradientScoringSession(torch.nn.Identity(), [], lambda *_: False), + ): + assert torch.backends.cuda.cudnn_sdp_enabled() == cudnn_enabled + assert others == ( + torch.backends.cuda.flash_sdp_enabled(), + torch.backends.cuda.mem_efficient_sdp_enabled(), + torch.backends.cuda.math_sdp_enabled(), + ) + if fail: + raise RuntimeError("scoring failed") + finally: + restored = torch.backends.cuda.cudnn_sdp_enabled() + torch.backends.cuda.enable_cudnn_sdp(original) + assert restored == cudnn_enabled + + +@pytest.mark.parametrize("bad_gradient", [float("nan"), float("inf"), -float("inf")]) +def test_auto_quantize_fails_fast_on_nonfinite_gradients(bad_gradient): + model = SimpleLinear() + original_cudnn = torch.backends.cuda.cudnn_sdp_enabled() + + def loss_func(output, _data): + output.register_hook(lambda grad: torch.full_like(grad, bad_gradient)) + return output.sum() + + with pytest.raises(RuntimeError, match="Non-finite output gradients in module") as exc_info: + mtq.auto_quantize( + model, + constraints={"effective_bits": 12.0}, + quantization_formats=[mtq.INT8_DEFAULT_CFG], + data_loader=[model.get_input()], + forward_step=lambda model, batch: model(batch), + loss_func=loss_func, + num_calib_steps=1, + num_score_steps=1, + ) + + assert "torch.backends.cuda.enable_cudnn_sdp(False)" in str(exc_info.value) + assert torch.backends.cuda.cudnn_sdp_enabled() == original_cudnn + assert any( + f"module '{name}'" in str(exc_info.value) + for name, module in model.named_modules() + if getattr(module, "_hparams_for_scoring", []) + ) + for name, module in model.named_modules(): + for hparam in getattr(module, "_hparams_for_scoring", []): + assert "forward" not in module.__dict__ + assert hparam.active == hparam.original + for recipe in hparam.choices: + importance = hparam._importance_dict[recipe][module] + if f"module '{name}'" in str(exc_info.value): + assert importance is None + else: + assert importance is None or torch.isfinite(importance).all() + + def test_backward_scoring_session_restores_partial_setup(): model = torch.nn.Sequential(torch.nn.Linear(4, 4)) score_module = model[0] From 763cf956badecba695fc14f4be44f0f38e406ba5 Mon Sep 17 00:00:00 2001 From: yeyu-nvidia Date: Thu, 17 Sep 2026 10:54:30 -0700 Subject: [PATCH 17/28] fix(specdec): resolve the eagle aux-layer preset in the vLLM hidden-state dump (#2410) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix Fixes `nvbugs/6753684`, filed against #2080 by the ModelOpt QA Sentinel. The vLLM offline hidden-state dump rejected `--aux-layers eagle` — **the flag's own default** — so the documented invocation aborted before writing any state: ``` File "collect_hidden_states/compute_hidden_states_vllm.py", line 76, in _resolve_aux_layers_standalone ids = sorted({int(t) for t in aux_layers.split(',') if t.strip()}) ValueError: invalid literal for int() with base 10: 'eagle' ``` **Root cause.** `compute_hidden_states_vllm.py` runs in a stock vLLM container, where importing `modelopt.torch` fails (the full init chain pulls in omegaconf and friends). It therefore carries `_resolve_aux_layers_standalone`, a local copy of the preset logic in `common.resolve_aux_layers`. That copy implemented the `dflash` preset and explicit id lists, but never `eagle` — while `add_aux_layers_args` defaults to `eagle`. The HF and TRT-LLM dumps call the shared helper and were unaffected; only the vLLM path forked, and nothing compared the fork against its source. This PR resolves `eagle` inline, mirroring `hf_eagle.default_eagle_aux_layer_ids`. It also fixes a second defect the bug exposes: the function already had a message naming the accepted values, but it was unreachable, because `int()` raised first. An unrecognised preset now reports what it accepts instead of surfacing the raw `int()` error — which is what made the original failure opaque. ### Usage The previously-broken documented invocation now works: ```bash cd examples/speculative_decoding python collect_hidden_states/compute_hidden_states_vllm.py \ --model Qwen/Qwen2.5-0.5B-Instruct \ --input-data ../dataset/synthetic_conversations_1k.jsonl \ --output-dir /tmp/hs_vllm \ --max-seq-len 512 --tp 1 ``` `--aux-layers dflash` and explicit lists such as `--aux-layers 2,5,8` are unchanged. ### Testing Added `tests/unit/examples/test_vllm_hidden_states_aux_layers.py`, which pins the standalone copy to the shared implementation it mirrors: - `eagle` matches `hf_eagle.default_eagle_aux_layer_ids` across layer counts 4, 6, 8, 12, 24, 28, 32, 36, 48, 52, 61, 80 — deliberately including counts small enough that the `max(0, ...)` clamps collapse ids together. - A named regression case for `nvbugs/6753684`. - `dflash` and explicit-list behaviour unchanged. - Unknown specs (`bogus`, `EAGLE3`, `eagle3`, empty) raise the actionable message. - Out-of-range ids still rejected. Divergence here is silent — the dump would write plausible-looking hidden states from the *wrong* layers, surfacing much later as a poor acceptance rate. Hence pinning to the reference rather than asserting hardcoded lists alone. All 20 assertions verified and every pre-commit hook passes (`ruff`, `mypy`, `bandit`, RST lint, license headers). One caveat worth stating plainly: **pytest could not be run locally.** `tests/unit/conftest.py` imports `modelopt.torch.utils.distributed`, which needs `CPUOffloadPolicy` from `torch.distributed.fsdp` — absent in this machine's torch. Each assertion was executed directly against the real module instead, but CI is the first genuine pytest run. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ — strictly widens accepted input; `dflash` and explicit lists behave identically. - 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)?: ✅ — bug fix for a defect present in a previous release. - Did you get Claude approval on this PR?: ❌ — not yet run. ### Additional Information The underlying fragility is the duplicated implementation, not this one missing branch. The function's own `TODO: drop this once common.resolve_aux_layers is decoupled from the heavy modelopt.torch import chain` is the real fix; the new test narrows the gap but does not close it. Worth tracking separately if the vLLM dump is expected to keep pace with new presets. 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **Bug Fixes** * Fixed `--aux-layers eagle` for vLLM offline hidden-state collection. * Added support for the documented `eagle` preset alongside `dflash` and explicit layer IDs. * Improved invalid-option errors to clearly list accepted formats. * Rejects `dflash` configurations when the target model has too few layers. * Continues rejecting layer IDs outside the model’s available range. * **Documentation** * Added a v0.48.0 changelog entry for the fix. --------- Signed-off-by: Ye Yu Co-authored-by: Claude Opus 5 Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Signed-off-by: Chad Voegele --- CHANGELOG.rst | 1 + .../compute_hidden_states_vllm.py | 38 ++++-- .../test_vllm_hidden_states_aux_layers.py | 110 ++++++++++++++++++ 3 files changed, 141 insertions(+), 8 deletions(-) create mode 100644 tests/examples/speculative_decoding/test_vllm_hidden_states_aux_layers.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 5af8f8588b9..bcc4f7ef9ed 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -79,6 +79,7 @@ Changelog - Fix EAGLE-3 training with context parallelism (``--cp_size > 1`` in ``examples/speculative_decoding``), which failed to start on ``accelerate >= 1.13`` and then raised ``got mixed torch.Tensor and DTensor``. - Polygraphy minimum dependency upgraded to ``0.53.4`` to solve ONNX AutoCast failures when marking optional graph outputs. - Fix ``--kv_cache_free_gpu_memory_fraction`` having no effect on the ``lm_eval`` task of ``examples/hf_ptq/scripts/huggingface_example.sh``, where the KV cache always took TensorRT-LLM's default 90% of free GPU memory and evaluation could run out of memory. ``examples/llm_eval/lm_eval_trtllm.py`` now takes ``kv_cache_free_gpu_memory_fraction`` in ``--model_args``, defaulting to 0.8. +- Fix ``--aux-layers eagle`` failing in the vLLM offline hidden-state dump (``examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py``). ``eagle`` is the flag's default, but the dump's standalone resolver -- a copy kept so the script runs in a stock vLLM container without ModelOpt -- only handled ``dflash`` and explicit id lists, so the documented invocation aborted with ``invalid literal for int(): 'eagle'`` before any state was written. An unrecognised preset now reports which values are accepted instead of surfacing the raw ``int()`` error. 0.46 (2026-08-17) ^^^^^^^^^^^^^^^^^ diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py index 77441f8f858..c232852a4c8 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py @@ -55,9 +55,10 @@ def _resolve_aux_layers_standalone( This dump runs in a stock vLLM container. ``common.resolve_aux_layers`` resolves the 'dflash'/'eagle' presets by importing ``modelopt.torch.speculative.plugins`` — which pulls in the full ``modelopt.torch`` init chain (omegaconf, etc.) that the vLLM - container does not have, so the import fails. Resolve the 'dflash' preset inline - (mirroring ``modeling_dflash.build_target_layer_ids`` for ``num_draft`` draft layers) - and accept an explicit comma-separated int list. ``num_draft`` MUST match the recipe's + container does not have, so the import fails. Resolve both presets inline -- 'eagle' + mirrors ``hf_eagle.default_eagle_aux_layer_ids``, 'dflash' mirrors + ``modeling_dflash.build_target_layer_ids`` for ``num_draft`` draft layers -- and + accept an explicit comma-separated int list. ``num_draft`` MUST match the recipe's ``dflash.dflash_architecture_config.num_hidden_layers`` (pass --num-draft-layers) or the dumped aux layers silently mis-align with what the draft consumes at training time. Keep in sync with modelopt. @@ -65,15 +66,39 @@ def _resolve_aux_layers_standalone( TODO: drop this once ``common.resolve_aux_layers`` is decoupled from the heavy ``modelopt.torch`` import chain so it can be reused directly in a vLLM container. """ + # One literal: a future preset means editing this message once, not twice. The bug + # this guards against was precisely a copy drifting from its original. + unsupported = ( + f"--aux-layers={aux_layers!r}: in the stock vLLM container (no modelopt) only the " + "'eagle' / 'dflash' presets or an explicit comma-separated layer-id list " + "are supported." + ) spec = aux_layers.strip().lower() + if spec == "eagle": + # Mirrors hf_eagle.default_eagle_aux_layer_ids: three layers near the start, + # middle, and end of the stack. This is add_aux_layers_args' default, so the + # dump must resolve it without modelopt or the documented invocation fails. + return sorted({1, max(0, num_hidden_layers // 2 - 1), max(0, num_hidden_layers - 4)}) if spec == "dflash": + # build_target_layer_ids raises here; without the same guard the interpolation + # below silently collapses to a short deduplicated list (e.g. [1] for a 4-layer + # target), dumping fewer aux layers than the draft consumes. + if num_hidden_layers < num_draft: + raise ValueError( + f"num_target_layers ({num_hidden_layers}) must be >= num_draft_layers ({num_draft})" + ) if num_draft == 1: return [num_hidden_layers // 2] start = min(1, num_hidden_layers - 1) end = max(start, num_hidden_layers - 3) span = end - start return sorted({round(start + (i * span) / (num_draft - 1)) for i in range(num_draft)}) - ids = sorted({int(t) for t in aux_layers.split(",") if t.strip()}) + try: + ids = sorted({int(t) for t in aux_layers.split(",") if t.strip()}) + except ValueError as exc: + # An unrecognised preset name would otherwise surface as a bare + # "invalid literal for int()", which hides what the caller should pass. + raise ValueError(unsupported) from exc # Match the shared helper's contract: ids must be valid layer indices. out_of_range = [i for i in ids if not 0 <= i < num_hidden_layers] if out_of_range: @@ -82,10 +107,7 @@ def _resolve_aux_layers_standalone( f"for a {num_hidden_layers}-layer model." ) if not ids: - raise ValueError( - f"--aux-layers={aux_layers!r}: in the stock vLLM container (no modelopt) only the " - "'dflash' preset or an explicit comma-separated layer-id list are supported." - ) + raise ValueError(unsupported) return ids diff --git a/tests/examples/speculative_decoding/test_vllm_hidden_states_aux_layers.py b/tests/examples/speculative_decoding/test_vllm_hidden_states_aux_layers.py new file mode 100644 index 00000000000..52f3c2b96f2 --- /dev/null +++ b/tests/examples/speculative_decoding/test_vllm_hidden_states_aux_layers.py @@ -0,0 +1,110 @@ +# 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. + +"""The vLLM dump resolves aux layers without modelopt -- so it can drift from it. + +``compute_hidden_states_vllm.py`` runs in a stock vLLM container where importing +``modelopt.torch`` fails, so it carries its own copy of the preset logic. A copy that +nothing compares against is a copy that silently diverges: the 'eagle' preset -- which +is ``add_aux_layers_args``' *default* -- was missing from it entirely, so the documented +invocation died on ``int('eagle')`` before the dump began. + +These tests pin the standalone resolver to the real ``hf_eagle`` helper. Unlike the +dump script, the unit suite *can* import modelopt, so the comparison is against the +actual source of truth rather than a third transcription of the formula -- a drift +guard written against its own copy would guard nothing. +""" + +import importlib.util +import sys + +import pytest +from _test_utils.examples.run_command import MODELOPT_ROOT + +from modelopt.torch.speculative.plugins.hf_eagle import default_eagle_aux_layer_ids +from modelopt.torch.speculative.plugins.modeling_dflash import build_target_layer_ids + +_COLLECT = MODELOPT_ROOT / "examples" / "speculative_decoding" / "collect_hidden_states" +sys.path.insert(0, str(_COLLECT)) + +_SPEC = importlib.util.spec_from_file_location( + "chs_vllm", _COLLECT / "compute_hidden_states_vllm.py" +) +assert _SPEC is not None and _SPEC.loader is not None +chs_vllm = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(chs_vllm) + +resolve = chs_vllm._resolve_aux_layers_standalone + +# Spans tiny models through 80-layer frontier stacks; includes counts small enough that +# the formula's max(0, ...) clamps collapse ids together. +LAYER_COUNTS = [4, 6, 8, 12, 24, 28, 32, 36, 48, 52, 61, 80] + + +@pytest.mark.parametrize("num_layers", LAYER_COUNTS) +def test_eagle_preset_matches_the_shared_helper(num_layers): + """The standalone copy must agree with modelopt's own EAGLE layer selection. + + Disagreement is silent: the dump writes plausible-looking hidden states from the + wrong layers, and only a poor acceptance rate much later reveals it. + """ + assert resolve("eagle", num_layers) == default_eagle_aux_layer_ids(num_layers) + + +def test_eagle_is_resolvable_because_it_is_the_documented_default(): + """Regression for nvbugs/6753684: `--aux-layers eagle` raised ValueError. + + ``add_aux_layers_args`` defaults to 'eagle', so omitting the flag entirely used to + crash the dump before it started. + """ + assert resolve("eagle", 24) == [1, 11, 20] + + +def test_dflash_preset_still_resolves(): + assert resolve("dflash", 32, num_draft=5) == [1, 8, 15, 22, 29] + + +@pytest.mark.parametrize("num_layers", range(1, 40)) +def test_dflash_preset_matches_the_shared_helper(num_layers): + """Same pinning as EAGLE, including where the shared helper *refuses*. + + A target with fewer layers than the draft has no valid assignment, so + ``build_target_layer_ids`` raises. Without the matching guard the standalone + copy silently deduplicated down to a short list -- e.g. ``[1]`` for a 4-layer + target -- and the dump would write fewer aux layers than the draft consumes. + """ + try: + expected = sorted(set(build_target_layer_ids(num_layers, 5))) + except ValueError: + with pytest.raises(ValueError, match="must be >="): + resolve("dflash", num_layers, num_draft=5) + else: + assert resolve("dflash", num_layers, num_draft=5) == expected + + +def test_explicit_id_list_still_resolves(): + assert resolve("2,5,8", 32) == [2, 5, 8] + + +@pytest.mark.parametrize("spec", ["bogus", "EAGLE3", "", "eagle3"]) +def test_unknown_spec_explains_what_is_accepted(spec): + """A bare ``invalid literal for int()`` hides what the caller should have passed.""" + with pytest.raises(ValueError, match="'eagle' / 'dflash' presets"): + resolve(spec, 32) + + +def test_out_of_range_ids_are_rejected(): + with pytest.raises(ValueError, match="out of range"): + resolve("2,99", 32) From 3d9872f9288b46a3316cee81fb7256f9f2e40f19 Mon Sep 17 00:00:00 2001 From: kinjalpatel27 <31936134+kinjalpatel27@users.noreply.github.com> Date: Thu, 17 Sep 2026 14:12:34 -0700 Subject: [PATCH 18/28] Fix vLLM fakequant calibration for hybrid attention models (#2414) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix Fix fakequant calibration for hybrid attention/Mamba models, including NVIDIA Nemotron-3-Nano, on vLLM 0.26 and 0.28. The manual calibration scheduler path previously submitted requests with empty KV-cache block tables. Hybrid models require scheduler-compatible cache state during prefill; on current vLLM releases the empty tables caused the Mamba state to use the reserved null block and calibration activations became NaN. Request cleanup also no longer matched the vLLM 0.28 execution lifecycle, which could leave request-scoped state in the persistent batch. This PR: - Allocates non-null scratch blocks for every KV-cache group using the vLLM warmup reservation policy. - Supports both the vLLM 0.28 reservation helper and the equivalent vLLM 0.26 calculation. - Passes newly allocated blocks through `new_block_ids_to_zero` when that scheduler field is available. - Validates that the calibration batch fits in the configured cache and reports how to reduce calibration demand if it does not. - Cleans up calibration requests through a zero-token scheduler step on current vLLM, with a direct cleanup fallback for older runners. - Updates the example Dockerfile to default to vLLM 0.28.0 while retaining vLLM 0.26.0 through `VLLM_VERSION`. - Documents the validated Nemotron-3-Nano NVFP4 KV-cache workflow and clarifies that reducing `--max-num-batched-tokens` is not required. ### Usage Build the default vLLM 0.28.0 image: ```bash docker build -f examples/vllm_serve/Dockerfile \ -t vllm-modelopt:v0.28.0 . ``` Build with vLLM 0.26.0: ```bash docker build --build-arg VLLM_VERSION=0.26.0 \ -f examples/vllm_serve/Dockerfile \ -t vllm-modelopt:v0.26.0 . ``` Calibrate and serve Nemotron-3-Nano with NVFP4 KV-cache fakequant: ```bash KV_QUANT_CFG=NVFP4_KV_CFG QUANT_CALIB_SIZE=512 \ python examples/vllm_serve/vllm_serve_fakequant.py \ \ --trust-remote-code --enforce-eager -tp 8 \ --max-model-len 8192 --host 0.0.0.0 --port 8000 ``` ### Testing Validated on omniml-a0 with `NVIDIA-Nemotron-3-Nano-30B-A3B-BF16`, tensor parallel size 8, `NVFP4_KV_CFG`, `QUANT_CALIB_SIZE=512`, and `--max-model-len 8192`. No `--max-num-batched-tokens` override was used. - vLLM 0.28.0: - All 512 calibration samples completed. - No NaNs or cache-cleanup warnings were observed. - The server started and `/health` passed. - An OpenAI-compatible completion request returned coherent generated text. - vLLM 0.26.0: - Repeated the same 512-sample TP8 calibration with the official `vllm/vllm-openai:v0.26.0` image. - No NaNs were observed. - The server started, passed `/health`, and returned coherent generated text. - Docker: - Built and verified the updated vLLM 0.28.0 image. - Focused tests: - `tests/examples/vllm_serve/test_vllm_mlflow_utils.py`: 32 passed. - Cleanup failure, missing legacy API, and legacy fallback tests: 5 passed on both vLLM 0.26.0 and 0.28.0. - Repository hooks: - Targeted pre-commit hooks for every changed Python, Markdown, and Docker file: passed. - `git diff --check`: passed. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - 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?: ✅ — added focused coverage for fail-closed cleanup, exception chaining, and the legacy cleanup fallback; the full regression was also validated end to end. - Did you update Changelog?: N/A - Did you get Claude approval on this PR?: N/A ### Additional Information The change is quantization-format agnostic. It corrects the calibration scheduler and cache lifecycle rather than special-casing `NVFP4_KV_CFG` or using an NVFP4 cast path. ## Summary by CodeRabbit - **New Features** - Added support for configuring the vLLM version through `VLLM_VERSION`, with vLLM 0.28.0 as the default. - Added calibration and serving guidance for hybrid attention/Mamba models, including Nemotron 3 Nano with NVFP4 KV-cache fake quantization. - **Bug Fixes** - Improved calibration block handling across supported vLLM versions. - Improved calibration cleanup to preserve original errors and provide reliable fallback behavior when standard cleanup is unavailable. - **Documentation** - Documented tested versions, direct installation commands, ModelOpt setup, serving options, and guidance to avoid NaNs during batched serving. --------- Signed-off-by: Kinjal Patel Signed-off-by: Chad Voegele --- examples/vllm_serve/Dockerfile | 3 +- examples/vllm_serve/README.md | 38 ++- examples/vllm_serve/vllm_ptq_utils.py | 182 ++++++++++--- .../quantization/test_vllm_dynamic_modules.py | 240 ++++++++++++++++++ 4 files changed, 428 insertions(+), 35 deletions(-) diff --git a/examples/vllm_serve/Dockerfile b/examples/vllm_serve/Dockerfile index 5fe0799c427..a4c26ab1a7b 100644 --- a/examples/vllm_serve/Dockerfile +++ b/examples/vllm_serve/Dockerfile @@ -1,4 +1,5 @@ -FROM vllm/vllm-openai:v0.26.0 +ARG VLLM_VERSION=0.28.0 +FROM vllm/vllm-openai:v${VLLM_VERSION} # Set environment variables ENV PIP_NO_CACHE_DIR=off \ diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index fc4e8a0ebcc..bc2b0fb53f3 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -4,17 +4,35 @@ This is a simple example to demonstrate calibrating and serving ModelOpt fakequa Compared with realquant, fakequant is 2-5x slower, but doesn't require dedicated kernel support and facilitates research. -The general fakequant example is tested with vLLM 0.9.0, 0.19.1, and 0.26.0. The compact -NVFP4 attention worker documented below requires vLLM 0.15.0 or newer. +The general fakequant example is tested with vLLM 0.9.0, 0.19.1, 0.26.0, and 0.28.0. The +compact NVFP4 attention worker documented below requires vLLM 0.15.0 or newer. ## Prepare environment -Follow the following instruction to build a docker environment, or install vllm with pip. +Use the Dockerfile to build an environment with vLLM 0.28.0: ```bash -docker build -f examples/vllm_serve/Dockerfile -t vllm-modelopt . +docker build -f examples/vllm_serve/Dockerfile -t vllm-modelopt:v0.28.0 . ``` +To build the same environment with another tested vLLM release, override `VLLM_VERSION`: + +```bash +docker build --build-arg VLLM_VERSION=0.26.0 \ + -f examples/vllm_serve/Dockerfile -t vllm-modelopt:v0.26.0 . +``` + +For a direct installation from the ModelOpt repository root, install the tested vLLM +release and the ModelOpt extras used by this example: + +```bash +python3 -m pip install "vllm==0.28.0" +python3 -m pip install -e ".[all,mlflow]" +``` + +See the [ModelOpt installation guide](../../docs/source/getting_started/_installation_for_Linux.rst) +for details about installing partial dependency sets. + ## Calibrate and serve fake quant model in vLLM Step 1: Configure quantization settings. @@ -39,6 +57,18 @@ Step 2: Run the following command, with all supported flag as `vllm serve`: python vllm_serve_fakequant.py -tp 8 --host 0.0.0.0 --port 8000 ``` +Hybrid attention/Mamba models such as Nemotron 3 Nano are supported on vLLM 0.26.0 and +0.28.0. For example, calibrate and serve with NVFP4 KV-cache fakequant as follows: + +```bash +KV_QUANT_CFG=NVFP4_KV_CFG QUANT_CALIB_SIZE=512 \ + python vllm_serve_fakequant.py -tp 8 \ + --max-model-len 8192 --enforce-eager --host 0.0.0.0 --port 8000 +``` + +Calibration uses dedicated scratch KV-cache blocks, so reducing `--max-num-batched-tokens` +is not required to avoid NaNs. + For vLLM versions that expose `--moe-backend`, this launcher defaults to `--moe-backend triton`. ModelOpt expert fakequant needs a decomposed MoE backend so both expert GEMMs are visible during calibration. diff --git a/examples/vllm_serve/vllm_ptq_utils.py b/examples/vllm_serve/vllm_ptq_utils.py index 709d6532fb3..ca9e276b9b4 100644 --- a/examples/vllm_serve/vllm_ptq_utils.py +++ b/examples/vllm_serve/vllm_ptq_utils.py @@ -36,8 +36,131 @@ def _create_new_data_cls(data_cls, **kwargs): return data_cls(**filtered_kwargs) +def _get_calibration_block_count( + model_runner: Any, +) -> Callable[[int, Any], int] | None: + """Return the block reservation policy supported by the installed vLLM.""" + vllm_config = model_runner.vllm_config + + try: + from vllm.v1.worker.gpu.warmup import _reserved_block_count + except ImportError: + try: + from vllm.utils.math_utils import cdiv + from vllm.v1.kv_cache_interface import CrossAttentionSpec, MambaSpec + except ImportError: + return None + + def block_count(num_tokens: int, kv_cache_spec: Any) -> int: + """Calculate the vLLM 0.26 warmup block reservation.""" + # vLLM 0.26's warmup reservation policy. + if isinstance(kv_cache_spec, CrossAttentionSpec): + num_tokens = 0 + num_blocks = cdiv(num_tokens, kv_cache_spec.block_size) + if isinstance(kv_cache_spec, MambaSpec) and kv_cache_spec.mamba_cache_mode == "align": + num_blocks += kv_cache_spec.num_speculative_blocks + return num_blocks + + else: + + def block_count(num_tokens: int, kv_cache_spec: Any) -> int: + """Calculate the current vLLM warmup block reservation.""" + # Calibration runs before model_state is initialized, so call the + # underlying reservation policy rather than _warmup_block_counter. + return _reserved_block_count( + num_tokens, + kv_cache_spec, + num_lookahead_tokens=vllm_config.num_lookahead_tokens, + max_model_len=model_runner.max_model_len, + max_encoder_len=0, + ) + + return block_count + + +def _allocate_calibration_blocks( + self: Any, sequence_lengths: list[int] +) -> tuple[list[tuple[list[int], ...]], list[int] | None]: + """Allocate scheduler-compatible scratch blocks for calibration requests. + + vLLM 0.28 treats block 0 as the null block. Its GPU runner expects real block + tables for hybrid attention/Mamba models, even for one-shot prefill requests. + Use vLLM's warmup reservation policy so this stays aligned with each cache + group's KVCacheSpec. + """ + kv_cache_config = self.model_runner.kv_cache_config + kv_cache_groups = kv_cache_config.kv_cache_groups + block_count = _get_calibration_block_count(self.model_runner) + + if block_count is None: + warnings.warn( + "vLLM warmup block reservation helpers were not found; falling back to " + "empty block tables. Hybrid attention/Mamba models may produce NaNs.", + stacklevel=2, + ) + return [tuple([] for _ in kv_cache_groups) for _ in sequence_lengths], None + + next_block_id = 1 # Block 0 is reserved as the null block. + block_ids_batch: list[tuple[list[int], ...]] = [] + allocated_block_ids: list[int] = [] + + for sequence_length in sequence_lengths: + request_block_ids = [] + for group in kv_cache_groups: + num_blocks = block_count(sequence_length, group.kv_cache_spec) + block_ids = list(range(next_block_id, next_block_id + num_blocks)) + next_block_id += num_blocks + allocated_block_ids.extend(block_ids) + request_block_ids.append(block_ids) + block_ids_batch.append(tuple(request_block_ids)) + + if next_block_id > kv_cache_config.num_blocks: + raise RuntimeError( + "Calibration batch requires " + f"{next_block_id - 1} KV cache blocks, but only " + f"{kv_cache_config.num_blocks - 1} non-null blocks are available." + ) + + scheduler_fields = {field.name for field in dataclasses.fields(SchedulerOutput)} + if "new_block_ids_to_zero" in scheduler_fields: + blocks_to_zero = ( + allocated_block_ids if getattr(kv_cache_config, "needs_kv_cache_zeroing", False) else [] + ) + else: + blocks_to_zero = None + return block_ids_batch, blocks_to_zero + + +def _cleanup_calibration_requests( + self: Any, + cleanup_output: SchedulerOutput, + calibration_error: BaseException | None, +) -> None: + """Clean request state without hiding an active calibration error.""" + try: + # Zero-token steps return before forward/sampling, so no sample_tokens call is needed. + self.execute_model(cleanup_output) + except Exception as execute_error: + finish_requests = getattr(self.model_runner, "finish_requests", None) + if finish_requests is None: + if calibration_error is not None: + raise calibration_error from execute_error + raise + + try: + finish_requests(cleanup_output) + except Exception as finish_error: + if calibration_error is not None: + finish_error.__cause__ = execute_error + raise calibration_error from finish_error + raise finish_error from execute_error + + def calibrate_fun(calib_dataloader: DataLoader, self: Any) -> Callable[[Any], None]: + """Create a calibration loop backed by the vLLM worker scheduler.""" + def calibrate_loop(model: Any) -> None: + """Calibrate the model with batches submitted through the scheduler.""" for batch_idx, batch in tqdm(enumerate(calib_dataloader)): input_ids_batch = batch["input_ids"] @@ -56,7 +179,9 @@ def calibrate_loop(model: Any) -> None: input_ids_list_batch = [input_ids_list_batch] num_groups = len(self.model_runner.kv_cache_config.kv_cache_groups) - empty_block_ids = tuple([] for _ in range(num_groups)) + block_ids_batch, new_block_ids_to_zero = _allocate_calibration_blocks( + self, [len(input_ids) for input_ids in input_ids_list_batch] + ) scheduled_new_reqs = [] num_scheduled_tokens = {} @@ -74,7 +199,7 @@ def calibrate_loop(model: Any) -> None: mm_features=[], sampling_params=SamplingParams(max_tokens=1), pooling_params=None, - block_ids=empty_block_ids, + block_ids=block_ids_batch[seq_idx], num_computed_tokens=0, lora_request=None, ) @@ -96,40 +221,36 @@ def calibrate_loop(model: Any) -> None: kv_connector_metadata=None, structured_output_request_ids={}, grammar_bitmask=None, + new_block_ids_to_zero=new_block_ids_to_zero, + ) + # Submit a zero-token scheduler step after the request has been + # registered. This is the vLLM 0.28 cleanup path and removes + # request-scoped attention/Mamba state from the persistent batch. + cleanup_output = _create_new_data_cls( + type(scheduler_output), + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens={}, + total_num_scheduled_tokens=0, + scheduled_spec_decode_tokens={}, + scheduled_encoder_inputs={}, + num_common_prefix_blocks=[0] * num_groups, + finished_req_ids=set(num_scheduled_tokens), + free_encoder_mm_hashes=[], + kv_connector_metadata=None, + structured_output_request_ids={}, + grammar_bitmask=None, ) try: output = self.execute_model(scheduler_output) if hasattr(self, "sample_tokens"): if output is None: # TODO: make this default when vllm <= 0.11 is outdated self.sample_tokens(None) - finally: - # finish_requests runs before add_requests inside execute_model, so - # req IDs aren't registered yet at that point — call it directly after. - # Wrap in try/except so a cleanup error never masks the original exception. - try: - if hasattr(self.model_runner, "finish_requests"): - cleanup_output = _create_new_data_cls( - type(scheduler_output), - scheduled_new_reqs=[], - scheduled_cached_reqs=scheduler_output.scheduled_cached_reqs, - num_scheduled_tokens={}, - total_num_scheduled_tokens=0, - scheduled_spec_decode_tokens={}, - scheduled_encoder_inputs={}, - num_common_prefix_blocks=scheduler_output.num_common_prefix_blocks, - finished_req_ids=set(num_scheduled_tokens.keys()), - free_encoder_mm_hashes=[], - kv_connector_metadata=None, - structured_output_request_ids={}, - grammar_bitmask=None, - ) - self.model_runner.finish_requests(cleanup_output) - else: - warnings.warn( - "model_runner.finish_requests not found; request state may leak during calibration." - ) - except Exception: - warnings.warn("Failed to clean up request state after calibration batch.") + except BaseException as calibration_error: + _cleanup_calibration_requests(self, cleanup_output, calibration_error) + raise + + _cleanup_calibration_requests(self, cleanup_output, calibration_error=None) return calibrate_loop @@ -171,6 +292,7 @@ def update_kv_cfg_for_mla(model: torch.nn.Module, kv_quant_cfg: list) -> list: def get_quant_config(quant_config: dict[str, Any], model: Any) -> dict[str, Any]: + """Resolve and merge model and KV-cache quantization configuration.""" import copy if quant_config["recipe_path"]: diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py index 038bbd8e978..0f440b93deb 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py @@ -28,6 +28,7 @@ from __future__ import annotations +import builtins import gc import importlib.util from pathlib import Path @@ -69,6 +70,245 @@ def _load_example_module(name: str): return module +def _calibration_worker( + num_blocks: int, + *, + cache_specs=("attention", "mamba"), + needs_kv_cache_zeroing=True, +): + cache_groups = [SimpleNamespace(kv_cache_spec=spec) for spec in cache_specs] + return SimpleNamespace( + model_runner=SimpleNamespace( + kv_cache_config=SimpleNamespace( + kv_cache_groups=cache_groups, + num_blocks=num_blocks, + needs_kv_cache_zeroing=needs_kv_cache_zeroing, + ) + ) + ) + + +def _patch_vllm_imports(monkeypatch, modules): + real_import = builtins.__import__ + + def fake_import(name, globals=None, locals=None, fromlist=(), level=0): + if name in modules: + imported = modules[name] + if isinstance(imported, BaseException): + raise imported + return imported + return real_import(name, globals, locals, fromlist, level) + + monkeypatch.setattr(builtins, "__import__", fake_import) + + +def test_get_calibration_block_count_uses_vllm_028_reservation_helper(monkeypatch): + """The current vLLM adapter must forward every warmup reservation argument.""" + module = _load_example_module("vllm_ptq_utils") + reserved_block_count = Mock(return_value=4) + _patch_vllm_imports( + monkeypatch, + {"vllm.v1.worker.gpu.warmup": SimpleNamespace(_reserved_block_count=reserved_block_count)}, + ) + model_runner = SimpleNamespace( + vllm_config=SimpleNamespace(num_lookahead_tokens=3), + max_model_len=2048, + ) + kv_cache_spec = object() + + block_count = module._get_calibration_block_count(model_runner) + + assert block_count is not None + assert block_count(128, kv_cache_spec) == 4 + reserved_block_count.assert_called_once_with( + 128, + kv_cache_spec, + num_lookahead_tokens=3, + max_model_len=2048, + max_encoder_len=0, + ) + + +def test_get_calibration_block_count_uses_vllm_026_reservation_policy(monkeypatch): + """The vLLM 0.26 adapter must preserve its cross-attention and Mamba rules.""" + module = _load_example_module("vllm_ptq_utils") + + class CrossAttentionSpec: + block_size = 16 + + class MambaSpec: + block_size = 16 + mamba_cache_mode = "align" + num_speculative_blocks = 2 + + cdiv = Mock( + side_effect=lambda numerator, denominator: (numerator + denominator - 1) // denominator + ) + _patch_vllm_imports( + monkeypatch, + { + "vllm.v1.worker.gpu.warmup": ImportError("0.28 helper unavailable"), + "vllm.utils.math_utils": SimpleNamespace(cdiv=cdiv), + "vllm.v1.kv_cache_interface": SimpleNamespace( + CrossAttentionSpec=CrossAttentionSpec, + MambaSpec=MambaSpec, + ), + }, + ) + model_runner = SimpleNamespace( + vllm_config=SimpleNamespace(), + max_model_len=2048, + ) + + block_count = module._get_calibration_block_count(model_runner) + + assert block_count is not None + assert block_count(33, SimpleNamespace(block_size=16)) == 3 + assert block_count(33, CrossAttentionSpec()) == 0 + assert block_count(33, MambaSpec()) == 5 + assert cdiv.call_args_list == [ + ((33, 16),), + ((0, 16),), + ((33, 16),), + ] + + +def test_allocate_calibration_blocks_assigns_non_null_blocks(monkeypatch): + """Scratch block tables must use unique non-null blocks for every request and group.""" + module = _load_example_module("vllm_ptq_utils") + block_count = Mock(side_effect=[1, 2, 2, 1]) + monkeypatch.setattr( + module, + "_get_calibration_block_count", + Mock(return_value=block_count), + ) + + block_tables, blocks_to_zero = module._allocate_calibration_blocks( + _calibration_worker(num_blocks=7), + sequence_lengths=[8, 16], + ) + + assert block_tables == [ + ([1], [2, 3]), + ([4, 5], [6]), + ] + assert block_count.call_args_list == [ + ((8, "attention"),), + ((8, "mamba"),), + ((16, "attention"),), + ((16, "mamba"),), + ] + + scheduler_fields = {field.name for field in module.dataclasses.fields(module.SchedulerOutput)} + expected_blocks_to_zero = ( + [1, 2, 3, 4, 5, 6] if "new_block_ids_to_zero" in scheduler_fields else None + ) + assert blocks_to_zero == expected_blocks_to_zero + + +def test_allocate_calibration_blocks_skips_zeroing_for_attention_only_cache(monkeypatch): + """Attention-only caches have no block zeroer and must receive an empty zeroing list.""" + module = _load_example_module("vllm_ptq_utils") + monkeypatch.setattr( + module, + "_get_calibration_block_count", + Mock(return_value=Mock(return_value=1)), + ) + + block_tables, blocks_to_zero = module._allocate_calibration_blocks( + _calibration_worker( + num_blocks=4, + cache_specs=("attention",), + needs_kv_cache_zeroing=False, + ), + sequence_lengths=[8], + ) + + assert block_tables == [([1],)] + scheduler_fields = {field.name for field in module.dataclasses.fields(module.SchedulerOutput)} + expected_blocks_to_zero = [] if "new_block_ids_to_zero" in scheduler_fields else None + assert blocks_to_zero == expected_blocks_to_zero + + +def test_allocate_calibration_blocks_rejects_insufficient_capacity(monkeypatch): + """Scratch block allocation must account for block 0 being unavailable.""" + module = _load_example_module("vllm_ptq_utils") + monkeypatch.setattr( + module, + "_get_calibration_block_count", + Mock(return_value=Mock(side_effect=[1, 2, 2, 1])), + ) + + with pytest.raises( + RuntimeError, + match=( + r"Calibration batch requires 6 KV cache blocks, " + r"but only 5 non-null blocks are available\." + ), + ): + module._allocate_calibration_blocks( + _calibration_worker(num_blocks=6), + sequence_lengths=[8, 16], + ) + + +@pytest.mark.parametrize("has_calibration_error", [False, True]) +def test_cleanup_failure_preserves_calibration_error(has_calibration_error): + """Cleanup must fail closed without replacing an active calibration error.""" + module = _load_example_module("vllm_ptq_utils") + execute_error = RuntimeError("scheduler cleanup failed") + finish_error = RuntimeError("legacy cleanup failed") + calibration_error = ValueError("calibration failed") if has_calibration_error else None + worker = SimpleNamespace( + execute_model=Mock(side_effect=execute_error), + model_runner=SimpleNamespace(finish_requests=Mock(side_effect=finish_error)), + ) + + expected_error = calibration_error or finish_error + with pytest.raises(type(expected_error)) as raised: + module._cleanup_calibration_requests(worker, object(), calibration_error) + + assert raised.value is expected_error + if calibration_error is not None: + assert calibration_error.__cause__ is finish_error + assert finish_error.__cause__ is execute_error + + +@pytest.mark.parametrize("has_calibration_error", [False, True]) +def test_cleanup_without_legacy_fallback_preserves_primary_error(has_calibration_error): + """Missing legacy cleanup must preserve the most useful primary error.""" + module = _load_example_module("vllm_ptq_utils") + execute_error = RuntimeError("scheduler cleanup failed") + calibration_error = ValueError("calibration failed") if has_calibration_error else None + worker = SimpleNamespace( + execute_model=Mock(side_effect=execute_error), + model_runner=SimpleNamespace(), + ) + + expected_error = calibration_error or execute_error + with pytest.raises(type(expected_error)) as raised: + module._cleanup_calibration_requests(worker, object(), calibration_error) + + assert raised.value is expected_error + if calibration_error is not None: + assert calibration_error.__cause__ is execute_error + + +def test_cleanup_uses_legacy_fallback(): + """A successful legacy cleanup may recover from an unsupported scheduler step.""" + module = _load_example_module("vllm_ptq_utils") + cleanup_output = object() + finish_requests = Mock() + worker = SimpleNamespace( + execute_model=Mock(side_effect=RuntimeError("unsupported scheduler cleanup")), + model_runner=SimpleNamespace(finish_requests=finish_requests), + ) + + module._cleanup_calibration_requests(worker, cleanup_output, calibration_error=None) + + finish_requests.assert_called_once_with(cleanup_output) + + class _NativeAttention(torch.nn.Module): def forward(self, query, key, value, *args, **kwargs): return query, key, value From 961741c0009fde21b676b6a0226bc1c45ea8abe1 Mon Sep 17 00:00:00 2001 From: yeyu-nvidia Date: Thu, 17 Sep 2026 16:35:00 -0700 Subject: [PATCH 19/28] docs(recipes): point chat_template at where the templates actually live (#2460) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Documentation `dflash.yaml` tells the reader that chat templates live in `modelopt_recipes/general/speculative_decoding/chat_templates/`. That directory does not exist, and this comment is its only mention anywhere in the tree: ``` $ git grep -n 'speculative_decoding/chat_templates' modelopt_recipes/general/speculative_decoding/dflash.yaml:19: # Templates are in modelopt_recipes/general/speculative_decoding/chat_templates/ ``` Templates actually sit beside each launcher example — `tools/launcher/examples/Qwen/Qwen3-8B/chat_template_train.jinja`, `.../MiniMax/MiniMax-M2.7-DFlash/chat_template_train.jinja`, and so on. Split out of #2201, where this two-line comment was the only reason `modelopt-recipes-codeowners` was a required reviewer on a skills-documentation PR. ### Usage No behaviour change — comment only. ### Testing None needed; the file's only change is a YAML comment. `pre-commit` passes. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - 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?: N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - Did you get Claude approval on this PR?: N/A 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit - **Documentation** - Updated speculative decoding guidance to clarify where each model’s chat template is maintained alongside its launcher example. Signed-off-by: Ye Yu Co-authored-by: Claude Opus 5 Signed-off-by: Chad Voegele --- modelopt_recipes/general/speculative_decoding/dflash.yaml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/modelopt_recipes/general/speculative_decoding/dflash.yaml b/modelopt_recipes/general/speculative_decoding/dflash.yaml index 021cccd475d..737789dc4ed 100644 --- a/modelopt_recipes/general/speculative_decoding/dflash.yaml +++ b/modelopt_recipes/general/speculative_decoding/dflash.yaml @@ -16,7 +16,8 @@ data: offline_data_path: # Jinja chat template with {% generation %} tags for answer_only_loss. # Required when answer_only_loss=true. Set in per-model launcher YAML. - # Templates are in modelopt_recipes/general/speculative_decoding/chat_templates/ + # Each model keeps its own beside its launcher example, e.g. + # tools/launcher/examples/Qwen/Qwen3-8B/chat_template_train.jinja chat_template: # maps to TrainingArguments (main.py) From 0470b12677bd7365d0563478e52a3e5dbf08bc0b Mon Sep 17 00:00:00 2001 From: Ajinkya Rasane <131806219+ajrasane@users.noreply.github.com> Date: Thu, 17 Sep 2026 23:47:25 -0400 Subject: [PATCH 20/28] [6771663] Preserve ONNX API output types when wiring casts (#2451) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix Preserves the public ONNX graph I/O types captured at the API boundary when `PrecisionConverter` wires output casts. Type inference can change the working graph's output declaration before conversion; consulting that mutated declaration caused the required cast back to the original public type to be discarded and metadata restoration to fail. The converter now derives its I/O type map from the preserved boundary metadata and uses that map when deciding whether a cast should become a public graph output. A regression test covers an FP32 output whose working declaration is inferred as FP16, and the changelog records the corrected behavior. ```python converted = convert_to_f16(model, keep_io_types=True) ``` - Ran `pytest tests/unit/onnx/autocast/test_precisionconverter.py` (186 passed). - Ran `pytest tests/unit/onnx/autocast` (249 passed). - Ran Ruff check and format validation on the changed Python files. - Verified the original minimal end-to-end reproduction with Python 3.12 and TensorRT 10.16.1.11; conversion now completes without the output-metadata restoration error. - Verified the full CLI quantization path proceeds through the formerly failing one-Q/DQ scheme and successfully benchmarks the generated TensorRT engine. Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ - 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?: N/A Tracking: [6771663] - **Bug Fixes** - Fixed ONNX FP16 conversion to preserve public graph output types when type inference changes internal declarations. - Ensured output casts are inserted correctly when preserving input/output types is enabled. - **Tests** - Added regression coverage confirming preserved output types, correct cast insertion, and valid ONNX model generation. Signed-off-by: Ajinkya Rasane Co-authored-by: Ajinkya Rasane Signed-off-by: Chad Voegele --- CHANGELOG.rst | 1 + modelopt/onnx/autocast/precisionconverter.py | 19 +++++++------ .../onnx/autocast/test_precisionconverter.py | 27 +++++++++++++++++++ 3 files changed, 39 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index bcc4f7ef9ed..f17b219f478 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -56,6 +56,7 @@ Changelog **Bug Fixes** +- Fix ONNX FP16 conversion failing to preserve public output types when type inference changes a graph output declaration before output casts are inserted. - Fail fast on non-finite AutoQuantize output gradients with an actionable error before accumulating sensitivity scores, without changing attention backend settings. - Fix ONNX INT8 entropy calibration failing or producing invalid quantization parameters for FP16 activations. - Speed up ``mtq.quantize`` on FSDP2-sharded fused-MoE models. Promoting static-block weight quantizers gathered each expert's slice of the fused weight across ranks even though only quantizer state is read, adding a collective per expert to calibration. diff --git a/modelopt/onnx/autocast/precisionconverter.py b/modelopt/onnx/autocast/precisionconverter.py index 7db660b29b4..a78249fdca4 100644 --- a/modelopt/onnx/autocast/precisionconverter.py +++ b/modelopt/onnx/autocast/precisionconverter.py @@ -142,13 +142,6 @@ def __init__( self.low_precision_type = PRECISION_MAP[low_precision_type] self.high_precision_type = PRECISION_MAP["fp32"] - # Preserve original network inputs and outputs for sanity checks - self.original_network_io = { - io.name: io.type.tensor_type.elem_type for io in self.model.graph.input - } - self.original_network_io.update( - {io.name: io.type.tensor_type.elem_type for io in self.model.graph.output} - ) self.original_network_io_metadata = ( { "input": [deepcopy(io) for io in self.model.graph.input], @@ -160,6 +153,13 @@ def __init__( for field, values in original_network_io_metadata.items() } ) + # Preserve the public I/O types captured at the API boundary. Type inference may have + # changed the working model's declarations before the converter is initialized. + self.original_network_io = { + io.name: io.type.tensor_type.elem_type + for values in self.original_network_io_metadata.values() + for io in values + } self.min_opset = min_opset self.max_ir_version = max_ir_version self.trt_plugins = trt_plugins @@ -1440,7 +1440,10 @@ def _add_cast( # Update network output for output in self.model.graph.output: if output.name == tensor_name and ( - (self.keep_io_types and cast_to.onnx_type == output.type.tensor_type.elem_type) + ( + self.keep_io_types + and cast_to.onnx_type == self.original_network_io.get(tensor_name) + ) or ( not self.keep_io_types and cast_to.onnx_type == self.low_precision_type.onnx_type diff --git a/tests/unit/onnx/autocast/test_precisionconverter.py b/tests/unit/onnx/autocast/test_precisionconverter.py index b480bb7c24e..361cef26a5e 100644 --- a/tests/unit/onnx/autocast/test_precisionconverter.py +++ b/tests/unit/onnx/autocast/test_precisionconverter.py @@ -2220,6 +2220,33 @@ def test_convert_to_f16_restores_public_io_metadata_from_entry_boundary(): onnx.checker.check_model(converted, full_check=True) +def test_convert_to_f16_preserves_declared_output_type_after_inference_changes_it(): + graph_input = helper.make_tensor_value_info("X", TensorProto.FLOAT16, [2, 3]) + graph_output = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2, 3]) + node = helper.make_node("Identity", ["X"], ["Y"], name="Identity_0") + graph = helper.make_graph([node], "inferred_output_type", [graph_input], [graph_output]) + model = helper.make_model( + graph, + producer_name="inferred_output_type", + opset_imports=[helper.make_opsetid("", 19)], + ir_version=10, + ) + + converted = convert_to_f16( + model, keep_io_types=True, op_block_list=[], trt_plugins=[], opset=19 + ) + + output = next(vi for vi in converted.graph.output if vi.name == "Y") + assert output.type.tensor_type.elem_type == TensorProto.FLOAT + output_producers = [node for node in converted.graph.node if "Y" in node.output] + assert len(output_producers) == 1 + assert output_producers[0].op_type == "Cast" + assert next(attr.i for attr in output_producers[0].attribute if attr.name == "to") == ( + TensorProto.FLOAT + ) + onnx.checker.check_model(converted, full_check=True) + + def test_convert_to_f16_refreshes_gathernd_pre_cast_declaration(monkeypatch): def discover_test_plugins_without_trt(self): self.custom_ops = { From a2cee6abc7076721db74c241f8555733225e7be2 Mon Sep 17 00:00:00 2001 From: Ajinkya Rasane <131806219+ajrasane@users.noreply.github.com> Date: Fri, 18 Sep 2026 13:22:28 -0400 Subject: [PATCH 21/28] [5565357] Fix SDXL NVFP4 export and performance (#2336) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix Adds a compact SDXL and SDXL-Turbo mixed-precision FP4 recipe: - block-16 NVFP4 for non-QKV Linear/GEMM layers; - FP8 for Conv2d layers; - high-precision Q/K/V projection Linears to preserve TensorRT horizontal fusion; - optional FP8 MHA quantization. For SDXL FP4 export, Conv2d quantizers export directly through the shared FP8 custom-op path. The previous `generate_fp8_scales` plus `convert_zp_fp8` INT8 zero-point workaround is removed. The graph then uses the existing FP8 Q/DQ normalization and `NVFP4QuantExporter` lowering, with opset 23 for FLOAT4 support. Flux FP8 export also saves the graph returned by its RoPE weight conversion. This PR also changes shared exporter behavior: - `_fp8_quantize` refreshes ONNX shape/type inference after applying the custom FP8 operator's uint8 output metadata, affecting all FP8 ONNX exports through this symbolic. - `_quantized_sdpa` derives `disable_fp8_mha` from the live Q/K/V quantizer state instead of a restored private module flag. Other model recipe configurations remain unchanged. ```bash python quantize.py \ --model sdxl-1.0 \ --model-dtype Half \ --trt-high-precision-dtype Half \ --format fp4 \ --block-size 16 \ --batch-size 2 \ --calib-size 128 \ --n-steps 20 \ --quantized-torch-ckpt-save-path ./sdxl-fp4 \ --onnx-dir ./onnx-sdxl-fp4 ``` - CPU-only focused and generic NVFP4 exporter tests: 44 passed in 4.35 seconds. - Focused Flux returned-graph save test: 1 passed. - Required Linux unit CI at `034fe23ec` passed with the `all` dependency set, including `tests/unit/examples/test_diffusers_fp4.py`. - Latest changed-file pre-commit checks: all passed. - TensorRT 10.14 on a B200 GPU: - 302 native block-scaled NVFP4 GEMM tactics; - 38 native FP8 Conv tactics; - no FP4 Q/K/V projections; - all 11 FP16 Q/K/V projection-fusion groups preserved; - three alternating batch-2 profiles measured 18.614 ms FP4 versus 20.028 ms FP16 median UNet latency, a 7.06% reduction. - FP8 SDXL/SD3 ONNX-to-TensorRT end-to-end runs were not executed because they require explicit approval. The existing end-to-end test matrix now includes SD3 FP8 alongside SDXL FP8. Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ — no public API or CLI flags change; the shared changes preserve the intended FP8 export and attention behavior. - 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)?: ✅ — the shared NVFP4 opset, FP8 shape-inference, and Diffusers attention-policy changes are recorded under bug fixes. - Did you get Claude approval on this PR?: N/A Tracking: [5565357] - **New Features** - Added SDXL support for mixed NVFP4/FP8 quantization, including convolution and softmax handling. - Added an SDXL quantization preset for streamlined post-training quantization workflows. - Expanded FP4 ONNX export support to Flux and SDXL, with improved FP4/FP8 graph processing and export reliability. - Added automatic quantization policy and format restoration from checkpoints. - **Documentation** - Documented SDXL layer behavior, optional FP8 attention quantization, and Blackwell/TensorRT requirements for FP4 and FP8 deployment. > 🤖 _Generated by Codex (AI agent)._ --------- Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> Co-authored-by: Codex Signed-off-by: Chad Voegele --- CHANGELOG.rst | 1 + .../quantization/ONNX-TRT-Deployment.md | 4 +- examples/diffusers/quantization/config.py | 3 + .../quantization/onnx_utils/export.py | 27 +- .../onnx_utils/fp8_onnx_graphsurgeon.py | 23 -- examples/diffusers/quantization/quantize.py | 72 +++--- examples/diffusers/quantization/utils.py | 3 - modelopt/onnx/export/nvfp4_exporter.py | 6 + modelopt/torch/quantization/export_onnx.py | 5 +- .../plugins/diffusion/diffusers.py | 10 +- .../ptq/presets/diffusers/nvfp4_fp8_conv.yaml | 53 ++++ tests/examples/diffusers/test_diffusers.py | 24 ++ .../quantization/test_onnx_export_cuda.py | 4 - tests/unit/examples/test_diffusers_fp4.py | 234 ++++++++++++++++++ .../unit/onnx/quantization/test_qdq_utils.py | 6 +- .../quantization/test_onnx_export_cpu.py | 58 ++++- 16 files changed, 428 insertions(+), 105 deletions(-) create mode 100644 modelopt_recipes/configs/ptq/presets/diffusers/nvfp4_fp8_conv.yaml create mode 100644 tests/unit/examples/test_diffusers_fp4.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index f17b219f478..4bdc9ae0564 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -56,6 +56,7 @@ Changelog **Bug Fixes** +- Fix shared ONNX export metadata and Diffusers attention policy: every ``NVFP4QuantExporter`` post-process now upgrades the default-domain opset to at least 23, all FP8 custom-op exports re-run ONNX shape/type inference after setting output metadata, and quantized SDPA derives FP8 MHA enablement from the live Q/K/V quantizers instead of honoring a caller-set ``_disable_fp8_mha`` attribute. - Fix ONNX FP16 conversion failing to preserve public output types when type inference changes a graph output declaration before output casts are inserted. - Fail fast on non-finite AutoQuantize output gradients with an actionable error before accumulating sensitivity scores, without changing attention backend settings. - Fix ONNX INT8 entropy calibration failing or producing invalid quantization parameters for FP16 activations. diff --git a/examples/diffusers/quantization/ONNX-TRT-Deployment.md b/examples/diffusers/quantization/ONNX-TRT-Deployment.md index 57448b8a38e..b6933c88f72 100644 --- a/examples/diffusers/quantization/ONNX-TRT-Deployment.md +++ b/examples/diffusers/quantization/ONNX-TRT-Deployment.md @@ -28,12 +28,12 @@ python quantize.py \ #### FLUX-Dev|SDXL|SDXL-Turbo|LTX-Video FP8/FP4 [Script](./quantize.py) -*In our example code, FP4 is only supported for Flux. However, you can modify our script to enable FP4 format support for your own model.* +FP4 ONNX export is supported for Flux and SDXL. ```sh python quantize.py \ --model {flux-dev|sdxl-1.0|sdxl-turbo|ltx-video-dev} --model-dtype {Half|BFloat16} --trt-high-precision-dtype {Half|BFloat16} \ - --format {fp8|fp4} --batch-size 2 --calib-size {128|256} --quantize-mha \ + --format {fp8|fp4} --batch-size 2 --calib-size {128|256} \ --n-steps 20 --quantized-torch-ckpt-save-path ./{MODEL_NAME}.pt --collect-method default \ --onnx-dir {ONNX_DIR} ``` diff --git a/examples/diffusers/quantization/config.py b/examples/diffusers/quantization/config.py index cb8fdf3a5da..3ab8f5db87c 100644 --- a/examples/diffusers/quantization/config.py +++ b/examples/diffusers/quantization/config.py @@ -31,6 +31,9 @@ NVFP4_FP8_MHA_CONFIG = load_config( "configs/ptq/presets/diffusers/nvfp4_fp8_mha", schema_type=QuantizeConfig ).model_dump(exclude_unset=True) +NVFP4_FP8_CONV_CONFIG = load_config( + "configs/ptq/presets/diffusers/nvfp4_fp8_conv", schema_type=QuantizeConfig +).model_dump(exclude_unset=True) def set_quant_config_attr(quant_config, trt_high_precision_dtype, quant_algo, **kwargs): diff --git a/examples/diffusers/quantization/onnx_utils/export.py b/examples/diffusers/quantization/onnx_utils/export.py index 5da795f0f48..246dd1a9881 100644 --- a/examples/diffusers/quantization/onnx_utils/export.py +++ b/examples/diffusers/quantization/onnx_utils/export.py @@ -51,8 +51,6 @@ from modelopt.torch.quantization.export_onnx import configure_linear_module_onnx_quantizers from modelopt.torch.utils import torch_to -from .fp8_onnx_graphsurgeon import convert_zp_fp8 - MODEL_ID_TO_DYNAMIC_AXES = { "sdxl-1.0": { "sample": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}, @@ -124,18 +122,6 @@ def flux_convert_rope_weight_type(onnx_graph): return gs.export_onnx(graph) -def generate_fp8_scales(backbone): - # temporary solution due to a known bug in torch.onnx._dynamo_export - for _, module in backbone.named_modules(): - if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d)) and ( - hasattr(module.input_quantizer, "_amax") and module.input_quantizer is not None - ): - module.input_quantizer._num_bits = 8 - module.weight_quantizer._num_bits = 8 - module.input_quantizer._amax = module.input_quantizer._amax * (127 / 448.0) - module.weight_quantizer._amax = module.weight_quantizer._amax * (127 / 448.0) - - def _gen_dummy_inp_and_dyn_shapes_sdxl(backbone, min_bs=1, opt_bs=1): assert isinstance(backbone, UNet2DConditionModel) or isinstance( backbone._orig_mod, UNet2DConditionModel @@ -469,7 +455,6 @@ def modelopt_export_sd(backbone, onnx_dir, model_name, precision): tmp_subfolder = tempfile.mkdtemp(prefix="myapp_") tmp_output = Path(f"{tmp_subfolder}/{model_file_name}") q_output = Path(f"{onnx_dir}/{model_file_name}") - quantizer_context = ( configure_linear_module_onnx_quantizers(backbone) if precision == "fp4" else nullcontext() ) @@ -536,16 +521,8 @@ def modelopt_export_sd(backbone, onnx_dir, model_name, precision): ) print(f"Saved at {tmp_output}") onnx_model = onnx.load(str(tmp_output), load_external_data=True) - if precision == "fp8": - if not model_name.startswith("flux"): - graph = gs.import_onnx(onnx_model) - graph.cleanup().toposort() - onnx_model = gs.export_onnx(graph) - onnx_model = convert_zp_fp8(onnx_model) - graph = gs.import_onnx(onnx_model) - onnx_model = gs.export_onnx(graph.cleanup()) - else: - flux_convert_rope_weight_type(onnx_model) + if precision == "fp8" and model_name.startswith("flux"): + onnx_model = flux_convert_rope_weight_type(onnx_model) if precision == "fp4": onnx_model = NVFP4QuantExporter.process_model(onnx_model) save_onnx(onnx_model, q_output) diff --git a/examples/diffusers/quantization/onnx_utils/fp8_onnx_graphsurgeon.py b/examples/diffusers/quantization/onnx_utils/fp8_onnx_graphsurgeon.py index 7194e672635..7904d27e0a3 100644 --- a/examples/diffusers/quantization/onnx_utils/fp8_onnx_graphsurgeon.py +++ b/examples/diffusers/quantization/onnx_utils/fp8_onnx_graphsurgeon.py @@ -97,29 +97,6 @@ def insert_cast(graph, input_tensor, attrs): next_node.inputs[idx] = output_tensor -def convert_zp_fp8(onnx_graph): - """ - Convert Q/DQ zero datatype from INT8 to FP8. - We use this WAR because FP8 Conv cannot be exported to ONNX directly. - The workaround is to first convert the FP8 QDQs into INT8 QDQs, - then modify the ONNX model afterward to change those INT8 QDQs back into FP8 QDQs. - """ - # Find all zero constant nodes - qdq_zero_nodes = set() - for node in onnx_graph.graph.node: - if node.op_type == "QuantizeLinear" and len(node.input) > 2: - qdq_zero_nodes.add(node.input[2]) - - print(f"[WAR], found {len(qdq_zero_nodes)} INT8 QDQ pairs, you can ignore this message..") - - # Convert zero point datatype from INT8 to FP8. - for node in onnx_graph.graph.node: - if node.output[0] in qdq_zero_nodes: - node.attribute[0].t.data_type = onnx.TensorProto.FLOAT8E4M3FN - - return onnx_graph - - def cast_resize_io(graph): """ After all activations and weights are converted to fp16, we will diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index 1d71c088652..9262adbaac1 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -27,6 +27,7 @@ FP8_DEFAULT_CONFIG, INT8_DEFAULT_CONFIG, NVFP4_DEFAULT_CONFIG, + NVFP4_FP8_CONV_CONFIG, NVFP4_FP8_MHA_CONFIG, reset_set_int8_config, set_quant_config_attr, @@ -55,6 +56,9 @@ import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.export import export_hf_checkpoint +from modelopt.torch.quantization.nn import TensorQuantizer + +_SDXL_MODEL_TYPES = (ModelType.SDXL_BASE, ModelType.SDXL_TURBO) def setup_logging(verbose: bool = False) -> logging.Logger: @@ -130,7 +134,9 @@ def get_quant_config(self, n_steps: int, backbone: torch.nn.Module) -> Any: elif self.config.format == QuantFormat.FP8: base_cfg = FP8_DEFAULT_CONFIG elif self.config.format == QuantFormat.FP4: - if self.model_config.model_type.value.startswith("flux"): + if self.model_config.model_type in _SDXL_MODEL_TYPES: + base_cfg = NVFP4_FP8_CONV_CONFIG + elif self.model_config.model_type.value.startswith("flux"): base_cfg = NVFP4_FP8_MHA_CONFIG else: base_cfg = NVFP4_DEFAULT_CONFIG @@ -271,23 +277,6 @@ def __init__( self.logger = logger self.pipeline_manager = pipeline_manager - def _has_conv_layers(self, model: torch.nn.Module) -> bool: - """ - Check if the model contains any convolutional layers. - - Args: - model: Model to check - - Returns: - True if model contains Conv layers, False otherwise - """ - for module in model.modules(): - if isinstance(module, (torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d)) and ( - module.input_quantizer.is_enabled or module.weight_quantizer.is_enabled - ): - return True - return False - def save_checkpoint( self, backbone: torch.nn.Module, @@ -335,15 +324,10 @@ def export_onnx( # Deferred: the ONNX stack (onnx, onnx_graphsurgeon, ...) is only needed # for --onnx-dir exports; HF-checkpoint-only runs must not require it. - from onnx_utils.export import generate_fp8_scales, modelopt_export_sd + from onnx_utils.export import modelopt_export_sd self.logger.info(f"Starting ONNX export to {self.config.onnx_dir}") - if quant_format == QuantFormat.FP8 and self._has_conv_layers(backbone): - self.logger.info( - "Detected quantizing conv layers in backbone. Generating FP8 scales..." - ) - generate_fp8_scales(backbone) self.logger.info("Preparing models for export...") pipe.to("cpu") torch.cuda.empty_cache() @@ -457,7 +441,7 @@ def create_argument_parser() -> argparse.ArgumentParser: %(prog)s --model ltx-video-dev --format fp8 --batch-size 1 --calib-size 32 --ltx-skip-upsampler # Restore and export a previously quantized model - %(prog)s --model flux-schnell --restore-from checkpoint.pt --onnx-dir ./exports/ + %(prog)s --model flux-schnell --restore-from ./checkpoints/ --onnx-dir ./exports/ """, ) model_group = parser.add_argument_group("Model Configuration") @@ -586,7 +570,9 @@ def create_argument_parser() -> argparse.ArgumentParser: help="Directory for HuggingFace checkpoint export", ) export_group.add_argument( - "--restore-from", type=str, help="Path to restore from previous checkpoint" + "--restore-from", + type=str, + help="Checkpoint directory; quantization format and MHA policy are restored automatically", ) export_group.add_argument( "--trt-high-precision-dtype", @@ -600,6 +586,25 @@ def create_argument_parser() -> argparse.ArgumentParser: return parser +def _infer_restored_quantization_format( + backbones: list[tuple[str, torch.nn.Module]], +) -> QuantFormat: + has_nvfp4 = False + has_fp8 = False + + for _, backbone in backbones: + for module in backbone.modules(): + if isinstance(module, TensorQuantizer) and module.is_enabled: + has_nvfp4 |= module.is_nvfp4_dynamic or module.is_nvfp4_static + has_fp8 |= module.is_fp8 + + if has_nvfp4: + return QuantFormat.FP4 + if has_fp8: + return QuantFormat.FP8 + return QuantFormat.INT8 + + def main() -> None: from diffusers.models.normalization import RMSNorm as DiffuserRMSNorm @@ -674,9 +679,9 @@ def main() -> None: ) logger.info("Validating configurations...") - quant_config.validate() export_config.validate() if not export_config.restore_from: + quant_config.validate() calib_config.validate() pipeline_manager = PipelineManager(model_config, logger) @@ -685,8 +690,12 @@ def main() -> None: export_manager = ExportManager(export_config, logger, pipeline_manager) - if export_config.restore_from and export_config.restore_from.exists(): + if export_config.restore_from: export_manager.restore_checkpoint() + quant_config.format = _infer_restored_quantization_format( + list(pipeline_manager.iter_backbones()) + ) + logger.info(f"Detected restored quantization format: {quant_config.format.value}") else: logger.info("Initializing calibration...") @@ -716,11 +725,12 @@ def forward_loop(mod): mtq.compress(backbone) logger.info(f"{backbone_name} compression completed") - # For VAE backbones, skip check_conv_and_mha — the whole point - # of VAE quantization is to quantize Conv layers. if backbone_name not in ("video_decoder", "vae"): check_conv_and_mha( - backbone, quant_config.format == QuantFormat.FP4, quant_config.quantize_mha + backbone, + quant_config.format == QuantFormat.FP4 + and model_config.model_type not in _SDXL_MODEL_TYPES, + quant_config.quantize_mha, ) export_manager.save_checkpoint(backbone, backbone_name) diff --git a/examples/diffusers/quantization/utils.py b/examples/diffusers/quantization/utils.py index c3cfdcd5cdd..b7a79e49e70 100644 --- a/examples/diffusers/quantization/utils.py +++ b/examples/diffusers/quantization/utils.py @@ -64,11 +64,8 @@ def check_conv_and_mha(backbone, if_fp4, quantize_mha): ): if hasattr(module, attr): getattr(module, attr).disable() - setattr(module, "_disable_fp8_mha", True) print(f"Disabled Attention layer quantization for layer {name}") - else: - setattr(module, "_disable_fp8_mha", False) def filter_func_ltx_video(name: str) -> bool: diff --git a/modelopt/onnx/export/nvfp4_exporter.py b/modelopt/onnx/export/nvfp4_exporter.py index 338e2725b14..42598af5475 100644 --- a/modelopt/onnx/export/nvfp4_exporter.py +++ b/modelopt/onnx/export/nvfp4_exporter.py @@ -430,4 +430,10 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str): utils.topologically_sort_graph_nodes(graph) + if fp4_qdq_nodes: + default_opset = next( + opset for opset in onnx_model.opset_import if opset.domain in {"", "ai.onnx"} + ) + default_opset.version = max(default_opset.version, 23) + return onnx_model diff --git a/modelopt/torch/quantization/export_onnx.py b/modelopt/torch/quantization/export_onnx.py index e5778c3c96b..2ad3383e572 100644 --- a/modelopt/torch/quantization/export_onnx.py +++ b/modelopt/torch/quantization/export_onnx.py @@ -225,9 +225,12 @@ def _fp8_quantize( "Constant", value_t=torch.tensor(scale_inv).to(torch_dtype_map[inputs.type().scalarType()]), ) - return g.op("trt::TRT_FP8QuantizeLinear", inputs, scale).setType( + quantized = g.op("trt::TRT_FP8QuantizeLinear", inputs, scale).setType( inputs.type().with_dtype(torch.uint8).with_sizes(output_shape) ) + # PyTorch runs shape inference before setType for custom ops, so refresh its reliability state. + torch._C._jit_pass_onnx_node_shape_type_inference(quantized.node(), g.params_dict, g.opset) + return quantized def _fp8_dequantize( diff --git a/modelopt/torch/quantization/plugins/diffusion/diffusers.py b/modelopt/torch/quantization/plugins/diffusion/diffusers.py index f2f6a702479..e92c775c05b 100644 --- a/modelopt/torch/quantization/plugins/diffusion/diffusers.py +++ b/modelopt/torch/quantization/plugins/diffusion/diffusers.py @@ -141,6 +141,14 @@ def _quantized_sdpa(self, *args, **kwargs): q_quantized_scale = self.q_bmm_quantizer._get_amax(query) k_quantized_scale = self.k_bmm_quantizer._get_amax(key) v_quantized_scale = self.v_bmm_quantizer._get_amax(value) + disable_fp8_mha = not all( + quantizer.is_enabled and quantizer.is_fp8 + for quantizer in ( + self.q_bmm_quantizer, + self.k_bmm_quantizer, + self.v_bmm_quantizer, + ) + ) # We don't need to calibrate the output of softmax return self.bmm2_output_quantizer( @@ -155,7 +163,7 @@ def _quantized_sdpa(self, *args, **kwargs): self.q_bmm_quantizer.trt_high_precision_dtype if hasattr(self.q_bmm_quantizer, "trt_high_precision_dtype") else "Half", - self._disable_fp8_mha if hasattr(self, "_disable_fp8_mha") else True, + disable_fp8_mha, ) ) diff --git a/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4_fp8_conv.yaml b/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4_fp8_conv.yaml new file mode 100644 index 00000000000..16e2ebbebd2 --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4_fp8_conv.yaml @@ -0,0 +1,53 @@ +# 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. + +# Diffusers SDXL preset with dynamic NVFP4 Linears and per-tensor FP8 Conv2d layers. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig +imports: + base_disable_all: configs/ptq/units/base_disable_all + fp8: configs/numerics/fp8 + nvfp4: configs/numerics/nvfp4 + +algorithm: max +quant_cfg: + - $import: base_disable_all + - parent_class: nn.Linear + quantizer_name: '*weight_quantizer' + cfg: + $import: nvfp4 + - parent_class: nn.Linear + quantizer_name: '*input_quantizer' + cfg: + $import: nvfp4 + - parent_class: nn.Linear + quantizer_name: '*to_[qkv].input_quantizer' + enable: false + - parent_class: nn.Linear + quantizer_name: '*to_[qkv].weight_quantizer' + enable: false + - parent_class: nn.Conv2d + quantizer_name: '*weight_quantizer' + cfg: + $import: fp8 + - parent_class: nn.Conv2d + quantizer_name: '*input_quantizer' + cfg: + $import: fp8 + - quantizer_name: '*output_quantizer' + enable: false + - quantizer_name: '*softmax_quantizer' + cfg: + $import: fp8 diff --git a/tests/examples/diffusers/test_diffusers.py b/tests/examples/diffusers/test_diffusers.py index 15c5eb44934..979894819de 100644 --- a/tests/examples/diffusers/test_diffusers.py +++ b/tests/examples/diffusers/test_diffusers.py @@ -117,6 +117,17 @@ def inference(self, tmp_path: Path) -> None: quant_algo="smoothquant", collect_method="min-mean", ), + pytest.param( + DiffuserModel( + name="sd3-medium", + path=SD3_PATH, + dtype="Half", + format_type="fp8", + quant_algo="max", + collect_method="default", + ), + marks=minimum_sm(89), + ), pytest.param( DiffuserModel( name="sdxl-1.0", @@ -128,6 +139,17 @@ def inference(self, tmp_path: Path) -> None: ), marks=minimum_sm(89), ), + pytest.param( + DiffuserModel( + name="sdxl-1.0", + path=SDXL_PATH, + dtype="Half", + format_type="fp4", + quant_algo="max", + collect_method="default", + ), + marks=minimum_sm(100), + ), DiffuserModel( name="sdxl-1.0", path=SDXL_PATH, @@ -140,7 +162,9 @@ def inference(self, tmp_path: Path) -> None: ids=[ "flux_schnell_bf16_int8_smoothquant_3.0_min_mean", "sd3_medium_fp16_int8_smoothquant_3.0_min_mean", + "sd3_medium_fp16_fp8_max_3.0_default", "sdxl_1.0_fp16_fp8_max_3.0_default", + "sdxl_1.0_fp16_fp4_max_3.0_default", "sdxl_1.0_fp16_int8_smoothquant_3.0_min_mean", ], ) diff --git a/tests/gpu/torch/quantization/test_onnx_export_cuda.py b/tests/gpu/torch/quantization/test_onnx_export_cuda.py index 300abc52e9e..39f422c75a9 100644 --- a/tests/gpu/torch/quantization/test_onnx_export_cuda.py +++ b/tests/gpu/torch/quantization/test_onnx_export_cuda.py @@ -17,7 +17,6 @@ import pytest import torch -import torch.nn as nn from _test_utils.torch.quantization.onnx_export import TEST_MODELS, onnx_export_tester @@ -40,7 +39,4 @@ def test_onnx_export_cuda(model_cls, num_bits, per_channel_quantization, constan torch.manual_seed(0) model = model_cls() - for _, module in model.named_modules(): - if isinstance(module, nn.Conv2d) and num_bits == (4, 3): - pytest.skip("Conv2d with FP8 quantization is not supported yet") onnx_export_tester(model, "cuda", num_bits, per_channel_quantization, constant_folding, dtype) diff --git a/tests/unit/examples/test_diffusers_fp4.py b/tests/unit/examples/test_diffusers_fp4.py new file mode 100644 index 00000000000..16023e509ad --- /dev/null +++ b/tests/unit/examples/test_diffusers_fp4.py @@ -0,0 +1,234 @@ +# 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. + +import importlib.util +import logging +import sys +from pathlib import Path +from unittest.mock import Mock + +import pytest +import torch +from torch import nn + +pytest.importorskip("onnx") +pytest.importorskip("onnx_graphsurgeon") +pytest.importorskip("diffusers") + +import modelopt.torch.quantization as mtq +from examples.diffusers.quantization.onnx_utils import export as diffusion_export +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.plugins.diffusion import diffusers as diffusers_plugin + +_QUANTIZATION_EXAMPLE = ( + Path(__file__).resolve().parents[3] / "examples" / "diffusers" / "quantization" +) +_LOCAL_IMPORT_NAMES = ( + "calib.plugin_calib", + "calib", + "calibration", + "config", + "models_utils", + "pipeline_manager", + "quantize_config", + "utils", +) + + +def _load_quantize_example(): + spec = importlib.util.spec_from_file_location( + "diffusers_quantize_example", _QUANTIZATION_EXAMPLE / "quantize.py" + ) + assert spec is not None and spec.loader is not None + + original_modules = { + name: sys.modules.pop(name) for name in _LOCAL_IMPORT_NAMES if name in sys.modules + } + sys.path.insert(0, str(_QUANTIZATION_EXAMPLE)) + try: + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + finally: + sys.path.pop(0) + for name in _LOCAL_IMPORT_NAMES: + sys.modules.pop(name, None) + sys.modules.update(original_modules) + return module + + +_quantize = _load_quantize_example() +ModelType = _quantize.ModelType +ModelConfig = _quantize.ModelConfig +QuantFormat = _quantize.QuantFormat +QuantizationConfig = _quantize.QuantizationConfig +Quantizer = _quantize.Quantizer +_infer_restored_quantization_format = _quantize._infer_restored_quantization_format + + +class _RecipeBackbone(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(16, 16, bias=False) + self.attn = nn.Module() + self.attn.to_q = nn.Linear(16, 16, bias=False) + self.attn.to_k = nn.Linear(16, 16, bias=False) + self.attn.to_v = nn.Linear(16, 16, bias=False) + self.conv = nn.Conv2d(4, 4, kernel_size=1, bias=False) + + +def _quantizer(*, num_bits, enabled=True, block_sizes=None): + quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=num_bits, axis=None, block_sizes=block_sizes) + ) + quantizer.amax = torch.tensor(448.0) + if not enabled: + quantizer.disable() + return quantizer + + +_FP8_QUANTIZER_CONFIG = {"num_bits": (4, 3)} +_NVFP4_QUANTIZER_CONFIG = { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, +} + + +@pytest.mark.parametrize("model_type", [ModelType.SDXL_BASE, ModelType.SDXL_TURBO]) +def test_sdxl_fp4_recipe(model_type): + model = _RecipeBackbone() + config = Quantizer( + QuantizationConfig(format=QuantFormat.FP4), + ModelConfig(model_type=model_type), + logging.getLogger(__name__), + ).get_quant_config(n_steps=1, backbone=model) + + mtq.replace_quant_module(model) + mtq.set_quantizer_by_cfg(model, config["quant_cfg"]) + + for quantizer in (model.linear.input_quantizer, model.linear.weight_quantizer): + assert quantizer.is_enabled + assert quantizer.is_nvfp4_dynamic + assert quantizer.block_sizes[-1] == 16 + for projection in (model.attn.to_q, model.attn.to_k, model.attn.to_v): + assert not projection.input_quantizer.is_enabled + assert not projection.weight_quantizer.is_enabled + for quantizer in (model.conv.input_quantizer, model.conv.weight_quantizer): + assert quantizer.is_enabled + assert quantizer.is_fp8 + + +@pytest.mark.parametrize( + ("format_config", "mha_config", "expected_format", "disable_fp8_mha"), + [ + pytest.param( + _NVFP4_QUANTIZER_CONFIG, + _FP8_QUANTIZER_CONFIG, + QuantFormat.FP4, + False, + id="mixed-fp4", + ), + pytest.param( + _FP8_QUANTIZER_CONFIG, + _FP8_QUANTIZER_CONFIG, + QuantFormat.FP8, + False, + id="fp8", + ), + pytest.param( + {"num_bits": 8}, + {**_FP8_QUANTIZER_CONFIG, "enabled": False}, + QuantFormat.INT8, + True, + id="int8-disabled-fp8", + ), + pytest.param( + {"num_bits": 8}, + {"num_bits": 8}, + QuantFormat.INT8, + True, + id="int8-mha", + ), + ], +) +def test_restored_quantizer_state_drives_format_and_fp8_mha( + monkeypatch, format_config, mha_config, expected_format, disable_fp8_mha +): + backbone = nn.Module() + backbone.quantizer = _quantizer(**format_config) + backbone.attention = nn.Module() + for name in ("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer"): + setattr(backbone.attention, name, _quantizer(**mha_config)) + backbone.attention.bmm2_output_quantizer = lambda output: output + + fp8_sdpa = Mock(return_value=torch.empty(0)) + monkeypatch.setattr(diffusers_plugin.FP8SDPA, "apply", fp8_sdpa) + monkeypatch.setattr(torch.onnx, "is_in_onnx_export", lambda: True) + + assert _infer_restored_quantization_format([("transformer", backbone)]) == expected_format + diffusers_plugin._quantized_sdpa(backbone.attention, *(torch.empty(1) for _ in range(3))) + assert fp8_sdpa.call_args.args[-1] is disable_fp8_mha + + +def test_restore_infers_checkpoint_format_for_export(monkeypatch, tmp_path): + backbone = nn.Module() + backbone.quantizer = _quantizer(**_NVFP4_QUANTIZER_CONFIG) + + pipeline_manager = Mock() + pipeline_manager.create_pipeline.return_value = object() + pipeline_manager.iter_backbones.return_value = [("transformer", backbone)] + export_manager = Mock() + monkeypatch.setattr(_quantize, "PipelineManager", lambda *args: pipeline_manager) + monkeypatch.setattr(_quantize, "ExportManager", lambda *args: export_manager) + monkeypatch.setattr( + sys, + "argv", + [ + "quantize.py", + "--model", + "flux-schnell", + "--restore-from", + str(tmp_path), + "--onnx-dir", + str(tmp_path / "onnx"), + ], + ) + + _quantize.main() + + export_manager.restore_checkpoint.assert_called_once_with() + assert export_manager.export_onnx.call_args.args[-1] == QuantFormat.FP4 + export_manager.export_hf_ckpt.assert_called_once() + + +def test_flux_fp8_export_saves_converted_rope_graph(monkeypatch, tmp_path): + original_model = Mock() + converted_model = Mock() + monkeypatch.setattr( + diffusion_export, + "generate_dummy_kwargs_and_dynamic_axes_and_shapes", + lambda *args: ({}, {}, None), + ) + monkeypatch.setattr(diffusion_export, "onnx_export", lambda *args, **kwargs: None) + monkeypatch.setattr(diffusion_export.onnx, "load", lambda *args, **kwargs: original_model) + convert_rope_weight_type = Mock(return_value=converted_model) + monkeypatch.setattr(diffusion_export, "flux_convert_rope_weight_type", convert_rope_weight_type) + save_onnx = Mock() + monkeypatch.setattr(diffusion_export, "save_onnx", save_onnx) + + diffusion_export.modelopt_export_sd(nn.Module(), tmp_path, "flux-dev", "fp8") + + convert_rope_weight_type.assert_called_once_with(original_model) + save_onnx.assert_called_once_with(converted_model, tmp_path / "model.onnx") diff --git a/tests/unit/onnx/quantization/test_qdq_utils.py b/tests/unit/onnx/quantization/test_qdq_utils.py index 5cdaf9203d0..95944e1f7ab 100644 --- a/tests/unit/onnx/quantization/test_qdq_utils.py +++ b/tests/unit/onnx/quantization/test_qdq_utils.py @@ -42,6 +42,7 @@ replace_zero_scale_with_smallest_nonzero, ) from modelopt.onnx.quantization.quant_utils import pack_float32_to_4bit_cpp_based +from modelopt.onnx.utils import get_opset_version def create_test_model_with_int4_dq_reshape_transpose_matmul(constant_scale: bool = False): @@ -343,8 +344,7 @@ def create_test_model_with_nvfp4_qdq(with_transpose: bool = False): value_info=value_info, ) - model = helper.make_model(graph) - return model + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 20)]) class TestQuantizeWeightsToInt4: @@ -666,6 +666,8 @@ def test_fp4qdq_conversion(self, with_transpose): # Run FP4QDQ to 2DQ conversion converted_model = NVFP4QuantExporter.process_model(model) + assert get_opset_version(converted_model) == 23 + # Verify TRT_FP4QDQ node is removed fp4qdq_nodes = [node for node in converted_model.graph.node if node.op_type == "TRT_FP4QDQ"] assert len(fp4qdq_nodes) == 0 diff --git a/tests/unit/torch/quantization/test_onnx_export_cpu.py b/tests/unit/torch/quantization/test_onnx_export_cpu.py index ce2ef626d63..9b470f7b019 100644 --- a/tests/unit/torch/quantization/test_onnx_export_cpu.py +++ b/tests/unit/torch/quantization/test_onnx_export_cpu.py @@ -39,6 +39,15 @@ from modelopt.torch.quantization.utils import is_quantized_linear +def _export_to_onnx(model, sample_input, **kwargs): + buffer = io.BytesIO() + if "enable_onnx_checker" in inspect.signature(torch.onnx.export).parameters: + kwargs["enable_onnx_checker"] = False + torch.onnx.export(model, sample_input, buffer, dynamo=False, **kwargs) + buffer.seek(0) + return onnx.load_model_from_string(buffer.read()) + + @pytest.mark.parametrize("model_cls", TEST_MODELS) @pytest.mark.parametrize( ("num_bits", "per_channel_quantization", "constant_folding"), @@ -59,6 +68,41 @@ def test_onnx_export_cpu(model_cls, num_bits, per_channel_quantization, constant ) +def test_fp8_conv_export_preserves_custom_qdq_and_kernel_shape(): + model = torch.nn.Conv2d(3, 4, 3, bias=False).eval() + sample_input = torch.randn(1, 3, 8, 8) + model = mtq.quantize( + model, + mtq.FP8_DEFAULT_CFG, + forward_loop=lambda quantized_model: quantized_model(sample_input), + ) + + exported_model = _export_to_onnx(model, sample_input, opset_version=20) + producers = {output: node for node in exported_model.graph.node for output in node.output} + conv = next(node for node in exported_model.graph.node if node.op_type == "Conv") + + for conv_input in conv.input[:2]: + dequantize = producers[conv_input] + quantize = producers[dequantize.input[0]] + assert dequantize.op_type == "TRT_FP8DequantizeLinear" + assert quantize.op_type == "TRT_FP8QuantizeLinear" + + value_info = {value.name: value for value in exported_model.graph.value_info} + weight_dequantize = producers[conv.input[1]] + weight_quantize = producers[weight_dequantize.input[0]] + for value_name in (*weight_quantize.output, *weight_dequantize.output): + shape = [ + dimension.dim_value for dimension in value_info[value_name].type.tensor_type.shape.dim + ] + assert shape == [4, 3, 3, 3] + + kernel_shape = next( + attribute for attribute in conv.attribute if attribute.name == "kernel_shape" + ) + assert list(kernel_shape.ints) == [3, 3] + onnx.checker.check_model(exported_model) + + def test_nvfp4_exported_onnx_is_topologically_sorted(monkeypatch): def forward_loop(model): model(sample_input) @@ -78,26 +122,14 @@ def cpu_dynamic_block_quantize(inputs, *args): module.input_quantizer.disable() module.weight_quantizer._onnx_quantizer_type = "static" - buffer = io.BytesIO() - if "enable_onnx_checker" in inspect.signature(torch.onnx.export).parameters: - kwargs = {"enable_onnx_checker": False} - else: - kwargs = {} - - torch.onnx.export( + exported_model = _export_to_onnx( model, sample_input, - buffer, input_names=["input"], output_names=["output"], export_params=True, opset_version=21, - dynamo=False, - **kwargs, ) - - buffer.seek(0) - exported_model = onnx.load_model_from_string(buffer.read()) assert any(node.op_type == "TRT_FP4QDQ" for node in exported_model.graph.node) converted_model = NVFP4QuantExporter.process_model(exported_model) From a48fd620a88f4a4f511c1d4991f3bec624f96b5d Mon Sep 17 00:00:00 2001 From: Jenny Chen Date: Fri, 18 Sep 2026 14:43:16 -0400 Subject: [PATCH 22/28] Move Nemotron Nano offline KD example (#2471) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix ### Usage ```python # Add a code snippet demonstrating how to use this ``` ### Testing ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ / ❌ / N/A - 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?: ✅ / ❌ / N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ / ❌ / N/A - Did you get Claude approval on this PR?: ✅ / ❌ / N/A ### Additional Information ## Summary by CodeRabbit * **Documentation** * Updated the NVIDIA Nemotron offline usage example to use the BF16 source-directory pipeline path. * Pipeline configuration remains unchanged. Signed-off-by: Jennifer Chen Signed-off-by: Chad Voegele --- .../offline_kd_qad.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) rename tools/launcher/examples/nvidia/{NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4 => NVIDIA-Nemotron-3-Nano-30B-A3B-BF16}/offline_kd_qad.yaml (99%) diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/offline_kd_qad.yaml similarity index 99% rename from tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml rename to tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/offline_kd_qad.yaml index 29cb5280ab8..741235d253f 100644 --- a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/offline_kd_qad.yaml @@ -34,7 +34,7 @@ # # Usage: # source .env-slurm && cd tools/launcher -# uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4/offline_kd_qad.yaml --yes +# uv run launch.py --yaml examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/offline_kd_qad.yaml --yes job_name: NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4-QAD-OfflineKD pipeline: From 415e8266e1aab4dcb7140bc0bc783df443a2006a Mon Sep 17 00:00:00 2001 From: didi Date: Fri, 18 Sep 2026 22:32:06 +0300 Subject: [PATCH 23/28] docs: clarify canonical pruning documentation source (#1871) (#2469) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #1871 The pruning documentation is split between `docs/source/guides/3_pruning.rst` and `examples/pruning/README.md`, causing confusion about which is authoritative. The examples/pruning/README.md is the comprehensive, up-to-date reference covering Minitron, Puzzletron, FastNAS, support matrix, guidelines, and distillation hyperparameters. This PR adds a note to the RST guide making clear: - The README is canonical for Minitron and Puzzletron (LLM/VLM pruning) - The guide covers FastNAS for Computer Vision models Signed-off-by: Diya ## Summary by CodeRabbit - **Documentation** - Updated the pruning guide’s introductory content for clearer separation of general guidance and the related Minitron/Puzzletron note. - Clarified that the guide focuses on FastNAS pruning for computer vision models. - Added references to the Pruning README for Minitron and Puzzletron API examples, support information, guidelines, and distillation hyperparameters. --------- Signed-off-by: didi Signed-off-by: Chad Voegele --- docs/source/guides/3_pruning.rst | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/docs/source/guides/3_pruning.rst b/docs/source/guides/3_pruning.rst index a65890625ba..779e219f247 100644 --- a/docs/source/guides/3_pruning.rst +++ b/docs/source/guides/3_pruning.rst @@ -7,6 +7,13 @@ Pruning Checkout `Megatron-Bridge Minitron Pruning & Distillation `_ and `ResNet20 on CIFAR-10 Notebook `_ for an end-to-end example of pruning. +.. note:: + + For Minitron (LLM/VLM pruning via Megatron-Bridge/Megatron-LM) and Puzzletron pruning, + the canonical reference is the + `Pruning README `_ + which contains up-to-date API examples, support matrix, guidelines, and distillation + hyperparameters. This guide covers FastNAS pruning for Computer Vision models. ModelOpt provides three main pruning methods (aka ``mode``) - Minitron, Puzzletron, and FastNAS - via a unified API :meth:`mtp.prune `. Given a model, From 39a61af003528a97a030ffde08c02fc7b5b8af73 Mon Sep 17 00:00:00 2001 From: Zhiyu Date: Fri, 18 Sep 2026 18:21:07 -0700 Subject: [PATCH 24/28] Fix hf_ptq.py discarding completed PTQ run on sanity-generate() failure (#2480) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix `post_quantize()` in `examples/hf_ptq/hf_ptq.py` ran the optional post-quantization sanity-check `full_model.generate()` unguarded, directly before `export_quantized()`. Any exception raised there aborted the whole run and discarded a completed calibration without exporting a checkpoint. Root cause (traced from [NVBug 6752977](https://nvbugspro.nvidia.com/bug/6752977), DGX Spark GB10 / DeepSeek-R1-Distill-Llama-8B / NVFP4): `get_model()` loads with `device_map="auto"`, relying on `accelerate`'s `infer_auto_device_map`/`get_max_memory()` to decide GPU vs. CPU placement. On DGX Spark's unified-memory single-GPU host, that memory probe under-reports GPU capacity, so part of even an 8B model can land on CPU — and the existing fallback shrinks the GPU budget further (`* gpu_mem_percentage`), compounding it. Calibration survives this because it never invokes the real fake-quant kernel, but the post-PTQ sanity `generate()` does, and NVFP4's dynamic-block-quantize op (`modelopt/torch/quantization/tensor_quant.py`) hard-asserts `amax.is_cuda` with no CPU fallback, so any CPU-offloaded layer crashes there — after ~5.8 hours of calibration, before export. This PR does not attempt to fix the underlying `device_map`/memory-probing behavior (unverified without the actual hardware/logs, which weren't reachable from this environment). Instead it makes the failure mode safe: a failure in the optional sanity check now only skips that check and warns, and export always proceeds, regardless of why `generate()` failed. No new API. Behavior change only: `examples/hf_ptq/hf_ptq.py` now completes export even if the post-quantization sanity `generate()` call raises. - Added `tests/examples/hf_ptq/test_hf_ptq_args.py::test_post_quantize_export_survives_a_failed_sanity_generate`, which drives `post_quantize()` with a `full_model.generate()` that raises and asserts `export_quantized()` still runs. - Ran `pytest tests/examples/hf_ptq/test_hf_ptq_args.py` (48 passed). - Ran `pre-commit` on the changed files (`ruff-format` reformatted line wrapping only). - Is this change backward compatible?: ✅ - 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?: ❌ Fixes NVBug 6752977. Linked JIRA: OMNIML-5932. 🤖 Generated with [Claude Code](https://claude.com/claude-code) - **Bug Fixes** - Quantized checkpoint export now continues when the optional post-quantization generation check fails. - A warning is shown when the generation check cannot complete. Signed-off-by: Zhiyu Cheng Co-authored-by: Claude Sonnet 5 Signed-off-by: Chad Voegele --- CHANGELOG.rst | 1 + examples/hf_ptq/hf_ptq.py | 17 ++++++--- tests/examples/hf_ptq/test_hf_ptq_args.py | 42 +++++++++++++++++++++++ 3 files changed, 55 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4bdc9ae0564..2cf5401c331 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -58,6 +58,7 @@ Changelog - Fix shared ONNX export metadata and Diffusers attention policy: every ``NVFP4QuantExporter`` post-process now upgrades the default-domain opset to at least 23, all FP8 custom-op exports re-run ONNX shape/type inference after setting output metadata, and quantized SDPA derives FP8 MHA enablement from the live Q/K/V quantizers instead of honoring a caller-set ``_disable_fp8_mha`` attribute. - Fix ONNX FP16 conversion failing to preserve public output types when type inference changes a graph output declaration before output casts are inserted. +- Fix ``examples/hf_ptq/hf_ptq.py`` discarding a completed PTQ run (no checkpoint exported) when the optional post-quantization sanity-check ``generate()`` call raised, for example because ``device_map="auto"`` placed part of the model on CPU. That failure is now caught and only skips the sanity check; export proceeds regardless. - Fail fast on non-finite AutoQuantize output gradients with an actionable error before accumulating sensitivity scores, without changing attention backend settings. - Fix ONNX INT8 entropy calibration failing or producing invalid quantization parameters for FP16 activations. - Speed up ``mtq.quantize`` on FSDP2-sharded fused-MoE models. Promoting static-block weight quantizers gathered each expert's slice of the fused weight across ranks even though only quantizer state is read, adding a collective per expert to calibration. diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index a56a62b54b5..5cd3fc82fdc 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -1020,11 +1020,18 @@ def post_quantize( pass elif model_type != "llama4" and not is_nemotron_vl_model: # Our fake quantizer may not be fully compatible with torch.compile. - generated_ids_after_ptq = full_model.generate( - preview_input_ids, - attention_mask=preview_attention_mask, - max_new_tokens=100, - ) + # This is a best-effort sanity check: e.g. a `device_map="auto"` load that offloads + # part of the model to CPU (seen on unified-memory single-GPU hosts) can make a + # quantized layer run on CPU, which some kernels (e.g. NVFP4 dynamic block + # quantization) don't support. Don't let that discard the completed calibration. + try: + generated_ids_after_ptq = full_model.generate( + preview_input_ids, + attention_mask=preview_attention_mask, + max_new_tokens=100, + ) + except Exception as e: + warnings.warn(f"Post-quantization generation sanity check failed, skipping it: {e}") elif is_nemotron_vl_model and tokenizer is not None: generated_ids_after_ptq = run_nemotron_vl_preview( full_model, diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 6a2c36e4a7c..90bf671abd6 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -20,6 +20,7 @@ from types import SimpleNamespace import pytest +import torch import yaml from modelopt.recipe import load_recipe @@ -402,3 +403,44 @@ def test_mlflow_checkpoint_tag_is_absolute(monkeypatch, example_utils): ) assert Path(example_utils._mlflow_run_tags(args)["checkpoint_path"]).is_absolute() + + +# --- post-quantization sanity-check generate() must not block export ---------------------------- + + +def test_post_quantize_export_survives_a_failed_sanity_generate(monkeypatch): + """A failed optional sanity generation must not prevent checkpoint export.""" + hf_ptq = _import_hf_ptq(monkeypatch) + + full_model = SimpleNamespace( + generate=lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom")) + ) + export_calls = [] + monkeypatch.setattr( + hf_ptq, + "export_quantized", + lambda *a, **k: export_calls.append((a, k)), + ) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + + args = SimpleNamespace(specdec_offline_dataset=None, verbose=False) + + with pytest.warns(UserWarning, match="Post-quantization generation sanity check failed"): + hf_ptq.post_quantize( + args=args, + full_model=full_model, + language_model=full_model, + model_type="llama", + tokenizer=None, + processor=None, + preview_input_ids=torch.zeros(1, 4, dtype=torch.long), + preview_attention_mask=None, + generated_ids_before_ptq=torch.zeros(1, 4, dtype=torch.long), + is_nemotron_vl_model=False, + first_text_speech_dataset=None, + default_padding_side="right", + default_pad_token=None, + calib_dataloader=None, + ) + + assert len(export_calls) == 1 From b9351263d70168dc453df48b105be64cc1db4e24 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Sat, 19 Sep 2026 13:45:21 +0530 Subject: [PATCH 25/28] Fix hybrid stack spec serialization in Megatron-Bridge checkpoints (#2452) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix Hybrid (e.g. Nemotron-H) checkpoints saved by the `examples/megatron_bridge` scripts could not be reloaded or exported to HuggingFace: TypeError: MLPSubmodules.__init__() missing 2 required positional arguments: 'linear_fc1' and 'linear_fc2' **Root cause.** Megatron-LM's YAML writer represents a `functools.partial` via `_partial_representer`, which passes each keyword value through `represent_data`. A dataclass instance has no representer, so it falls through to `_safe_object_representer`, which emits only `{_target_, _call_}` and drops every field. The default hybrid stack spec builds its dense-MLP and MoE layers as exactly such partials, and `set_moe_expert_layout()` stored the *built* `ModuleSpec` on the provider — which is serialized into every checkpoint's `run_config.yaml`. So `MLPSubmodules` / `MoESubmodules` were written with no fields at all. Both export paths hit it: `convert.sh` via `from_auto_config`, and `export_distilled_megatron_to_hf.py` via `export_ckpt → load_megatron_model`. Every hybrid provider is affected, dense or MoE. **Fix.** `set_moe_expert_layout()` stores a named, zero-argument *factory* instead. The provider already calls a callable spec at build time (`_resolve_hybrid_stack_spec`), so model construction is unchanged — only the serialized form differs: ```yaml hybrid_stack_spec: _call_: false _target_: megatron.bridge.models.hybrid.hybrid_provider.transformer_engine_hybrid_stack_spec ``` The grouped-GEMM factory is Megatron-Bridge's own `transformer_engine_hybrid_stack_spec`, so stock tooling (`scripts/conversion/convert.sh`) resolves it without importing ModelOpt. **Known limitation — the two layouts are not symmetric.** The SequentialMLP layout has no bridge-side equivalent (the upstream TE hybrid spec hardcodes `TEGroupedMLP`), so it serializes a ModelOpt target, which `instantiate` only accepts in a process that has imported `mbridge.py` and thereby run `register_allowed_target_prefix`. A SequentialMLP hybrid checkpoint therefore converts through the ModelOpt entrypoints but not through stock `convert.sh`, where it fails on the disallowed prefix instead of on `MLPSubmodules` — no regression, but that path stays broken for this one layout. The reach is narrow: `use_moe_grouped_gemm()` returns True for any architecture with a grouped-expert export rule, NemotronH included, so SequentialMLP requires an explicit `--no_moe_grouped_gemm`. Closing it properly needs an upstream `moe_grouped_gemm`-aware factory in Megatron-Bridge. Both spec builders also move out of `nas/plugins/megatron.py`, which never used them, into a new `utils/plugins/megatron_layer_specs.py` beside the other Megatron-Core-only helpers. Not into `mbridge.py`: that module needs `megatron.bridge`, while `get_te_hybrid_stack_spec` is reached by 16 test files through `tests/_test_utils/torch/megatron/models.py`, which is bridge-free. The underlying defect is upstream in `megatron/training/config/yaml_utils.py`; this only stops ModelOpt from stepping on it, so it is worth a separate Megatron-LM issue. No API change — hybrid checkpoints saved after this fix convert with the existing commands: ```bash torchrun --nproc_per_node 1 examples/megatron_bridge/export_distilled_megatron_to_hf.py \ --student_hf_path \ --megatron_path /checkpoints \ --hf_export_path \ --export_iterations all ``` Verified in `nemo:26.08` (megatron-core 0.19.0) against a 30B-A3B Nemotron-3.5-Lightning pruned+distilled run: - **Round trip, both MoE layouts.** Ran `set_moe_expert_layout` on a real `HybridModelProvider`, dumped it through `dump_dataclass_to_yaml` (the writer used for `run_config.yaml`), reloaded via `instantiate`, resolved. `moe_grouped_gemm=True` → `TELayerNormColumnParallelLinear`/`TERowParallelLinear` + `TEGroupedMLP`; `False` → same MLP + `SequentialMLP`. The field stays callable after `finalize()` and `_resolve_hybrid_stack_spec()`, so a saved config cannot regress. - Applying the equivalent `run_config.yaml` fix to 32 iteration checkpoints: all 32 rebuild the provider (52 layers, hidden 2304, 104 experts) with populated `MLPSubmodules` / `MoESubmodules`. - **End-to-end exports**, 6 iterations, all rc 0, each producing exactly the source model's 5139 tensor keys (0 missing, 0 extra), 9 shards / 41.5 GiB, all weights finite, drift from the base rising monotonically with iteration (lm_head 0.030 → 0.092). Covered `convert.sh` CPU, `convert.sh` GPU (4×GB300, TP=4), and `export_distilled_megatron_to_hf.py`. Same iteration and wrapper: CPU 123 s vs GPU 134 s — GPU is not faster, since with TP=4 each rank still builds 20.9 B of 22.3 B params and the cost is I/O plus CPU-side conversion. - `ruff check` / `ruff format --check` passed on the source files before the module move. **Not yet run:** `tests/gpu_megatron/torch/utils/plugins/test_mbridge.py` (added here) — the GPU allocation expired. It asserts the round-trip property verified manually above, but its `HybridModelProvider(num_layers=2, hidden_size=64, num_attention_heads=4)` construction is unverified. The module move is verified only by reference grep and syntax check, so please also run one test that uses `tests/_test_utils/torch/megatron/models.py`. `pre-commit` was not run either (unavailable in the environment used). - Is this change backward compatible?: ✅ behavior; note `get_te_hybrid_stack_spec` moved module (`modelopt.torch.nas.plugins.megatron` → `modelopt.torch.utils.plugins.megatron_layer_specs`), and a checkpoint from 0.46.1/0.47.0 needs the `run_config.yaml` edit described in the changelog. - 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?: ✅ (added, not yet executed — see Testing) - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ❌ — will run `/claude review` before marking ready. * **New Features** * Hybrid checkpoints now record complete layer specifications in `run_config.yaml`. * Recorded specifications support conversion to Hugging Face format. * Hybrid MoE configurations support grouped-GEMM and sequential-MLP modes. * Configuration-based reconstruction preserves the selected MoE layout. * **Compatibility** * Checkpoints from earlier releases may require manually setting the hybrid layer specification before conversion. * **Tests** * Added coverage confirming hybrid specifications survive configuration serialization and can be recreated successfully. Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) Signed-off-by: Chad Voegele --- CHANGELOG.rst | 2 + modelopt/torch/nas/plugins/megatron.py | 23 +----- modelopt/torch/utils/plugins/__init__.py | 6 +- modelopt/torch/utils/plugins/mbridge.py | 34 ++++++++- .../utils/plugins/megatron_layer_specs.py | 46 ++++++++++++ tests/_test_utils/torch/megatron/models.py | 13 ++-- .../torch/utils/plugins/test_mbridge.py | 74 +++++++++++++++++++ 7 files changed, 166 insertions(+), 32 deletions(-) create mode 100644 modelopt/torch/utils/plugins/megatron_layer_specs.py create mode 100644 tests/gpu_megatron/torch/utils/plugins/test_mbridge.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 2cf5401c331..fc6b74812f1 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -36,6 +36,7 @@ Changelog **Backward Breaking Changes** +- ``get_te_hybrid_stack_spec`` was removed from ``modelopt.torch.nas.plugins.megatron``; it had no use outside tests. Use ``modelopt.torch.utils.plugins.megatron_layer_specs.te_hybrid_stack_spec_sequential_mlp`` for the SequentialMLP layout, or ``megatron.core.models.hybrid.hybrid_layer_specs.hybrid_stack_spec`` for grouped GEMM. - Migrate the FAR3D ONNX PTQ example to the shared evaluator and ModelOpt containers and ``quantize_vovnet.py``. Only the encoder supports INT8 and FP8; decoder calibration, quantization, and related CLI flags are removed, and the decoder remains in its exported mixed FP16/FP32 precision. - Image-text calibration with ``--calib_with_images`` now forwards multimodal batches through the complete VLM for all VLM families, so existing non-Nemotron commands may produce different language-model activation ranges and output scales. Recipe-based VLM PTQ also targets the complete VLM: vision modules stay in high precision by default and are quantized only when a model-specific recipe enables them, so custom recipes must explicitly exclude vision modules when required. - Move the checkpoint-mirror recipe tier from ``huggingface/models///`` to the top-level ``models///``, keyed by each recipe's canonical Hugging Face Hub id — so the Step 3.5 Flash recipe moves to ``models/stepfun-ai/Step-3.5-Flash/ptq/`` and the NVIDIA Nemotron recipes gain the ``NVIDIA-`` prefix (e.g. ``models/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16/ptq/nvfp4-mse``). Update any saved ``--recipe`` paths for these checkpoint recipes accordingly; the per-``model_type`` recipes under ``huggingface/`` are unchanged. @@ -59,6 +60,7 @@ Changelog - Fix shared ONNX export metadata and Diffusers attention policy: every ``NVFP4QuantExporter`` post-process now upgrades the default-domain opset to at least 23, all FP8 custom-op exports re-run ONNX shape/type inference after setting output metadata, and quantized SDPA derives FP8 MHA enablement from the live Q/K/V quantizers instead of honoring a caller-set ``_disable_fp8_mha`` attribute. - Fix ONNX FP16 conversion failing to preserve public output types when type inference changes a graph output declaration before output casts are inserted. - Fix ``examples/hf_ptq/hf_ptq.py`` discarding a completed PTQ run (no checkpoint exported) when the optional post-quantization sanity-check ``generate()`` call raised, for example because ``device_map="auto"`` placed part of the model on CPU. That failure is now caught and only skips the sanity check; export proceeds regardless. +- Hybrid (e.g. Nemotron-H) checkpoints saved by the ``examples/megatron_bridge`` scripts now record their layer spec in ``run_config.yaml`` in a form that reloads, so they can be converted to HuggingFace; a checkpoint saved by an earlier release still needs its ``model.hybrid_stack_spec`` block replaced by hand. - Fail fast on non-finite AutoQuantize output gradients with an actionable error before accumulating sensitivity scores, without changing attention backend settings. - Fix ONNX INT8 entropy calibration failing or producing invalid quantization parameters for FP16 activations. - Speed up ``mtq.quantize`` on FSDP2-sharded fused-MoE models. Promoting static-block weight quantizers gathered each expert's slice of the fused weight across ranks even though only quantizer state is read, adding a collective per expert to calibration. diff --git a/modelopt/torch/nas/plugins/megatron.py b/modelopt/torch/nas/plugins/megatron.py index 2376422a998..c2d31b6f3de 100644 --- a/modelopt/torch/nas/plugins/megatron.py +++ b/modelopt/torch/nas/plugins/megatron.py @@ -15,7 +15,6 @@ """Plugin to add NAS/Pruning support for megatron-core Language models like GPT and Mamba.""" -import copy import types from abc import ABC from collections.abc import Callable, Sequence @@ -35,10 +34,6 @@ ) from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding from megatron.core.models.gpt import GPTModel -from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec -from megatron.core.models.hybrid.hybrid_layer_specs import ( - hybrid_stack_spec as _te_hybrid_stack_spec, -) from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.parallel_state import is_pipeline_first_stage, is_pipeline_last_stage from megatron.core.ssm.gated_delta_net import GatedDeltaNet @@ -56,7 +51,6 @@ from megatron.core.transformer.moe.router import TopKRouter from megatron.core.transformer.moe.shared_experts import SharedExpertMLP from megatron.core.transformer.multi_latent_attention import MLASelfAttention -from megatron.core.transformer.spec_utils import ModuleSpec from megatron.core.transformer.transformer_layer import TransformerLayer from modelopt.torch.nas.modules import DynamicModuleList @@ -91,21 +85,8 @@ # Attention module types that _DynamicTransformerLayer converts. _ATTENTION_TYPES: tuple[type, ...] = (SelfAttention, MLASelfAttention, GatedDeltaNet) -__all__ = ["get_te_hybrid_stack_spec"] - - -def get_te_hybrid_stack_spec(moe_grouped_gemm: bool = False) -> ModuleSpec: - """Return the TE Hybrid stack spec.""" - if moe_grouped_gemm: - return _te_hybrid_stack_spec - - # The upstream TE hybrid stack spec hardcodes TEGroupedMLP for MoE. - # Replace it with SequentialMLP (TE linear layers, no grouped gemm dependency). - te_hybrid_stack_spec = copy.deepcopy(_te_hybrid_stack_spec) - te_hybrid_stack_spec.submodules.moe_layer.submodules.mlp = get_moe_module_spec( - use_te=True, num_experts=8, moe_grouped_gemm=False - ) - return te_hybrid_stack_spec +# This module only registers DynamicModules; it exports no public API. +__all__ = [] # Local Parallel Linear DynamicModules ########################################################################## diff --git a/modelopt/torch/utils/plugins/__init__.py b/modelopt/torch/utils/plugins/__init__.py index da40fe9e565..ff251eaf133 100644 --- a/modelopt/torch/utils/plugins/__init__.py +++ b/modelopt/torch/utils/plugins/__init__.py @@ -23,6 +23,9 @@ with import_plugin("megatron_generate"): from .megatron_generate import * +with import_plugin("megatron_layer_specs"): + from .megatron_layer_specs import * + with import_plugin("megatron_mmlu"): from .megatron_mmlu import * @@ -33,6 +36,7 @@ from .prepare_megatron_data_blend import * # NOTE: Dont pre-import megatron bridge plugin here to avoid circular dependency issues. -# We dont register anything so this isnt a problem. +# It registers an instantiate allowlist prefix on import, which only the ModelOpt entrypoints +# that import it need, so leaving it out here is still fine. # with import_plugin("megatron bridge"): # from .mbridge import * diff --git a/modelopt/torch/utils/plugins/mbridge.py b/modelopt/torch/utils/plugins/mbridge.py index 05713c7b282..7b85fb9d4fb 100644 --- a/modelopt/torch/utils/plugins/mbridge.py +++ b/modelopt/torch/utils/plugins/mbridge.py @@ -22,13 +22,17 @@ from megatron.bridge import AutoBridge from megatron.bridge.models.gpt_provider import GPTModelProvider from megatron.bridge.models.hf_pretrained.utils import is_safe_repo -from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider +from megatron.bridge.models.hybrid.hybrid_provider import ( + HybridModelProvider, + transformer_engine_hybrid_stack_spec, +) from megatron.bridge.training.checkpointing import _load_model_weights_from_checkpoint from megatron.bridge.training.post_training.checkpointing import ( _get_modelopt_checkpoint_path, has_modelopt_state, load_modelopt_state, ) +from megatron.bridge.utils.instantiate_utils import register_allowed_target_prefix from megatron.core.models.gpt import GPTModel from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.transformer.module import MegatronModule @@ -37,8 +41,16 @@ from transformers import AutoConfig, AutoTokenizer from modelopt.torch.export.plugins.mcore_common import all_mcore_hf_export_mapping -from modelopt.torch.nas.plugins.megatron import get_te_hybrid_stack_spec from modelopt.torch.utils import print_rank_0, warn_rank_0 +from modelopt.torch.utils.plugins.megatron_layer_specs import te_hybrid_stack_spec_sequential_mlp + +# ``set_moe_expert_layout`` records ``te_hybrid_stack_spec_sequential_mlp`` in a SequentialMLP +# checkpoint's ``run_config.yaml``, and rebuilding that config resolves the target against this +# allowlist. Only a process that imports this module registers it, so such a checkpoint must be +# converted through a ModelOpt entrypoint, not stock ``scripts/conversion/convert.sh``. The prefix +# covers all of ``modelopt`` rather than one module, which assumes a ``run_config.yaml`` is trusted +# input -- it comes from a checkpoint the caller is already choosing to load. +register_allowed_target_prefix("modelopt.") __all__ = [ "get_language_model", @@ -114,10 +126,24 @@ def set_moe_expert_layout(provider, moe_grouped_gemm: bool) -> None: Set ``moe_grouped_gemm`` on the provider (the bridge's native, possibly custom/hybrid spec reads it at build time) rather than replacing the whole layer spec -- overwriting it would drop custom layers (e.g. Qwen3.5's GatedDeltaNet or Gemma3's custom spec). A hybrid provider - additionally needs its stack spec rebuilt, since the native one pins ``TEGroupedMLP``. + additionally has its stack spec set, since the native one pins ``TEGroupedMLP``: the bridge's + own factory for grouped GEMM, a ModelOpt one that swaps in SequentialMLP otherwise. + + Assign a *factory function*, never a built ``ModuleSpec``: the provider is serialized into + every checkpoint's ``run_config.yaml``, and Megatron-LM's YAML writer drops the fields of a + dataclass nested inside a ``functools.partial`` keyword, which is how a stack spec holds + ``MLPSubmodules`` / ``MoESubmodules``. Such a checkpoint cannot be reloaded or exported. The + provider calls the factory at build time, so behavior is unchanged. """ if isinstance(provider, HybridModelProvider): - provider.hybrid_stack_spec = get_te_hybrid_stack_spec(moe_grouped_gemm=moe_grouped_gemm) + # The grouped-GEMM factory is Megatron-Bridge's, and returns Megatron-Core's + # ``hybrid_stack_spec`` unchanged -- the layer composition is identical either way. It is + # named from the bridge so stock tooling resolves the target without importing ModelOpt. + provider.hybrid_stack_spec = ( + transformer_engine_hybrid_stack_spec + if moe_grouped_gemm + else te_hybrid_stack_spec_sequential_mlp + ) provider.moe_grouped_gemm = moe_grouped_gemm elif (provider.num_moe_experts or 0) > 0: provider.moe_grouped_gemm = moe_grouped_gemm diff --git a/modelopt/torch/utils/plugins/megatron_layer_specs.py b/modelopt/torch/utils/plugins/megatron_layer_specs.py new file mode 100644 index 00000000000..d9b08b7b134 --- /dev/null +++ b/modelopt/torch/utils/plugins/megatron_layer_specs.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Megatron-Core layer specs used to build models for ModelOpt workflows.""" + +import copy + +from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec +from megatron.core.models.hybrid.hybrid_layer_specs import ( + hybrid_stack_spec as _te_hybrid_stack_spec, +) +from megatron.core.transformer.spec_utils import ModuleSpec + +__all__ = ["te_hybrid_stack_spec_sequential_mlp"] + + +def te_hybrid_stack_spec_sequential_mlp() -> ModuleSpec: + """Return the TE Hybrid stack spec with SequentialMLP MoE experts. + + Named and zero-argument so a provider can store this function instead of the ModuleSpec it + builds; see ``set_moe_expert_layout`` for why a built spec cannot be serialized. + + Its module path and name are written into ``run_config.yaml`` as a ``_target_``, so moving or + renaming it breaks every SequentialMLP hybrid checkpoint already saved. + """ + # The upstream TE hybrid stack spec hardcodes TEGroupedMLP for MoE. + # Replace it with SequentialMLP (TE linear layers, no grouped gemm dependency). + # num_experts only has to be non-zero to select the MoE branch; the real count comes from the + # model config at build time. + te_hybrid_stack_spec = copy.deepcopy(_te_hybrid_stack_spec) + te_hybrid_stack_spec.submodules.moe_layer.submodules.mlp = get_moe_module_spec( + use_te=True, num_experts=8, moe_grouped_gemm=False + ) + return te_hybrid_stack_spec diff --git a/tests/_test_utils/torch/megatron/models.py b/tests/_test_utils/torch/megatron/models.py index b2bd2e90921..efdc030a9dc 100644 --- a/tests/_test_utils/torch/megatron/models.py +++ b/tests/_test_utils/torch/megatron/models.py @@ -25,6 +25,7 @@ get_gpt_layer_with_transformer_engine_spec, get_gpt_mtp_block_spec, ) +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec as te_hybrid_stack_spec from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.parallel_state import ( get_pipeline_model_parallel_rank, @@ -37,7 +38,7 @@ from megatron.core.transformer.transformer_config import MLATransformerConfig, TransformerConfig from modelopt.torch.export.unified_export_megatron import import_mcore_gpt_from_hf -from modelopt.torch.nas.plugins.megatron import get_te_hybrid_stack_spec +from modelopt.torch.utils.plugins.megatron_layer_specs import te_hybrid_stack_spec_sequential_mlp try: from megatron.core.extensions.transformer_engine import TENorm @@ -444,11 +445,11 @@ def get_mcore_hybrid_model( "share_embeddings_and_output_weights": False, "position_embedding_type": "none", } - spec = ( - get_te_hybrid_stack_spec(moe_grouped_gemm) - if transformer_impl == "transformer_engine" - else get_hybrid_stack_modelopt_spec(remap_te_layernorm=True) - ) + if transformer_impl == "transformer_engine": + # The upstream TE hybrid stack spec hardcodes TEGroupedMLP for MoE. + spec = te_hybrid_stack_spec if moe_grouped_gemm else te_hybrid_stack_spec_sequential_mlp() + else: + spec = get_hybrid_stack_modelopt_spec(remap_te_layernorm=True) model = HybridModel( hybrid_stack_spec=spec, hybrid_layer_pattern=hybrid_layer_pattern, **common_kwargs ) diff --git a/tests/gpu_megatron/torch/utils/plugins/test_mbridge.py b/tests/gpu_megatron/torch/utils/plugins/test_mbridge.py new file mode 100644 index 00000000000..3b5e7ea78c8 --- /dev/null +++ b/tests/gpu_megatron/torch/utils/plugins/test_mbridge.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +import pytest +import yaml +from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider +from megatron.bridge.utils.instantiate_utils import instantiate +from megatron.bridge.utils.yaml_utils import dump_dataclass_to_yaml + +from modelopt.torch.utils.plugins.mbridge import set_moe_expert_layout + + +def _round_trip(value): + """Serialize through the writer used for run_config.yaml, then reload.""" + node = yaml.safe_load(dump_dataclass_to_yaml({"spec": value}))["spec"] + return node["_target_"], instantiate(node) + + +@pytest.mark.parametrize( + ("moe_grouped_gemm", "expected_experts", "expected_target"), + [ + ( + True, + "TEGroupedMLP", + "megatron.bridge.models.hybrid.hybrid_provider.transformer_engine_hybrid_stack_spec", + ), + ( + False, + "SequentialMLP", + "modelopt.torch.utils.plugins.megatron_layer_specs.te_hybrid_stack_spec_sequential_mlp", + ), + ], +) +def test_set_moe_expert_layout_survives_run_config_round_trip( + moe_grouped_gemm, expected_experts, expected_target +): + """A provider's stack spec must still build real submodules after a run_config round trip. + + A built ``ModuleSpec`` loses its ``MLPSubmodules`` / ``MoESubmodules`` when written to + ``run_config.yaml``, so ``set_moe_expert_layout`` stores a factory function instead. + """ + provider = HybridModelProvider(num_layers=2, hidden_size=64, num_attention_heads=4) + set_moe_expert_layout(provider, moe_grouped_gemm=moe_grouped_gemm) + assert provider.moe_grouped_gemm == moe_grouped_gemm + + assert callable(provider.hybrid_stack_spec) + + target, factory = _round_trip(provider.hybrid_stack_spec) + # The target is an on-disk contract: renaming or moving the factory breaks saved checkpoints. + assert target == expected_target + + provider.hybrid_stack_spec = factory + spec = provider._resolve_hybrid_stack_spec() + + mlp = spec.submodules.mlp_layer.submodules.mlp.keywords["submodules"] + assert mlp.linear_fc1 is not None + assert mlp.linear_fc2 is not None + + moe = spec.submodules.moe_layer.submodules.mlp.keywords["submodules"] + assert moe.experts is not None + # Experts are built through a partial for the grouped-GEMM layout. + assert getattr(moe.experts, "func", moe.experts).__name__ == expected_experts From 042e1ff820eb5313906e438d37191b1c3c9a7231 Mon Sep 17 00:00:00 2001 From: yueshen2016 <39203804+yueshen2016@users.noreply.github.com> Date: Mon, 21 Sep 2026 10:47:34 -0700 Subject: [PATCH 26/28] Fix HF export crash when a dynamic-block quantizer has zero amax (#2438) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix `TensorQuantizer.export_amax()` early-returns `self.amax` unsanitized for dynamic-block quantizers, while the static path immediately below it has always substituted `maxbound` for zero/NaN entries. The `nvfp4` numerics unit sets `type: dynamic`, so a recipe that applies it to an *activation* quantizer — e.g. `general/ptq/nvfp4_mlp_only-kv_fp8_cast`, which targets `*mlp*input_quantizer` — feeds a raw `0.0` into `NVFP4QTensor.get_activation_scaling_factor`, whose assert aborts the entire export: ``` AssertionError: Failed to export module 'model.language_model.layers.37.mlp.gate_proj' (type=QuantLinear): activation scaling factor 0.0 not positive. ``` Calibration leaves `amax` at 0 whenever a layer — or an unrouted MoE expert — saw only zeros, so one dead layer costs the whole run at the final export step. This factors the substitution into `_sanitize_export_amax()` and calls it from both branches. Two details beyond de-duplication: - **Branch-free, so it survives a meta `amax`.** `torch.where` + `nan_to_num` both have meta kernels; `bool()` on a meta tensor raises. The layerwise and streaming export flows carry meta `amax` — `validate_attr` short-circuits on `is_meta` for exactly that reason — so only the warning is gated on a materialized tensor. - **No longer mutates calibrated state.** The old in-place `amax[amax == 0] = ...` wrote through a view of `self._amax`; `torch.where` returns a fresh tensor, so that hazard disappears. - **Warns, with a count.** The fix turns a loud failure into a silent one, and a zero amax means calibration never activated that layer — worth surfacing rather than papering over. The message reports how many entries were substituted, since per-location dedup otherwise collapses many dead experts into one uninformative message. A healthy model emits none. Scope: only the activation path is data-dependent and reachable this way. Weight-side `_amax` uses are left alone, since a weight amax of 0 would require an all-zero weight matrix. **Knowingly left as follow-up:** `export/quant_utils.py::get_scaling_factor` discards the sanitized `amax` when `num_bits == (2, 1)` and recomputes via `get_weights_scaling_factor_2_from_quantizer`, which reads `weight_quantizer._amax` raw — so a dynamic-NVFP4 *input* quantizer on a module whose *weight* quantizer is a different format (or disabled) can still trip `assert torch.all(scaling_factor > 0)`. Format dispatch is weight-driven, so the reported recipe does not reach that branch; fixing it properly changes a signature shared with the weight-side callers and is out of scope here. Not a regression. The dynamic early return, the `type: dynamic` numerics unit, and the recipe that combines them all ship in released 0.46.0 / 0.46.1. No new or changed API. Exports that previously aborted now complete and warn: ```python mtq.quantize(model, quant_cfg, forward_loop) export_hf_checkpoint(model, export_dir=out) # before: AssertionError; now: exports + UserWarning ``` - New `test_amax_export_unusable_amax`, parametrized over zero and NaN, covering the dynamic-NVFP4 and static per-tensor configs; asserts the exported scale is positive and that export leaves the calibrated `amax` untouched. Plus `test_amax_export_meta_amax`, pinning that a meta `amax` survives export rather than raising. Both run on CPU and CUDA via the shared tester. - `tests/unit/torch/quantization/test_tensor_quantizer_cpu.py` — 40 passed. `tests/gpu/torch/quantization/test_tensor_quantizer_cuda.py` — 40 passed (GB300). - End-to-end repro on GB300, small Llama with one MLP fed all-zero activations under `general/ptq/nvfp4_mlp_only-kv_fp8_cast`: dead layer `export_amax()` `0.0` → `6.0`, live layer unchanged at `3.921875`, and `export_hf_checkpoint` goes from the `AssertionError` above to writing `model.safetensors`. - Full `examples/hf_ptq/hf_ptq.py` with the reported recipe and flags on a healthy model (Qwen3-0.6B): exits 0 and writes the checkpoint, confirming the normal path is unaffected. - `pre-commit run` clean on all changed files. - Is this change backward compatible?: ✅ - 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?: ✅ — `/claude review` run; its one IMPORTANT finding (meta-tensor regression) and both SUGGESTIONs addressed or answered in 251f2e3d Fixes NVBug 6768300, reported against 0.47.0rc1 on GB200. The reporter also notes it passed on 0.47.0rc0; that is not explained by code — `git diff 0.47.0rc0..0.47.0rc1` touches `export/quant_utils.py` only in `get_kv_cache_scaling_factor` (new `clamp_fp8_scales` argument whose default preserves the old behaviour) and the INT4-AWQ packing path, neither of which is on the dense-HF NVFP4 activation-scale path. Whether `amax` lands on exactly 0 is calibration/model-state dependent, which is what makes it look version-flaky. Worth flagging separately: in the reported log the **pre-PTQ** sample output is already gibberish, so that BF16 checkpoint looks broken independently of quantization. This change stops the crash, but such a run will now export a valid-but-garbage checkpoint — the new warning is the signal to investigate. Suggest the `cherry-pick-0.47.0` label so this lands in the ongoing release. 🤖 Generated with [Claude Code](https://claude.com/claude-code) * **Bug Fixes** * Fixed Hugging Face checkpoint export when dynamic-block quantizers have zero or invalid calibration scales. * Exports now use a positive fallback scale and issue a warning instead of failing when applicable. * Export operations no longer modify the original calibrated quantizer state. * Meta-device exports remain non-erroring and preserve device placement. * **Tests** * Added coverage for zero- and invalid-scale exports across dynamic and static quantization modes. Signed-off-by: Yue Co-authored-by: Claude Opus 5 (1M context) Signed-off-by: Chad Voegele --- CHANGELOG.rst | 1 + .../nn/modules/tensor_quantizer.py | 33 ++++++++++++-- .../quantization/tensor_quantizer_common.py | 43 +++++++++++++++++++ 3 files changed, 74 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index fc6b74812f1..16c48f5bb27 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -63,6 +63,7 @@ Changelog - Hybrid (e.g. Nemotron-H) checkpoints saved by the ``examples/megatron_bridge`` scripts now record their layer spec in ``run_config.yaml`` in a form that reloads, so they can be converted to HuggingFace; a checkpoint saved by an earlier release still needs its ``model.hybrid_stack_spec`` block replaced by hand. - Fail fast on non-finite AutoQuantize output gradients with an actionable error before accumulating sensitivity scores, without changing attention backend settings. - Fix ONNX INT8 entropy calibration failing or producing invalid quantization parameters for FP16 activations. +- Fix HuggingFace checkpoint export failing with ``activation scaling factor 0.0 not positive`` when a dynamic-block quantizer (such as an NVFP4 input quantizer) ends calibration with an amax of zero because the calibration data never activated that layer or expert. Such a quantizer now exports a positive fallback scale and warns instead of crashing, matching what static quantizers already did; if you see the warning, check whether the layer is expected to be inactive and consider a larger calibration size. - Speed up ``mtq.quantize`` on FSDP2-sharded fused-MoE models. Promoting static-block weight quantizers gathered each expert's slice of the fused weight across ranks even though only quantizer state is read, adding a collective per expert to calibration. - Fix ONNX AutoCast failing on models with external initializers larger than 2 GiB. - Avoid querying CUDA/Blackwell capability when ``NVFP4QTensor.quantize`` uses its CPU path or has the optional TensorRT-LLM fast path disabled. diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index 18b97ac2774..a2349c6eb3f 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -1084,10 +1084,38 @@ def _check_per_channel_block_sizes(block_sizes): # remove block_sizes self._block_sizes = None + def _sanitize_export_amax(self, amax: torch.Tensor) -> torch.Tensor: + """Replace zero/NaN amax entries with ``maxbound`` so exported scales stay positive. + + A zero amax means calibration never activated the layer; downstream exporters divide by + the exported amax, so a zero would fail export or produce inf at inference. Kept + branch-free because export flows may pass a meta ``amax``; only the warning reads values. + """ + sanitized = torch.nan_to_num( + torch.where(amax == 0, torch.full_like(amax, self.maxbound), amax), + nan=self.maxbound, + ) + + if not amax.is_meta: + num_invalid = int((torch.isnan(amax) | (amax == 0)).sum()) + if num_invalid: + warnings.warn( + f"{num_invalid}/{amax.numel()} amax entries of this " + f"{type(self).__name__} are zero or NaN at export time, which means " + "calibration never activated the corresponding layer or expert (or saw NaN " + "activations). Substituting maxbound so the exported scaling factor stays " + "positive. Consider increasing the calibration size if the layer is expected " + "to be active.", + stacklevel=3, + ) + return sanitized + def export_amax(self) -> torch.Tensor | None: """Export correctly formatted/shaped amax.""" if self.block_sizes is not None and self.block_sizes.get("type", None) == "dynamic": - return self.amax + # Dynamic block quantizers keep a per-tensor amax (the NVFP4 second-level scale) that + # needs no reshaping, but it still has to be positive for the exporters. + return None if self.amax is None else self._sanitize_export_amax(self.amax) if self.amax is None: return None @@ -1096,8 +1124,7 @@ def export_amax(self) -> torch.Tensor | None: amax = self.amax else: amax = self.amax.reshape(self._amax_shape_for_export) - amax[amax == 0] = self.maxbound - amax = torch.nan_to_num(amax, nan=self.maxbound) + amax = self._sanitize_export_amax(amax) clamp_min, clamp_max = torch.finfo(amax.dtype).tiny, torch.finfo(amax.dtype).max amax = amax.clamp(min=clamp_min, max=clamp_max) diff --git a/tests/_test_utils/torch/quantization/tensor_quantizer_common.py b/tests/_test_utils/torch/quantization/tensor_quantizer_common.py index dd8d790ee3f..77d3acb7df3 100644 --- a/tests/_test_utils/torch/quantization/tensor_quantizer_common.py +++ b/tests/_test_utils/torch/quantization/tensor_quantizer_common.py @@ -417,6 +417,49 @@ def test_amax_export(self): amax = quantizer.export_amax() assert amax.shape == (1,) + # One config per ``export_amax`` branch. + UNUSABLE_AMAX_CFGS = { + "dynamic_nvfp4": QuantizerAttributeConfig( + num_bits=(2, 1), block_sizes={-1: 16, "type": "dynamic", "scale_bits": (4, 3)} + ), + "static_per_tensor": QuantizerAttributeConfig(num_bits=4), + } + + @pytest.mark.parametrize("cfg_name", [*UNUSABLE_AMAX_CFGS]) + @pytest.mark.parametrize("bad_value", [0.0, float("nan")]) + def test_amax_export_unusable_amax(self, bad_value, cfg_name): + """An unusable amax must export as a positive scale without mutating the quantizer. + + Regression test for NVBug 6768300. The NaN case pins the ``nan_to_num`` half, which the + zero case alone would not catch. + """ + quantizer = TensorQuantizer(self.UNUSABLE_AMAX_CFGS[cfg_name]).to(self.device) + quantizer.amax = torch.full((1,), bad_value).to(self.device) + + amax = quantizer.export_amax() + + assert torch.all(amax > 0), amax + assert torch.all(amax == quantizer.maxbound), amax + # export must leave the calibrated state alone + stored = quantizer.amax + if bad_value == 0.0: + assert torch.all(stored == 0), stored + else: + assert torch.all(torch.isnan(stored)), stored + + @pytest.mark.parametrize("cfg_name", [*UNUSABLE_AMAX_CFGS]) + def test_amax_export_meta_amax(self, cfg_name): + """``export_amax()`` must stay usable when amax is on the meta device.""" + quantizer = TensorQuantizer(self.UNUSABLE_AMAX_CFGS[cfg_name]) + quantizer.amax = torch.zeros(1, device="meta") + + amax = quantizer.export_amax() + + # Shape differs per branch (the per-tensor path unsqueezes), so pin only that it stays + # meta instead of raising. + assert amax.is_meta, amax + assert amax.numel() == 1, amax.shape + def test_save_restore(self): ref_quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=4, axis=0)) From 6412746f9a7c81b67e94b70c0875b83b728aeadc Mon Sep 17 00:00:00 2001 From: noeyy-mino <174223378+noeyy-mino@users.noreply.github.com> Date: Tue, 22 Sep 2026 08:55:33 +0800 Subject: [PATCH 27/28] Noeyy/fix bug 6701777 (#2402) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix: 6701777 Regression source: "[OMNIML-3349] Add FP8 MHA quantization support for HuggingFace ViT" (#1289), merged into 0.44.0rc3 via the batch cherry-pick #1350. This PR: 1. Registers nn.LayerNorm as a QuantModule for the first time (modelopt/torch/quantization/nn/modules/quant_layernorm.py), intended to let FP8_DEFAULT_CFG's BMM input / LayerNorm output quantizer rules apply to ViT. 2. Removes the prior forced Cast-alignment logic in export_onnx.py that used to normalize Q/DQ node dtypes to trt_high_precision_dtype. Root Cause: Once nn.LayerNorm became a registered QuantModule, these wildcards started unintentionally matching norm1.norm inside FLUX's AdaLayerNormZero block — an elementwise_affine=False LayerNorm with no learnable weight/bias. Its input got routed through NVFP4 Q/DQ (emitted as Float32) while its synthesized affine scale remained native BFloat16, producing the dtype mismatch. Chosen fix: Explicitly exclude nn.LayerNorm from the diffusers NVFP4 presets rather than touching the global QuantModuleRegistry (which ViT FP8 MHA still needs). Add, in both modelopt_recipes/configs/ptq/presets/diffusers/nvfp4.yaml and nvfp4_fp8_mha.yaml, after the existing weight/input wildcard rules (list order matters — later entries override earlier ones): - parent_class: 'nn.LayerNorm' quantizer_name: '*' enable: false ### Usage ``` python examples/diffusers/quantization/quantize.py --model flux-dev --format fp4 --batch-size 2 --percentile 1.0 --alpha 0.8 --quant-algo max --n-steps 20 --quantized-torch-ckpt-save-path /tmp/pytest-of-root/pytest-0/test_diffusers_quant_export_on0/flux-dev-fp4.pt --onnx-dir /tmp/pytest-of-root/pytest-0/test_diffusers_quant_export_on0/flux-dev-fp4 --collect-method default --calib-size 128 --model-dtype BFloat16 --trt-high-precision-dtype BFloat16 trtexec --onnx=/tmp/pytest-of-root/pytest-0/test_diffusers_quant_export_on0/flux-dev-fp4/model.onnx --builderOptimizationLevel=4 --saveEngine=/tmp/pytest-of-root/pytest-0/test_diffusers_quant_export_on0/flux-dev-fp4/model.plan --stronglyTyped --minShapes=hidden_states:1x1024x64,img_ids:1024x3,encoder_hidden_states:1x512x4096,txt_ids:512x3,timestep:1,pooled_projections:1x768,guidance:1 --optShapes=hidden_states:1x4096x64,img_ids:4096x3,encoder_hidden_states:1x512x4096,txt_ids:512x3,timestep:1,pooled_projections:1x768,guidance:1 --maxShapes=hidden_states:1x4096x64,img_ids:4096x3,encoder_hidden_states:1x512x4096,txt_ids:512x3,timestep:1,pooled_projections:1x768,guidance:1 ``` ### Testing The above test commands. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: N/A - 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?: N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - Did you get Claude approval on this PR?: N/A ### Additional Information N/A ## Summary by CodeRabbit * **Bug Fixes** * Improved Diffusers NVFP4 and NVFP4/FP8 MHA quantization presets by excluding LayerNorm modules from quantization. Signed-off-by: Chad Voegele --- modelopt_recipes/configs/ptq/presets/diffusers/nvfp4.yaml | 3 +++ .../configs/ptq/presets/diffusers/nvfp4_fp8_mha.yaml | 3 +++ 2 files changed, 6 insertions(+) diff --git a/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4.yaml b/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4.yaml index 691defb2ae2..47071a4ce9c 100644 --- a/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4.yaml +++ b/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4.yaml @@ -35,3 +35,6 @@ quant_cfg: - quantizer_name: '*softmax_quantizer' cfg: $import: fp8 + - parent_class: 'nn.LayerNorm' + quantizer_name: '*' + enable: false diff --git a/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4_fp8_mha.yaml b/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4_fp8_mha.yaml index ee8c2704fc5..d93166ff11e 100644 --- a/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4_fp8_mha.yaml +++ b/modelopt_recipes/configs/ptq/presets/diffusers/nvfp4_fp8_mha.yaml @@ -39,3 +39,6 @@ quant_cfg: - quantizer_name: '*bmm2_output_quantizer' cfg: $import: fp8 + - parent_class: 'nn.LayerNorm' + quantizer_name: '*' + enable: false From eb22ef0d3031ddb617057892605edc57189d9598 Mon Sep 17 00:00:00 2001 From: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Date: Tue, 22 Sep 2026 14:51:23 +0800 Subject: [PATCH 28/28] [https://nvbugspro.nvidia.com/bug/6778095] Fix fused P-QDQ to respect disabled quantization during calibration (#2434) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix During max calibration, `enable_stats_collection()` calls `disable_quant()`, which sets `_if_quant=False`. The fused causal P-QDQ attention paths bypass `TensorQuantizer.forward()` and previously selected the Triton/Kitchen path from the configured enabled state alone, so P quant-dequant could still execute while quantization was inactive. This change: - enters the Triton P-QDQ path only when `p_bmm_quantizer._if_quant` is true; - bypasses all fused P-QDQ paths when the P quantizer is disabled or quantization is inactive; - adds focused parameterized regression tests covering Triton and Kitchen dispatch across enabled, `disable()`, and `disable_quant()` states. ### Usage N/A. This restores the existing `disable_quant()` contract and does not introduce a new API. ### Testing - `pytest tests/unit/torch/quantization/plugins/test_attention_quant.py`: 14 passed - Targeted pre-commit checks passed - Kitchen coverage verifies the full `{disable(), disable_quant()} × {lazy, already initialized}` dispatch matrix remains bypassed while inactive and resumes after re-enabling - Reproduced with the exact same ModelOpt 0.47.0rc1 wheel on both sides on B300: [module regression build #121](http://dlswqa-nas.nvidia.com:18880/view/yiguo/job/modelopt-quant-module/121/) - Controlled B300 isolation passed when only the P-BMM quantizers were hard-disabled during calibration, and also passed when the existing single Triton attention configuration was forced - Fixed-code B300 validation passed in three independent runs: [Jenkins build 123](http://dlswqa-nas.nvidia.com:18880/job/modelopt-quant-module/123/), [Jenkins build 124](http://dlswqa-nas.nvidia.com:18880/job/modelopt-quant-module/124/), and [Jenkins build 125](http://dlswqa-nas.nvidia.com:18880/job/modelopt-quant-module/125/). Jenkins build 123 compared all 178 module outputs byte-identically and found 0/377 `amax` and 0/377 `scale` changes. The confirmed impact is incorrect calibration behavior plus unstable quantizer state and module outputs; downstream benchmark accuracy impact has not been established. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - 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)?: N/A - Did you get Claude approval on this PR?: N/A ### Additional Information The failure was isolated to fused P-QDQ runtime-state dispatch. Quantizer topology and configuration were identical in the failing same-wheel comparison. ## Summary by CodeRabbit - **Bug Fixes** - Corrected attention dispatch so disabled or inactive quantization uses the original attention implementation. - Preserved optimized quantized attention when quantization is enabled. - Ensured attention masks remain unchanged when using the original attention implementation. - Improved fallback behavior for fused attention paths, including correct initialization and reuse when quantization is re-enabled. - **Tests** - Added coverage for enabled and disabled quantization states, attention-mask handling, fallback selection, and result consistency. --------- Signed-off-by: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Signed-off-by: Chad Voegele --- .../torch/quantization/plugins/huggingface.py | 13 +- .../plugins/test_attention_quant.py | 126 ++++++++++++++++++ 2 files changed, 137 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 2515910ec6e..12e7c49b687 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -247,8 +247,17 @@ def _quantized_attention( return self._eager_p_qdq_attention( original_attention_interface, query_states, key_states, value_states, **kwargs ) - return self._triton_qdq_attention( - p_qdq, query_states, key_states, value_states, **kwargs + if self.p_bmm_quantizer._if_quant: + return self._triton_qdq_attention( + p_qdq, query_states, key_states, value_states, **kwargs + ) + + # Fused P-QDQ paths bypass TensorQuantizer.forward(). + if not self.p_bmm_quantizer.is_enabled or not self.p_bmm_quantizer._if_quant: + if args: + kwargs.pop("attention_mask", None) + return original_attention_interface( + self, query_states, key_states, value_states, *args, **kwargs ) if kitchen is not None and self.kitchen_attn_fn is None: diff --git a/tests/unit/torch/quantization/plugins/test_attention_quant.py b/tests/unit/torch/quantization/plugins/test_attention_quant.py index 702cf3ad1db..e7b8fb69d61 100644 --- a/tests/unit/torch/quantization/plugins/test_attention_quant.py +++ b/tests/unit/torch/quantization/plugins/test_attention_quant.py @@ -185,3 +185,129 @@ def test_p_qdq_mode_detection(): sq.block_sizes = None sq.disable() assert quant_attention._p_qdq_mode() is None + + +@pytest.mark.parametrize("quantization_active", [True, False]) +def test_causal_p_qdq_dispatch_respects_quant_state(monkeypatch, quantization_active): + """Causal attention must dispatch according to the P quantizer runtime state.""" + quant_attention = make_quant_attention() + for name in ("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer"): + getattr(quant_attention, name).disable() + + pq = quant_attention.p_bmm_quantizer + pq.num_bits = (4, 3) + pq.block_sizes = None + if quantization_active: + pq.enable_quant() + else: + pq.disable_quant() + + calls = [] + expected = object() + expected_attention_mask = object() + + def triton_attention(*args, **kwargs): + calls.append("triton") + assert kwargs["attention_mask"] is expected_attention_mask + return expected + + monkeypatch.setattr( + quant_attention, + "_triton_qdq_attention", + triton_attention, + ) + monkeypatch.setattr( + quant_attention, + "_init_kitchen_attn_fn", + lambda: pytest.fail("P quantizer dispatch reached Kitchen initialization"), + ) + + def original_attention(_self, _query, _key, _value, attention_mask): + calls.append("original") + assert attention_mask is expected_attention_mask + return expected + + states = torch.zeros(1, 4, 2, 32) + output = quant_attention._quantized_attention( + original_attention, + quant_attention, + states, + states, + states, + expected_attention_mask, + ) + + assert output is expected + assert calls == ["triton" if quantization_active else "original"] + + +@pytest.mark.parametrize( + ("disable_method", "enable_method"), + [("disable", "enable"), ("disable_quant", "enable_quant")], +) +@pytest.mark.parametrize("preinitialized", [False, True]) +def test_kitchen_dispatch_respects_quantizer_runtime_state( + monkeypatch, disable_method, enable_method, preinitialized +): + """Kitchen dispatch must stop while P quantization is inactive and resume afterward.""" + quant_attention = make_quant_attention() + for name in ("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer"): + getattr(quant_attention, name).disable() + + pq = quant_attention.p_bmm_quantizer + pq.num_bits = (4, 3) + pq.block_sizes = {-1: 32, "type": "dynamic", "scale_bits": (8, 0)} + + calls = [] + + def kitchen_attention(query, _key, _value): + calls.append("kitchen") + return query.flatten(2) + + def init_kitchen(): + calls.append("init") + quant_attention.use_kitchen = True + quant_attention.kitchen_attn_fn = kitchen_attention + + monkeypatch.setattr( + "modelopt.torch.quantization.plugins.huggingface.kitchen", + object(), + ) + monkeypatch.setattr(quant_attention, "_init_kitchen_attn_fn", init_kitchen) + + if preinitialized: + quant_attention.use_kitchen = True + quant_attention.kitchen_attn_fn = kitchen_attention + + expected = object() + + def original_attention(_self, _query, _key, _value): + calls.append("original") + return expected + + states = torch.zeros(1, 4, 2, 32) + getattr(pq, disable_method)() + output = quant_attention._quantized_attention( + original_attention, + quant_attention, + states, + states, + states, + ) + + assert output is expected + assert calls == ["original"] + + calls.clear() + getattr(pq, enable_method)() + output = quant_attention._quantized_attention( + original_attention, + quant_attention, + states, + states, + states, + ) + + assert output[0].shape == (1, 2, 4, 32) + assert output[1] is None + assert calls == (["kitchen"] if preinitialized else ["init", "kitchen"])