From d92db3e5a040c9830861f2e72a13b1efd20fbdba Mon Sep 17 00:00:00 2001 From: gcunhase <4861122+gcunhase@users.noreply.github.com> Date: Fri, 21 Aug 2026 23:57:49 +0000 Subject: [PATCH 1/8] onnx: add sensitivity primitive with exclusion picker for per-op-type / per-node PTQ ranking Adds ``modelopt.onnx.quantization.sensitivity``, a per-op-type or per-node accuracy sensitivity ranking primitive for ONNX PTQ, plus a coverage / threshold based exclusion picker that turns the ranking into an actionable ``--nodes_to_exclude`` or ``--op_types_to_exclude`` list. Primitive (sensitivity.score): - For each quantizable target (op type or individual node), invokes the existing ``modelopt.onnx.quantization.quantize`` entry point to insert calibrated Q/DQ on just that target, runs the reference and quantized ONNXs through ONNXRuntime on the same calibration inputs, and computes a proxy metric between the two graph-output activation sets. Higher score means the target adds more accuracy loss if quantized -- so callers keep high-scoring targets at higher precision. - Op-type granularity via ``--op_types_to_quantize`` (fast; ~10-15 probes on a typical graph); per-node granularity via ``--nodes_to_quantize`` regex (deep dive; N_nodes probes). - Three proxy metrics via ``metrics.py``: kl_div (default, softmax-normalized), mse (raw), cos_dist (1 - cosine_similarity). - Real calibration data via .npy / .npz / directory path, or synthetic random fallback (directional-only; warned in the CLI and marked in the output JSON's calibration_source field). - Default op_types_scope excludes layout / copy ops via ``modelopt.onnx.op_types.is_copy_op`` -- Transpose / Reshape / Concat and friends show up in ORT's default quantizable set, but their sensitivity signal reflects Q/DQ insertion at data-movement boundaries rather than any INT8-kernel trade-off, so ranking them clutters the output with "don't do this anyway" entries. Picker (sensitivity.suggest_exclusion, sensitivity.summarize_exclusion): - Coverage mode (default): return the largest target set whose cumulative sensitivity score stays at or below ``coverage * total_mass``. The actual coverage is always less than or equal to the requested value ("at most X%"), so the operator never gets more exclusion than they asked for. Architecture-portable because the target is a fraction, not an absolute number. - Threshold mode: return every target whose individual sensitivity score exceeds ``threshold``. Simpler and more predictable when the operator already knows what per-target sensitivity magnitude they consider "too sensitive to quantize" for a specific model. - ``near_tie_ratio`` (default 0.99) emits a logger.warning when the cut-off between included and excluded targets is a near-tie -- flags potential intra-group precision fragmentation. Set to ``None`` to disable. - Companion ``summarize_exclusion`` reports the effect of an exclusion set: coverage_pct, num_excluded, num_previously_quantized, num_remaining_quantized, excluded_mass, total_mass. Public surface: - Python: ``from modelopt.onnx.quantization.sensitivity import score, suggest_exclusion, summarize_exclusion``. - CLI: ``python -m modelopt.onnx.quantization.sensitivity --onnx_path=... --calibration_data_path=... --granularity=op_type --metric=kl_div``. - Output JSON schema includes scores, calibration_source, num_calibration_samples, metric, granularity, target_precision. Tests: - ``tests/gpu/onnx/quantization/test_sensitivity.py``: synthetic-graph tier (real deterministic inputs -> LayerNormalization scores above Conv) plus a synthetic-random regression tier (calibration_data=None still preserves the directional invariant), plus CoAtNet-0 op-type and per-node integration stubs marked @pytest.mark.manual (gated by --run-manual and a MODELOPT_SENSITIVITY_FIXTURES env var). - ``tests/unit/onnx/quantization/test_sensitivity_picker.py``: coverage / threshold / min_score_floor / max_nodes / near-tie warning tests. - ``tests/unit/onnx/quantization/test_nodes_to_quantize.py``: validates the existing ``--nodes_to_quantize`` include-only flag that per-node granularity relies on. Documentation: - ``docs/source/guides/_onnx_quantization.rst``: new "Quantization Sensitivity Scan" chapter plus a "Turning scores into an exclusion list" subsection with both policy modes, ``:ref:`` cross-reference to the metric options, and a note that the picker documentation assumes per-node granularity for simplicity but the same logic applies to per-op-type. Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com> Co-Authored-By: Claude Opus 4.7 --- CHANGELOG.rst | 1 + docs/source/guides/_onnx_quantization.rst | 211 +++++++++ modelopt/onnx/op_types.py | 1 + .../onnx/quantization/sensitivity/__init__.py | 43 ++ .../onnx/quantization/sensitivity/__main__.py | 254 ++++++++++ .../onnx/quantization/sensitivity/metrics.py | 114 +++++ .../onnx/quantization/sensitivity/picker.py | 246 ++++++++++ .../onnx/quantization/sensitivity/score.py | 443 ++++++++++++++++++ modelopt/onnx/utils.py | 12 + .../gpu/onnx/quantization/test_sensitivity.py | 297 ++++++++++++ .../quantization/test_nodes_to_quantize.py | 121 +++++ .../quantization/test_sensitivity_picker.py | 197 ++++++++ 12 files changed, 1940 insertions(+) mode change 100755 => 100644 CHANGELOG.rst create mode 100644 modelopt/onnx/quantization/sensitivity/__init__.py create mode 100644 modelopt/onnx/quantization/sensitivity/__main__.py create mode 100644 modelopt/onnx/quantization/sensitivity/metrics.py create mode 100644 modelopt/onnx/quantization/sensitivity/picker.py create mode 100644 modelopt/onnx/quantization/sensitivity/score.py create mode 100644 tests/gpu/onnx/quantization/test_sensitivity.py create mode 100644 tests/unit/onnx/quantization/test_nodes_to_quantize.py create mode 100644 tests/unit/onnx/quantization/test_sensitivity_picker.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst old mode 100755 new mode 100644 index 6687ebd31ea..44dfbbcec56 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -16,6 +16,7 @@ Changelog - Add SFT-masked data support to ``examples/megatron_bridge/distill.py``: ``--sft --sft_dataset_root `` distills on raw prompt-completion JSONL (``{"input", "output"}`` records) with the loss masked to the response tokens, using Megatron-Bridge's ``FinetuningDatasetConfig`` and the model's own HuggingFace tokenizer instead of the pretraining ``GPTDataset`` and ``NullTokenizer``. - Add per-expert weight quantization for Transformer Engine ``TEGroupedMLP`` (fused MoE experts): each expert now has its own ``weight_quantizer`` (a ``GroupedQuantizer`` holding one ``TensorQuantizer`` per expert) with an independent ``amax``, instead of a single shared ``amax`` across all experts. Applies to ``mtq.quantize`` calibration, HF / Megatron export, and QAD. - Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager. +- Add ``modelopt.onnx.quantization.sensitivity`` — per-op-type or per-node accuracy sensitivity ranking for ONNX PTQ, plus coverage or threshold-based exclusion picker that turns the ranking into an actionable ``--nodes_to_exclude`` or ``--op_types_to_exclude`` list, depending on granularity. *Misc* diff --git a/docs/source/guides/_onnx_quantization.rst b/docs/source/guides/_onnx_quantization.rst index e4d0c2d93d6..b62b768fbb5 100644 --- a/docs/source/guides/_onnx_quantization.rst +++ b/docs/source/guides/_onnx_quantization.rst @@ -121,3 +121,214 @@ The following command will build the engine using fp16 precision. After building .. note:: If you replace ``--fp16`` flag with ``--best`` flag, this command will create an int8 engine with TensorRT's implicit quantization. + +Quantization Sensitivity Scan +============================= + +Post-training quantization of any ONNX model often runs into the same friction: it is unclear +which ops or nodes destroy accuracy at INT8/FP8, and practitioners iterate through hand-crafted +exclusion policies until they find one that works. The +:func:`modelopt.onnx.quantization.sensitivity.score` primitive automates that investigation for +any ONNX model with a calibration dataset. It ranks quantizable targets (op types or individual +nodes) by a proxy metric between the reference and per-target quantized activations, so a +downstream picker can decide which ops to keep at higher precision. Works across CNN, +Transformer, and hybrid architectures alike -- the ranking reflects each model's own +precision-sensitive pathways (residual paths, normalization boundaries, SE gating, attention +projections, etc.) without any architecture-specific configuration. The primitive reuses +:func:`modelopt.onnx.quantization.quantize` internally for each per-target probe, so scales are +properly calibrated (not autotune's placement-only descriptors). + +.. _sensitivity-supported-options: + +Supported options +----------------- + +- ``granularity``: ``op_type`` (default; probes each quantizable op type once, ~10-15 probes) or + ``node`` (probes each ONNX node individually, N_nodes probes; slower but per-instance). +- ``metric``: ``kl_div`` (default; softmax-normalized KL divergence — recommended), ``mse`` + (raw mean squared error; cheaper but scale-sensitive) or ``cos`` (``1 - cosine_similarity``; + scale-invariant, robust to activation magnitude variance). +- ``target_precision``: ``int8`` (default) or ``fp8``. +- ``calibration_method``: passed through to the underlying quantize call — ``entropy`` (default), + ``max``, ``mse``, ``percentile``, etc. +- ``calibration_data``: sequence of input-dicts, path to real data (``.npy`` / ``.npz`` / + directory), or ``None`` to fall back to synthetic random tensors (directional-only; see note + below). +- ``op_types_scope``: optional whitelist of op types to probe. If omitted, defaults to the + intersection of ops actually present in the graph and the union of ORT's default quantizable + set, activation ops, normalization ops, and fusible reduction ops. Graph plumbing (``Cast`` / + ``Constant`` / ``Shape`` / ...) is skipped so wall-clock is not wasted on zero-drift probes. + Any ops that slip past the filter but still produce zero drift are hidden from the CLI table + by default (pass ``--show_zero_scores`` to see them; they always appear in the JSON). + +Python API: + +.. code-block:: python + + from modelopt.onnx.quantization.sensitivity import score + + result = score( + onnx_path="coatnet-0.onnx", + calibration_data="imagenet_calib_500.npz", + granularity="op_type", # or "node" + metric="kl_div", # or "mse" or "cos" + target_precision="int8", + ) + # result["scores"] is a dict {op_type_or_node_name: metric_value}, higher = more sensitive. + +The ``imagenet_calib_500.npz`` in the example above is a 500-sample ImageNet-1k calibration set +prepared with the same preprocessing as the exported ONNX. For a CoAtNet-0 checkpoint exported +from timm's ``coatnet_0_rw_224.sw_in1k`` (``pretrained=True``), the code looks like: + +.. code-block:: python + + import numpy as np, onnx, timm, torch + from datasets import load_dataset + from timm.data import resolve_model_data_config, create_transform + + # 1. Export the timm checkpoint to ONNX. + model = timm.create_model("coatnet_0_rw_224.sw_in1k", pretrained=True).eval() + cfg = resolve_model_data_config(model) + dummy = torch.randn(1, *cfg["input_size"]) # (1, 3, 224, 224) + torch.onnx.export( + model, dummy, "coatnet-0.onnx", + input_names=["input"], output_names=["output"], + opset_version=17, + ) + + # 2. Prepare the calibration NPZ with matching preprocessing. + m = onnx.load("coatnet-0.onnx") + input_name = m.graph.input[0].name + tfm = create_transform(**cfg, is_training=False) + ds = load_dataset("ILSVRC/imagenet-1k", split="validation", streaming=True) + samples = [tfm(ex["image"].convert("RGB")).numpy() + for i, ex in enumerate(ds) if i < 500] + np.savez("imagenet_calib_500.npz", + **{input_name: np.stack(samples).astype(np.float32)}) + +Use the analogous timm handle for any other model family (``resnet50``, ``mobilenetv3_large_100``, +``vit_base_patch16_224``, ...); the ``resolve_model_data_config`` +``create_transform`` pair keeps +preprocessing consistent with the exported ONNX regardless of architecture. + +Command line: + +.. code-block:: bash + + # Op-type ranking with real calibration data (one probe per op class; ~14 min on CoAtNet-0) + python -m modelopt.onnx.quantization.sensitivity \ + --onnx_path coatnet-0.onnx \ + --calibration_data_path imagenet_calib_500.npz \ + --granularity op_type \ + --metric kl_div + + # Per-node ranking with real calibration data (one probe per quantizable node; ~60 min on CoAtNet-0) + python -m modelopt.onnx.quantization.sensitivity \ + --onnx_path coatnet-0.onnx \ + --calibration_data_path imagenet_calib_500.npz \ + --granularity node \ + --metric kl_div + +Rendered ranking (CoAtNet-0, real 500-sample ImageNet calibration):: + + Sensitivity scan (int8 / kl_div / op_type): + Add 2.848 <-- highest impact + Mul 1.890 + LayerNormalization 1.653 + ReduceMean 1.570 + BatchNormalization 0.355 + Conv 0.181 + AveragePool 0.057 + Sigmoid 0.039 + MatMul 0.015 + Relu ~0 + Softmax ~0 + GlobalAveragePool ~0 + Gemm 0 <-- lowest impact + (1 target(s) with score 0.0 hidden; pass --show_zero_scores or read the JSON) + Wrote coatnet-0.sensitivity.json + +.. note:: + + Omitting ``--calibration_data_path`` falls back to synthetic random inputs. Absolute scores are + then directional-only and must not be paired with absolute thresholds -- attention-heavy models + are the highest-risk degradation case because random Q times K^T produces near-uniform softmax + that hides real-input MHA quantization pathology. The ``calibration_source`` field of the + output JSON records which mode was used. + +In per-node granularity the scanner iterates over every quantizable node in the graph and runs +one probe per node; each probe uses the existing ``--nodes_to_quantize `` flag on the main +quantize CLI to quantize that node alone (everything else stays FP16) so the resulting output +drift attributes to that specific node. + +Turning scores into an exclusion list +------------------------------------- + +The :func:`sensitivity.score` output is a dictionary from target name to sensitivity score +(see ``metric`` in :ref:`sensitivity-supported-options` above). The picker +function :func:`sensitivity.suggest_exclusion` turns that dictionary into an actionable +``--nodes_to_exclude`` or ``--op_types_to_exclude`` list, depending on granularity, for +:func:`modelopt.onnx.quantization.quantize`, and :func:`sensitivity.summarize_exclusion` +reports what the exclusion set covers. + +In the rest of this documentation, we'll assume ``per-node`` granularity for simplicity, +but the same logic goes for ``per-op-type`` granularity. + +Two policy modes are supported: + +- **Coverage mode** (default): return the largest node set whose cumulative sensitivity + score stays at or below ``coverage * total_mass``. The actual coverage is always less + than or equal to the requested value ("at most X%"). Architecture-portable because the + target is a fraction, not an absolute number -- ``coverage=0.90`` means the same thing + on any model regardless of sensitivity score magnitudes. +- **Threshold mode**: return every node whose individual sensitivity score exceeds + ``threshold`` (no cumulative-mass logic). Simpler and more predictable when the + operator already knows the per-node sensitivity score magnitude that separates + "quantize safely" from "keep at higher precision" for a specific model. Per-node + sensitivity score magnitudes are not portable across models. When ``threshold`` is + set, ``coverage`` is ignored. + +Python API -- coverage mode: + +.. code-block:: python + + from modelopt.onnx.quantization import quantize + from modelopt.onnx.quantization.sensitivity import ( + score, suggest_exclusion, summarize_exclusion, + ) + + result = score( + onnx_path="coatnet-0.onnx", + calibration_data="imagenet_calib_500.npz", + granularity="node", + ) + + # Leave at most 90% of the total sensitivity score mass at FP16; quantize the rest. + excluded = suggest_exclusion(result["scores"], coverage=0.90) + + quantize( + onnx_path="coatnet-0.onnx", + quantize_mode="int8", + calibration_data="imagenet_calib_500.npz", + nodes_to_exclude=excluded, + output_path="coatnet-0.quant.onnx", + ) + +Python API -- threshold mode: + +.. code-block:: python + + # The threshold value is determined empirically by looking at the per-node sensitivity scores. + # For CoAtNet-0, a threshold of 0.02 captures the load-bearing sensitivity + # (roughly the top 25 nodes as per the KL scores, ~89% of total mass). + excluded = suggest_exclusion(result["scores"], threshold=0.02) + +.. note:: + + The picker emits a ``logger.warning`` when the boundary between included and + excluded nodes is a near-tie -- specifically, if the first-excluded node's sensitivity + score is at least 99% of the last-included node's sensitivity score. In that case two + nodes with nearly + equivalent sensitivity end up in different precisions (one FP16, one INT8), which + can produce intra-group precision fragmentation. The warning suggests a slightly + larger ``coverage`` (or smaller ``threshold``) to include the near-tied node. Set + ``near_tie_ratio=None`` to disable the warning entirely. diff --git a/modelopt/onnx/op_types.py b/modelopt/onnx/op_types.py index 637c0ad7a45..f95537a0def 100644 --- a/modelopt/onnx/op_types.py +++ b/modelopt/onnx/op_types.py @@ -407,4 +407,5 @@ def get_activation_ops(): "Softsign", "Swish", "HardSwish", + "Gelu", } diff --git a/modelopt/onnx/quantization/sensitivity/__init__.py b/modelopt/onnx/quantization/sensitivity/__init__.py new file mode 100644 index 00000000000..8c1a64d433c --- /dev/null +++ b/modelopt/onnx/quantization/sensitivity/__init__.py @@ -0,0 +1,43 @@ +# 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. + +"""ONNX quantization sensitivity scan. + +Ranks quantization targets (op types or individual nodes) by the accuracy impact they would have if +quantized. The core primitive, :func:`score`, mutates the graph with a properly calibrated single- +target Q/DQ pass (via the standard :func:`modelopt.onnx.quantization.quantize` entry point), runs +both the FP16 reference and the quantized model through ONNXRuntime, and reports a proxy metric per +target so a downstream picker can decide which ops or nodes to keep at higher precision. +""" + +from modelopt.onnx.quantization.sensitivity.picker import ( + suggest_exclusion, + summarize_exclusion, +) +from modelopt.onnx.quantization.sensitivity.score import ( + CalibrationSource, + Granularity, + Metric, + score, +) + +__all__ = [ + "CalibrationSource", + "Granularity", + "Metric", + "score", + "suggest_exclusion", + "summarize_exclusion", +] diff --git a/modelopt/onnx/quantization/sensitivity/__main__.py b/modelopt/onnx/quantization/sensitivity/__main__.py new file mode 100644 index 00000000000..ba8ecb0f87d --- /dev/null +++ b/modelopt/onnx/quantization/sensitivity/__main__.py @@ -0,0 +1,254 @@ +# 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. + +"""Command-line entrypoint for the ONNX quantization sensitivity scan. + +Runs :func:`modelopt.onnx.quantization.sensitivity.score` and renders the ranked results to stderr +and to a JSON file. Mirrors the flag style of ``python -m modelopt.onnx.quantization.autotune``. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys + +import numpy as np + +from modelopt.onnx.logging_config import logger +from modelopt.onnx.quantization.sensitivity.score import ( + CalibrationSource, + Granularity, + Metric, + score, +) + + +def _default_output_json(onnx_path: str) -> str: + """Derive the default ``--output_json`` path next to the input ONNX file.""" + stem, _ = os.path.splitext(os.path.basename(onnx_path)) + return os.path.join(os.path.dirname(os.path.abspath(onnx_path)), f"{stem}.sensitivity.json") + + +def _load_calibration(path: str | None) -> str | dict | None: + """Return calibration input for :func:`score`. + + If ``path`` is a ``.npz`` file, load it eagerly so the caller sees a proper ``dict`` (matches + what the main quantize CLI does). Directories and ``.npy`` files are passed through as strings so + :func:`score` uses its path-loader. + + Args: + path: Filesystem location or ``None`` for the synthetic-random fallback. + + Returns: + The value to hand to :func:`score` as ``calibration_data``. + """ + if path is None: + return None + if os.path.isdir(path) or path.endswith(".npy"): + return path + if path.endswith(".npz"): + payload = np.load(path, allow_pickle=False) + return {key: payload[key] for key in payload.files} + raise ValueError(f"Unsupported calibration_data_path: {path}") + + +def _render_ranked_table(result: dict, show_zero_scores: bool = False) -> str: + """Format a sensitivity result as a two-column, high-to-low ranked table. + + Args: + result: The return value of :func:`score`. + show_zero_scores: If False (default), hide targets whose drift score is exactly ``0.0``. + Such targets typically indicate op types the underlying quantize call skipped (graph + plumbing like ``Cast`` or ``Reshape``); their zero score is legitimate but noisy in + the ranked table. All scores -- including zeros -- always appear in the JSON output. + + Returns: + A newline-joined string with a header, one row per non-hidden target, and highest / lowest + markers. A trailing footer notes the count of hidden zero-score rows when applicable. + """ + scores = result["scores"] + header = ( + f"Sensitivity scan ({result['target_precision']} / " + f"{result['metric']} / {result['granularity']}):" + ) + if not scores: + return header + "\n (no quantizable targets found)" + + ranked = sorted(scores.items(), key=lambda kv: kv[1], reverse=True) + hidden = 0 + if not show_zero_scores: + visible = [(n, v) for n, v in ranked if v != 0.0] + hidden = len(ranked) - len(visible) + ranked = visible + + if not ranked: + return header + f"\n (all {hidden} target(s) scored 0.0 -- pass --show_zero_scores to see them)" + + name_width = max(len(name) for name, _ in ranked) + lines = [header] + for i, (name, value) in enumerate(ranked): + marker = "" + if i == 0: + marker = " <-- highest impact" + elif i == len(ranked) - 1: + marker = " <-- lowest impact" + lines.append(f" {name:<{name_width}} {value:.3f}{marker}") + if hidden: + lines.append( + f" ({hidden} target(s) with score 0.0 hidden; pass --show_zero_scores or read the JSON)" + ) + return "\n".join(lines) + + +def get_parser() -> argparse.ArgumentParser: + """Build the argparse parser for the sensitivity CLI.""" + parser = argparse.ArgumentParser( + prog="modelopt.onnx.quantization.sensitivity", + description=( + "Rank ONNX quantization targets (op types or individual nodes) by their impact on " + "model output. Emits a ranked table to stderr and a JSON file for downstream tooling." + ), + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "--onnx_path", required=True, type=str, help="Path to the input ONNX model." + ) + parser.add_argument( + "--calibration_data_path", + type=str, + default=None, + help=( + "Real calibration data (.npy, .npz, or a directory of .npz files). If omitted, " + "falls back to synthetic random tensors and produces directional-only rankings." + ), + ) + parser.add_argument( + "--num_calib_samples", + type=int, + default=100, + help="Number of synthetic samples generated when --calibration_data_path is omitted.", + ) + parser.add_argument( + "--granularity", + type=str, + default=Granularity.OP_TYPE.value, + choices=[g.value for g in Granularity], + help="Scan granularity: 'op_type' (fast, one probe per type) or 'node' (per-instance).", + ) + parser.add_argument( + "--metric", + type=str, + default=Metric.KL_DIV.value, + choices=[m.value for m in Metric], + help="Proxy metric between FP-reference and quantized activations.", + ) + parser.add_argument( + "--target_precision", + type=str, + default="int8", + choices=["int8", "fp8"], + help="Precision to probe per target.", + ) + parser.add_argument( + "--calibration_method", + type=str, + default="entropy", + choices=["entropy", "max"], + help="Calibration method threaded through to quantize().", + ) + parser.add_argument( + "--calibration_eps", + type=str, + nargs="+", + default=["cuda:0", "cpu"], + help="ORT execution providers, in priority order.", + ) + parser.add_argument( + "--op_types_scope", + type=str, + nargs="+", + default=None, + help=( + "Optional whitelist of op types to probe. Defaults to every unique op type actually " + "present in the ONNX graph." + ), + ) + parser.add_argument( + "--output_json", + type=str, + default=None, + help="Where to write the sensitivity JSON. Defaults to .sensitivity.json.", + ) + parser.add_argument( + "--show_zero_scores", + action="store_true", + help=( + "Include zero-drift targets (op types the underlying quantize call could not affect) " + "in the stderr ranked table. They always appear in the JSON regardless." + ), + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + """Entry point. + + Args: + argv: Optional argument list (defaults to ``sys.argv[1:]``). Provided for programmatic use + from tests and other callers. + + Returns: + Process exit code: 0 on success, non-zero if :func:`score` raises. + """ + args = get_parser().parse_args(argv) + + if args.calibration_data_path is None: + logger.warning( + "Synthetic random calibration -- scores are directional-only; do not pair with " + "absolute thresholds. See calibration_source in the output JSON." + ) + + calibration_data = _load_calibration(args.calibration_data_path) + + result = score( + onnx_path=args.onnx_path, + calibration_data=calibration_data, + num_synthetic_samples=args.num_calib_samples, + target_precision=args.target_precision, + granularity=args.granularity, + metric=args.metric, + calibration_method=args.calibration_method, + calibration_eps=args.calibration_eps, + op_types_scope=args.op_types_scope, + ) + # Round-trip through str(CalibrationSource(...)) is unnecessary -- score() already emits a plain + # string. Assert here for documentation of the expected schema. + assert result["calibration_source"] in {c.value for c in CalibrationSource} + + payload = {"onnx_path": os.path.abspath(args.onnx_path), **result} + output_json = args.output_json or _default_output_json(args.onnx_path) + os.makedirs(os.path.dirname(os.path.abspath(output_json)) or ".", exist_ok=True) + with open(output_json, "w", encoding="utf-8") as f: + json.dump(payload, f, indent=2, sort_keys=True) + + print(_render_ranked_table(result, show_zero_scores=args.show_zero_scores), file=sys.stderr) + print(f"Wrote {output_json}", file=sys.stderr) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/modelopt/onnx/quantization/sensitivity/metrics.py b/modelopt/onnx/quantization/sensitivity/metrics.py new file mode 100644 index 00000000000..4b7bb135a66 --- /dev/null +++ b/modelopt/onnx/quantization/sensitivity/metrics.py @@ -0,0 +1,114 @@ +# 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. + +"""Proxy metrics for ONNX quantization sensitivity scoring. + +Each metric maps a pair ``(fp16_act, quant_act)`` of aligned activation tensors to a non-negative +scalar. Higher values mean the quantization target under test caused more distortion of the model's +output, so the caller ranks targets by increasing sensitivity to decide what to keep at higher +precision. Callers pass the raw activations exactly as ORT returned them; each metric normalizes +internally where relevant (e.g. softmax for KL) and averages across the leading batch dimension. +""" + +import numpy as np + +__all__ = ["cos_dist", "kl_div", "mse"] + +_EPS = 1e-12 + + +def _flatten_per_sample(tensor: np.ndarray) -> np.ndarray: + """Flatten every non-batch dimension into a single feature dim. + + Args: + tensor: Any-shape numpy array whose first axis is the sample/batch axis. Scalar tensors + (0-D) are treated as a single sample with one feature. + + Returns: + A ``(num_samples, num_features)`` array. + """ + arr = np.asarray(tensor) + if arr.ndim == 0: + return arr.reshape(1, 1) + return arr.reshape(arr.shape[0], -1) + + +def _softmax(logits: np.ndarray, axis: int = -1) -> np.ndarray: + """Numerically stable softmax along ``axis``.""" + shifted = logits - np.max(logits, axis=axis, keepdims=True) + exp = np.exp(shifted) + return exp / (np.sum(exp, axis=axis, keepdims=True) + _EPS) + + +def kl_div(fp16_act: np.ndarray, quant_act: np.ndarray) -> float: + """KL divergence between softmax-normalized FP16 and quantized activations. + + Both tensors are flattened per-sample and passed through softmax to obtain probability + distributions, then the KL divergence ``sum(p * log(p / q))`` is computed per sample and + averaged. This is the recommended default metric because it matches the intuition "output + distribution should be similar" and is robust to activation magnitude scale. + + Args: + fp16_act: FP16 reference activations, shape ``(num_samples, ...)``. + quant_act: Activations from the quantized model, shape ``(num_samples, ...)``. + + Returns: + Mean KL divergence across the sample axis, as a Python float. + """ + p = _softmax(_flatten_per_sample(fp16_act).astype(np.float64)) + q = _softmax(_flatten_per_sample(quant_act).astype(np.float64)) + per_sample = np.sum(p * (np.log(p + _EPS) - np.log(q + _EPS)), axis=-1) + return float(np.mean(per_sample)) + + +def mse(fp16_act: np.ndarray, quant_act: np.ndarray) -> float: + """Mean squared error on raw activation values. + + Cheap to compute but sensitive to activation magnitude scale; a target whose output happens to + be large in absolute value will look more sensitive under MSE than under KL / cosine. + + Args: + fp16_act: FP16 reference activations. + quant_act: Activations from the quantized model with the same shape as ``fp16_act``. + + Returns: + Mean squared error across all elements, as a Python float. + """ + diff = ( + _flatten_per_sample(fp16_act).astype(np.float64) + - _flatten_per_sample(quant_act).astype(np.float64) + ) + return float(np.mean(diff * diff)) + + +def cos_dist(fp16_act: np.ndarray, quant_act: np.ndarray) -> float: + """Cosine distance ``1 - cos(fp16, quant)`` averaged across samples. + + Scale-invariant: robust to models with wide activation-magnitude variance where MSE would be + dominated by the largest tensors. + + Args: + fp16_act: FP16 reference activations. + quant_act: Activations from the quantized model with the same shape as ``fp16_act``. + + Returns: + Mean cosine distance across the sample axis, as a Python float in ``[0, 2]``. + """ + p = _flatten_per_sample(fp16_act).astype(np.float64) + q = _flatten_per_sample(quant_act).astype(np.float64) + dot = np.sum(p * q, axis=-1) + norm = np.linalg.norm(p, axis=-1) * np.linalg.norm(q, axis=-1) + cos = dot / (norm + _EPS) + return float(np.mean(1.0 - cos)) diff --git a/modelopt/onnx/quantization/sensitivity/picker.py b/modelopt/onnx/quantization/sensitivity/picker.py new file mode 100644 index 00000000000..361d5994c08 --- /dev/null +++ b/modelopt/onnx/quantization/sensitivity/picker.py @@ -0,0 +1,246 @@ +# 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. + +"""Exclusion picker for the sensitivity primitive. + +Turns a per-node or per-op-type score dictionary produced by :func:`sensitivity.score` +into an actionable ``--nodes_to_exclude`` or ``--op_types_to_exclude`` list (depending on +granularity) for :func:`modelopt.onnx.quantization.quantize`. Supports two policy modes: + +* **Coverage mode** (default): pick the largest target set whose cumulative + sensitivity score stays at or below ``coverage * total_mass``. Portable + across architectures because the target is a fraction, not an absolute + number. +* **Threshold mode**: exclude every target whose individual sensitivity + score exceeds an absolute cutoff. Simpler and more predictable when the + operator already knows what per-target sensitivity score magnitude they + consider "too sensitive to quantize" for a given model. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +from modelopt.onnx.logging_config import logger + + +def suggest_exclusion( + scores: Mapping[str, float], + coverage: float = 0.90, + *, + threshold: float | None = None, + max_nodes: int | None = None, + min_score_floor: float = 0.0, + near_tie_ratio: float | None = 0.99, +) -> list[str]: + """Return an exclusion list from a per-target sensitivity score dictionary. + + Two policy modes are supported: + + * **Coverage mode** (the default): return the largest target set whose + cumulative sensitivity score stays at or below ``coverage * total_mass``. + Used when ``threshold`` is ``None``. The actual coverage will be less + than or equal to the requested value -- adding the next target in the + ranking would exceed the requested value, so the picker stops before + crossing it. + * **Threshold mode**: return every target whose sensitivity score + exceeds ``threshold``. Used when ``threshold`` is a float; + ``coverage`` is ignored in this mode. + + Coverage mode is architecture-portable (the target is a fraction of the + model's total mass, so the same ``coverage`` value produces + proportionally-sized exclusion sets on different models). Threshold mode + is simpler and more predictable when the operator already knows the + sensitivity score magnitude they consider "too sensitive to quantize" + for the specific model. + + Args: + scores: Per-target (node or op-type) sensitivity scores from + :func:`sensitivity.score` output. + coverage: Fraction of total sensitivity score mass to leave unquantized (coverage + mode only). Guidance: + + * ``0.85 - 0.90`` (default): balanced exploration. Recovers the + majority of the accuracy gap between default QDQ and the FP16 + reference while keeping the exclusion set small enough to + preserve most of the INT8 latency benefit. For architectures + with concentrated sensitivity distributions (e.g., + ResNet-family with sensitivity clustered in the first + bottleneck), ``0.80 - 0.85`` may produce equivalent accuracy + with a smaller exclusion set. + * ``0.95 - 0.99``: accuracy-critical deployments. Larger + exclusion set, approaches the FP16 accuracy ceiling, at the + cost of more Cast boundaries and reduced INT8 latency benefit. + * ``0.70 - 0.80``: performance-critical deployments. Smaller + exclusion set, maximizes INT8 coverage for latency at the + cost of a wider accuracy gap versus the FP16 reference. + + threshold: Absolute sensitivity score cutoff (threshold mode). When + set, every target with individual sensitivity score strictly + greater than ``threshold`` is excluded from quantization; + ``coverage`` is ignored. Set to ``None`` (default) to use + coverage mode. Guidance is model-dependent because per-target + sensitivity score magnitudes scale with model complexity: on + ResNet-50 a value of ``0.005 - 0.02`` picks up + the load-bearing targets; on CoAtNet-0 or larger models + ``0.05 - 0.5`` is a similar magnitude in relative terms. Use + coverage mode if you need portability across models. + max_nodes: Optional cap on the exclusion set size. Prevents + long-tail-heavy distributions from producing very large + exclusion sets that fragment the graph and hurt latency. + Applied in both modes; whichever limit triggers first stops + the accumulation. + min_score_floor: Targets with individual score below this value are + never included, even if the coverage target has not been + reached (coverage mode) or the target exceeds ``threshold`` + (threshold mode -- a defensive check). + near_tie_ratio: If the first-excluded target's sensitivity score is + at least this fraction of the last-included target's sensitivity + score, a warning is emitted via ``logger.warning`` recommending + the operator consider a slightly larger coverage / smaller + threshold to avoid intra-group precision fragmentation. Default + 0.99 (warn when the first-excluded target's sensitivity score is + within 1% of the last-included's). Set to ``None`` to disable + the warning entirely. + + Returns: + List of target names (from ``scores`` keys), sorted from highest to + lowest sensitivity score. Pass to + ``modelopt.onnx.quantization.quantize(..., nodes_to_exclude=...)`` if + ``scores`` came from per-node granularity, or to + ``modelopt.onnx.quantization.quantize(..., op_types_to_exclude=...)`` + if it came from per-op-type granularity. + """ + ranked = sorted(scores.items(), key=lambda kv: -kv[1]) + if not ranked: + return [] + + if threshold is not None: + # Threshold mode: pick every target whose sensitivity score strictly + # exceeds ``threshold``. Iteration order is highest-to-lowest score. + excluded: list[str] = [] + for name, score in ranked: + if score <= threshold or score < min_score_floor: + break + excluded.append(name) + if max_nodes is not None and len(excluded) >= max_nodes: + break + _warn_near_tie(ranked, excluded, near_tie_ratio, mode="threshold") + return excluded + + # Coverage mode: pick the largest target set whose cumulative sensitivity + # score stays at or below ``coverage * total_mass``. Stops BEFORE crossing + # the requested value, so the actual coverage is <= requested. Guarantees + # the operator never gets more exclusion than they asked for. + total = sum(scores.values()) + if total <= 0.0: + return [] + target = coverage * total + + cumulative = 0.0 + excluded = [] + for name, score in ranked: + if score < min_score_floor: + break + if cumulative + score > target: + # Adding this target would exceed the requested coverage; stop. + break + excluded.append(name) + cumulative += score + if max_nodes is not None and len(excluded) >= max_nodes: + break + + _warn_near_tie(ranked, excluded, near_tie_ratio, mode="coverage") + return excluded + + +def _warn_near_tie( + ranked: list[tuple[str, float]], + excluded: list[str], + near_tie_ratio: float | None, + mode: str, +) -> None: + """Emit a logger warning if the cut-off between included and excluded is a near-tie. + + A near-tie means the first-excluded target's sensitivity score is at + least ``near_tie_ratio`` of the last-included target's sensitivity score. + In that case, the two targets carry nearly equivalent sensitivity signal + but end up in different precisions (one FP16, one INT8), which can + produce intra-group fragmentation and unnecessary Cast overhead. The + operator can widen the coverage or lower the threshold to bring the + near-tied target into the exclusion set. + """ + if near_tie_ratio is None: + return + if not excluded or len(excluded) >= len(ranked): + return + last_included_kl = ranked[len(excluded) - 1][1] + if last_included_kl <= 0.0: + return + first_excluded_name, first_excluded_kl = ranked[len(excluded)] + ratio = first_excluded_kl / last_included_kl + if ratio < near_tie_ratio: + return + last_included_name = ranked[len(excluded) - 1][0] + logger.warning( + f"suggest_exclusion (mode={mode}): near-tie at the exclusion cut-off. " + f"Last included target '{last_included_name}' has score={last_included_kl:.5f}, " + f"first excluded target '{first_excluded_name}' has score={first_excluded_kl:.5f} " + f"({100.0 * ratio:.2f}% of last-included). " + f"Consider a slightly larger coverage / smaller threshold to include the " + f"near-tied target and avoid intra-group precision fragmentation." + ) + + +def summarize_exclusion( + scores: Mapping[str, float], + excluded: list[str], +) -> dict: + """Return a summary dictionary describing an exclusion set. + + Useful for logging or reporting the effect of :func:`suggest_exclusion` + before feeding the result into ``modelopt.onnx.quantization.quantize``. + + Args: + scores: The full per-target (node or op-type) sensitivity scores. + excluded: The list of target names that will be excluded from + quantization. + + Returns: + Dict with: + + * ``coverage_pct``: Percentage of total sensitivity score mass + captured by the exclusion set. + * ``num_excluded``: Number of targets to exclude from quantization. + * ``num_previously_quantized``: Total number of quantizable targets + the primitive probed (i.e., what would have been quantized + without the exclusion set). + * ``num_remaining_quantized``: How many targets will still be + quantized after the exclusion set is applied. + * ``excluded_mass``: Absolute cumulative sensitivity score + captured by the exclusion set. + * ``total_mass``: Sum of sensitivity scores across every probed target. + """ + total_mass = sum(scores.values()) + excluded_mass = sum(float(scores.get(name, 0.0)) for name in excluded) + coverage_pct = 100.0 * excluded_mass / total_mass if total_mass > 0.0 else 0.0 + return { + "coverage_pct": coverage_pct, + "num_excluded": len(excluded), + "num_previously_quantized": len(scores), + "num_remaining_quantized": len(scores) - len(excluded), + "excluded_mass": excluded_mass, + "total_mass": total_mass, + } diff --git a/modelopt/onnx/quantization/sensitivity/score.py b/modelopt/onnx/quantization/sensitivity/score.py new file mode 100644 index 00000000000..18e4251d71c --- /dev/null +++ b/modelopt/onnx/quantization/sensitivity/score.py @@ -0,0 +1,443 @@ +# 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. + +"""Core ONNX quantization sensitivity primitive: :func:`score`. + +For every quantization target (an op type or a single node), inserts calibrated Q/DQ nodes on just +that target via the standard :func:`modelopt.onnx.quantization.quantize` entry point, runs the +resulting ONNX and the unquantized reference through ONNXRuntime on the same calibration inputs, +and computes a proxy metric between the two graph-output activation sets. Higher score means the +target degrades the model more if quantized -- so callers keep high scores at higher precision. +""" + +from __future__ import annotations + +import glob +import os +import re +import tempfile +import time +from collections.abc import Callable, Sequence +from enum import Enum + +import numpy as np +import onnx + +from modelopt.onnx.logging_config import logger +from modelopt.onnx.op_types import ( + get_activation_ops, + is_copy_op, + is_default_quantizable_op_by_ort, + is_fusible_reduction_op, + is_normalization_op, +) +from modelopt.onnx.quantization.ort_utils import create_inference_session +from modelopt.onnx.quantization.quantize import quantize +from modelopt.onnx.quantization.sensitivity.metrics import cos_dist, kl_div, mse +from modelopt.onnx.utils import gen_random_inputs, get_input_names, get_op_types_in_graph + +__all__ = ["CalibrationSource", "Granularity", "Metric", "score"] + + +class Metric(str, Enum): + """Proxy metrics between FP16 and quantized activations.""" + + KL_DIV = "kl_div" + MSE = "mse" + COS = "cos" + + +class Granularity(str, Enum): + """Enumeration granularity for sensitivity targets.""" + + OP_TYPE = "op_type" + NODE = "node" + + +class CalibrationSource(str, Enum): + """Origin of the calibration data used for scoring.""" + + REAL = "real" + SYNTHETIC = "synthetic" + + +_METRIC_FUNCS: dict[str, Callable[[np.ndarray, np.ndarray], float]] = { + Metric.KL_DIV.value: kl_div, + Metric.MSE.value: mse, + Metric.COS.value: cos_dist, +} + +# Fixed seed for the synthetic-random calibration fallback so that repeated invocations produce +# identical inputs and, therefore, comparable rankings within one machine. +_SYNTHETIC_SEED = 0 + + +def _default_op_types_scope(onnx_model: onnx.ModelProto) -> set[str]: + """Return op types worth probing by default: present in the graph AND known-quantizable. + + Intersects the set of op types actually present in the graph with the union of ORT's default + quantizable ops, activation ops, normalization ops, and fusible reduction ops. Layout / copy + ops (``Transpose`` / ``Reshape`` / ``Concat`` / ...) are then excluded via + :func:`is_copy_op` -- they show up in ORT's default quantizable set but their sensitivity + signal reflects Q/DQ insertion at data-movement boundaries rather than any INT8-kernel + trade-off, and TensorRT never actually produces INT8 kernels for them, so ranking them + clutters the output with "don't do this anyway" entries. Graph plumbing (``Cast`` / + ``Constant`` / ``Shape`` / ...) not on any of the above lists is also skipped. + + Args: + onnx_model: Loaded ONNX model to enumerate. + + Returns: + Set of op-type strings to probe. + """ + activation_ops = get_activation_ops() + return { + op for op in get_op_types_in_graph(onnx_model) + if ( + is_default_quantizable_op_by_ort(op) + or op in activation_ops + or is_normalization_op(op) + or is_fusible_reduction_op(op) + ) + and not is_copy_op(op) + } + + +def score( + onnx_path: str, + calibration_data: ( + Sequence[dict[str, np.ndarray]] | dict[str, np.ndarray] | np.ndarray | str | None + ) = None, + *, + num_synthetic_samples: int = 100, + target_precision: str = "int8", + granularity: str = "op_type", + metric: str = "kl_div", + calibration_method: str = "entropy", + calibration_eps: Sequence[str] = ("cuda:0", "cpu"), + op_types_scope: Sequence[str] | None = None, + work_dir: str | None = None, +) -> dict: + """Rank quantization targets by their impact on model output. + + Runs one reference forward pass over calibration data on the unquantized ``onnx_path``, then for + each target (op type or node) invokes :func:`modelopt.onnx.quantization.quantize` to insert + calibrated Q/DQ nodes on just that target, re-runs the model, and computes ``metric`` between + the reference and quantized graph outputs. Scores are summed across output tensors and averaged + across the calibration samples inside each metric function; higher score means more accuracy + loss if the target is quantized. + + Args: + onnx_path: Path to the ONNX model to score. The model is treated as the FP-precision + reference and is quantized once per target below. + calibration_data: Calibration inputs. Accepts a ``dict[str, np.ndarray]`` (batch-first), + a ``Sequence[dict[str, np.ndarray]]`` of single-sample dicts, a raw ``np.ndarray`` + (single-input models only), or a path to real data on disk (``.npy`` file, ``.npz`` + file, or directory of ``.npz`` files). Passing ``None`` falls back to synthetic random + tensors of the ONNX's declared input shapes. Synthetic random calibration produces + directional rankings only; see :class:`CalibrationSource` in the returned dict. + num_synthetic_samples: Number of synthetic samples generated when + ``calibration_data is None``. Ignored otherwise. + target_precision: Quantization mode passed through to + :func:`modelopt.onnx.quantization.quantize` for each per-target probe. Supported values + are ``"int8"`` and ``"fp8"``. + granularity: ``"op_type"`` scores each quantizable op type once (one probe per type); + ``"node"`` scores each individual quantizable node (one probe per node), which is + substantially more expensive but pinpoints single-node offenders. + metric: One of :class:`Metric` values -- ``"kl_div"`` (default), ``"mse"``, or ``"cos"``. + calibration_method: Passed through to :func:`modelopt.onnx.quantization.quantize` (defaults + to ``"entropy"`` for int8/fp8). + calibration_eps: ONNXRuntime execution providers to use for both the reference and the + per-target forward passes, and for calibration inside :func:`quantize`. Same schema as + the ``--calibration_eps`` CLI flag. + op_types_scope: Optional whitelist of op types to probe. If omitted, defaults to the + intersection of ops present in ``onnx_path`` and the union of ORT's default + quantizable set, activation ops, normalization ops, and fusible reduction ops + (see :func:`_default_op_types_scope`). Graph plumbing (``Cast`` / ``Constant`` / + ``Shape`` / ...) is skipped by default because it produces zero-drift probes. + Ops that slip past the filter but that the underlying + :func:`modelopt.onnx.quantization.quantize` still cannot quantize are reported + with score ``0.0`` -- the CLI hides those from the pretty-printed table by + default but they always appear in the JSON output. + work_dir: Directory to place intermediate per-target quantized ONNX files. Defaults to a + fresh temporary directory that is removed after the call returns. + + Returns: + A dict with keys: + + * ``scores``: mapping of ``op_type`` (op-type granularity) or ``node_name`` (node + granularity) to the summed metric across graph outputs. + * ``calibration_source``: ``"real"`` if the caller supplied calibration data, ``"synthetic"`` + when the primitive fell back to random tensors. + * ``num_calibration_samples``: number of samples used for the scoring pass. + * ``metric``: the metric name as passed in. + * ``granularity``: ``"op_type"`` or ``"node"``. + * ``target_precision``: the requested quantization precision. + """ + if metric not in _METRIC_FUNCS: + raise ValueError( + f"Unknown metric '{metric}'. Expected one of {list(_METRIC_FUNCS.keys())}." + ) + if granularity not in (Granularity.OP_TYPE.value, Granularity.NODE.value): + raise ValueError( + f"Unknown granularity '{granularity}'. Expected 'op_type' or 'node'." + ) + if target_precision not in ("int8", "fp8"): + raise ValueError( + f"Unsupported target_precision '{target_precision}'. Expected 'int8' or 'fp8'." + ) + + onnx_model = onnx.load(onnx_path) + calib_dict, calibration_source = _resolve_calibration_data( + onnx_model, calibration_data, num_synthetic_samples + ) + num_samples = _num_samples(calib_dict) + logger.info( + f"Sensitivity scan on {onnx_path}: {calibration_source.value} calibration, " + f"{num_samples} samples, granularity={granularity}, metric={metric}, " + f"target_precision={target_precision}" + ) + + quantizable_ops = ( + set(op_types_scope) if op_types_scope else _default_op_types_scope(onnx_model) + ) + if granularity == Granularity.OP_TYPE.value: + targets = _enumerate_op_type_targets(onnx_model, quantizable_ops) + else: + targets = _enumerate_node_targets(onnx_model, quantizable_ops) + if not targets: + logger.warning("No quantizable targets found under the requested scope.") + + metric_fn = _METRIC_FUNCS[metric] + calibration_eps_list = list(calibration_eps) + ref_outputs = _run_inference(onnx_path, calib_dict, calibration_eps_list) + + scores: dict[str, float] = {} + use_tempdir = work_dir is None + tmp_ctx = tempfile.TemporaryDirectory() if use_tempdir else None + target_dir = tmp_ctx.name if tmp_ctx is not None else work_dir + assert target_dir is not None + try: + os.makedirs(target_dir, exist_ok=True) + wall_start = time.monotonic() + for idx, (target_name, quantize_kwargs) in enumerate(targets, start=1): + probe_path = os.path.join( + target_dir, f"probe_{_sanitize_filename(target_name)}.quant.onnx" + ) + step_start = time.monotonic() + try: + quantize( + onnx_path=onnx_path, + quantize_mode=target_precision, + calibration_data=calib_dict, + calibration_method=calibration_method, + calibration_eps=calibration_eps_list, + output_path=probe_path, + # Keep non-quantized ops at fp32 to avoid I/O dtype drift between the reference + # and quantized graphs -- the metric then reflects pure Q/DQ distortion. + high_precision_dtype="fp32", + keep_intermediate_files=False, + **quantize_kwargs, + ) + except Exception as e: + logger.warning( + f"[{idx}/{len(targets)}] quantize() failed for target '{target_name}': {e}" + ) + continue + quant_outputs = _run_inference(probe_path, calib_dict, calibration_eps_list) + scores[target_name] = _pair_metric(ref_outputs, quant_outputs, metric_fn) + logger.info( + f"[{idx}/{len(targets)}] scored '{target_name}' = {scores[target_name]:.6g} " + f"(step {time.monotonic() - step_start:.1f}s, total {time.monotonic() - wall_start:.1f}s)" + ) + finally: + if tmp_ctx is not None: + tmp_ctx.cleanup() + + return { + "scores": scores, + "calibration_source": calibration_source.value, + "num_calibration_samples": num_samples, + "metric": metric, + "granularity": granularity, + "target_precision": target_precision, + } + + +def _resolve_calibration_data( + onnx_model: onnx.ModelProto, + calibration_data: ( + Sequence[dict[str, np.ndarray]] | dict[str, np.ndarray] | np.ndarray | str | None + ), + num_synthetic_samples: int, +) -> tuple[dict[str, np.ndarray], CalibrationSource]: + """Normalize any accepted calibration input into a batch-first ``dict[str, ndarray]``. + + Args: + onnx_model: Loaded ONNX model, used to resolve input names and shapes when the caller + passes an ``ndarray`` (single-input models) or ``None`` (synthetic fallback). + calibration_data: One of the forms documented on :func:`score`. + num_synthetic_samples: Number of synthetic samples to generate when ``calibration_data`` is + ``None``. + + Returns: + A tuple ``(calib_dict, source)`` where ``calib_dict`` has each input as a batch-first + numpy array and ``source`` is either ``CalibrationSource.REAL`` or + ``CalibrationSource.SYNTHETIC``. + """ + input_names = get_input_names(onnx_model) + if calibration_data is None: + # np.random is used inside gen_random_inputs; reseed here so the fallback is deterministic + # across invocations on the same model. + np.random.seed(_SYNTHETIC_SEED) + samples = [gen_random_inputs(onnx_model) for _ in range(num_synthetic_samples)] + return _stack_sample_list(samples), CalibrationSource.SYNTHETIC + if isinstance(calibration_data, str): + return _load_calibration_from_path(calibration_data, input_names), CalibrationSource.REAL + if isinstance(calibration_data, np.ndarray): + assert len(input_names) == 1, ( + "ndarray calibration_data is only valid for single-input models." + ) + return {input_names[0]: calibration_data}, CalibrationSource.REAL + if isinstance(calibration_data, dict): + return {k: np.asarray(v) for k, v in calibration_data.items()}, CalibrationSource.REAL + # Sequence[dict] + return _stack_sample_list(list(calibration_data)), CalibrationSource.REAL + + +def _load_calibration_from_path( + path: str, input_names: list[str] +) -> dict[str, np.ndarray]: + """Load real calibration data from ``.npy``, ``.npz``, or a directory of ``.npz`` files. + + Args: + path: Filesystem location. + input_names: ONNX input names, used to attach ``.npy`` arrays to the sole input. + + Returns: + Batch-first ``dict[str, ndarray]``. + """ + if os.path.isdir(path): + files = sorted(glob.glob(os.path.join(path, "*.npz"))) + assert files, f"No .npz files found under directory {path}" + parts: dict[str, list[np.ndarray]] = {} + for f in files: + payload = np.load(f, allow_pickle=False) + for key in payload.files: + parts.setdefault(key, []).append(payload[key]) + return {k: np.concatenate(v, axis=0) for k, v in parts.items()} + if path.endswith(".npz"): + payload = np.load(path, allow_pickle=False) + return {key: payload[key] for key in payload.files} + if path.endswith(".npy"): + arr = np.load(path, allow_pickle=False) + assert len(input_names) == 1, ( + f"{path} is a single-tensor .npy but the model has {len(input_names)} inputs." + ) + return {input_names[0]: arr} + raise ValueError(f"Unsupported calibration_data path: {path}") + + +def _stack_sample_list(samples: Sequence[dict[str, np.ndarray]]) -> dict[str, np.ndarray]: + """Concatenate a sequence of single-sample dicts into one batch-first dict.""" + assert samples, "Empty calibration sample sequence." + keys = list(samples[0].keys()) + return {k: np.concatenate([np.asarray(s[k]) for s in samples], axis=0) for k in keys} + + +def _num_samples(calib_dict: dict[str, np.ndarray]) -> int: + """Return the batch-axis length of the first array in ``calib_dict``.""" + first = next(iter(calib_dict.values())) + return int(first.shape[0]) + + +def _enumerate_op_type_targets( + onnx_model: onnx.ModelProto, quantizable_ops: set[str] +) -> list[tuple[str, dict]]: + """Return one probe per op type present in the model and in ``quantizable_ops``. + + Args: + onnx_model: Loaded model to enumerate. + quantizable_ops: Whitelist of op types considered quantizable. + + Returns: + List of ``(op_type, quantize_kwargs)`` pairs where ``quantize_kwargs`` restricts + :func:`quantize` to that op type only. + """ + present = {node.op_type for node in onnx_model.graph.node} + scoped = sorted(present & quantizable_ops) + return [(op, {"op_types_to_quantize": [op]}) for op in scoped] + + +def _enumerate_node_targets( + onnx_model: onnx.ModelProto, quantizable_ops: set[str] +) -> list[tuple[str, dict]]: + """Return one probe per named quantizable node. + + Args: + onnx_model: Loaded model to enumerate. + quantizable_ops: Whitelist of op types considered quantizable. + + Returns: + List of ``(node_name, quantize_kwargs)`` pairs where ``quantize_kwargs`` restricts + :func:`quantize` to a regex matching that node only. + """ + targets: list[tuple[str, dict]] = [] + for node in onnx_model.graph.node: + if node.op_type not in quantizable_ops or not node.name: + continue + regex = f"^{re.escape(node.name)}$" + targets.append((node.name, {"nodes_to_quantize": [regex]})) + return targets + + +def _run_inference( + onnx_path: str, calib_dict: dict[str, np.ndarray], calibration_eps: list[str] +) -> list[np.ndarray]: + """Run every sample through ORT and stack outputs along the batch axis. + + Args: + onnx_path: ONNX file to load into an ORT ``InferenceSession``. + calib_dict: Batch-first input dict. + calibration_eps: ORT execution providers, same schema as + :func:`quantize`'s ``calibration_eps``. + + Returns: + List of numpy arrays, one per graph output, each shaped ``(num_samples, ...)``. + """ + session = create_inference_session(onnx_path, calibration_eps) + num_output = len(session.get_outputs()) + num_samples = _num_samples(calib_dict) + per_output: list[list[np.ndarray]] = [[] for _ in range(num_output)] + for i in range(num_samples): + feed = {name: arr[i : i + 1] for name, arr in calib_dict.items()} + outputs = session.run(None, feed) + for j, out in enumerate(outputs): + per_output[j].append(np.asarray(out)) + return [np.concatenate(chunks, axis=0) for chunks in per_output] + + +def _pair_metric( + ref_outputs: list[np.ndarray], + quant_outputs: list[np.ndarray], + metric_fn: Callable[[np.ndarray, np.ndarray], float], +) -> float: + """Sum the metric across matched graph outputs of the reference and quantized models.""" + return float(sum(metric_fn(ref, quant) for ref, quant in zip(ref_outputs, quant_outputs))) + + +def _sanitize_filename(name: str) -> str: + """Turn an arbitrary op/node name into a filesystem-safe token.""" + return re.sub(r"[^A-Za-z0-9._-]", "_", name)[:80] or "unnamed" diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py index f8b5a41a41a..70c1d8b001c 100644 --- a/modelopt/onnx/utils.py +++ b/modelopt/onnx/utils.py @@ -316,6 +316,18 @@ def get_tensor_by_name( return tensor_val or tensor_init or tensor_inp or tensor_out +def get_op_types_in_graph(onnx_model: onnx.ModelProto) -> set[str]: + """Return the set of unique op types that appear as nodes in the graph. + + Args: + onnx_model: Loaded ONNX model. + + Returns: + Set of unique op-type strings appearing in ``onnx_model.graph.node``. + """ + return {node.op_type for node in onnx_model.graph.node if node.op_type} + + def gen_random_inputs( model: onnx.ModelProto, shapes_spec: str | None = None ) -> dict[str, np.ndarray]: diff --git a/tests/gpu/onnx/quantization/test_sensitivity.py b/tests/gpu/onnx/quantization/test_sensitivity.py new file mode 100644 index 00000000000..ad726f3d6a0 --- /dev/null +++ b/tests/gpu/onnx/quantization/test_sensitivity.py @@ -0,0 +1,297 @@ +# 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 the ONNX quantization sensitivity primitive. + +Tiers: + +1. Synthetic-graph unit test with real deterministic inputs -- LayerNorm scores highest. +2. CoAtNet-0 op-type integration (``@pytest.mark.slow`` + real ImageNet calibration). +3. CoAtNet-0 per-node integration (``@pytest.mark.slow_gpu`` + real ImageNet calibration). +4. Synthetic-random calibration regression guard -- LayerNorm still > Conv directionally. + +Tiers 2 and 3 read a pre-staged CoAtNet-0 ONNX + calibration ``.npz`` from a fixtures directory +resolved via ``MODELOPT_ONNX_ACCURACY_MODELS_DIR`` (default ``/tmp``). Missing fixtures ``pytest.skip`` +cleanly. +""" + +from __future__ import annotations + +import os + +import numpy as np +import onnx +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.quantization.sensitivity import score + +_INPUT_NAME = "input" +_OUTPUT_NAME = "output" +_C_IN = 8 +_C_MID = 16 +_H = W = 16 +_MATMUL_DIM = _C_MID * _H * W +_LOGITS = 32 +_FIXTURE_DIR = os.environ.get("MODELOPT_ONNX_ACCURACY_MODELS_DIR", "/tmp") +# Ops covered by the synthetic Conv+MatMul+LN graph. Passed explicitly because the score() +# default -- get_autotuner_quantizable_ops() -- excludes LayerNormalization even though the ModelOpt +# quantize() path registers it via configure_ort. +_SYNTHETIC_OP_SCOPE = ["Conv", "MatMul", "LayerNormalization"] + + +def _build_conv_mm_ln_onnx(path: str, opset: int = 17) -> None: + """Build a small 2-Conv + 1-MatMul + 1-LayerNorm ONNX for deterministic sensitivity tests.""" + rng = np.random.default_rng(0) + w1 = rng.standard_normal((_C_MID, _C_IN, 3, 3)).astype(np.float32) * 0.1 + b1 = np.zeros((_C_MID,), dtype=np.float32) + w2 = rng.standard_normal((_C_MID, _C_MID, 3, 3)).astype(np.float32) * 0.1 + b2 = np.zeros((_C_MID,), dtype=np.float32) + mm = rng.standard_normal((_MATMUL_DIM, _LOGITS)).astype(np.float32) * 0.05 + ln_scale = np.ones((_LOGITS,), dtype=np.float32) + ln_bias = np.zeros((_LOGITS,), dtype=np.float32) + + initializers = [ + numpy_helper.from_array(w1, "w1"), + numpy_helper.from_array(b1, "b1"), + numpy_helper.from_array(w2, "w2"), + numpy_helper.from_array(b2, "b2"), + numpy_helper.from_array(mm, "mm_w"), + numpy_helper.from_array(ln_scale, "ln_scale"), + numpy_helper.from_array(ln_bias, "ln_bias"), + ] + + nodes = [ + helper.make_node( + "Conv", + ["input", "w1", "b1"], + ["conv1_out"], + name="conv_1", + pads=[1, 1, 1, 1], + strides=[1, 1], + ), + helper.make_node( + "Conv", + ["conv1_out", "w2", "b2"], + ["conv2_out"], + name="conv_2", + pads=[1, 1, 1, 1], + strides=[1, 1], + ), + helper.make_node( + "Flatten", ["conv2_out"], ["flat_out"], name="flatten_1", axis=1 + ), + helper.make_node( + "MatMul", ["flat_out", "mm_w"], ["mm_out"], name="matmul_1" + ), + helper.make_node( + "LayerNormalization", + ["mm_out", "ln_scale", "ln_bias"], + [_OUTPUT_NAME], + name="layernorm_1", + axis=-1, + epsilon=1e-5, + ), + ] + + graph = helper.make_graph( + nodes=nodes, + name="sens_test_graph", + inputs=[ + helper.make_tensor_value_info( + _INPUT_NAME, TensorProto.FLOAT, [1, _C_IN, _H, W] + ) + ], + outputs=[ + helper.make_tensor_value_info(_OUTPUT_NAME, TensorProto.FLOAT, [1, _LOGITS]) + ], + initializer=initializers, + ) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("", opset)], ir_version=8 + ) + onnx.save(model, path) + + +def _deterministic_calibration(num_samples: int = 8) -> dict[str, np.ndarray]: + """Fixed-seed calibration data for the synthetic sensitivity graph.""" + rng = np.random.default_rng(42) + return {_INPUT_NAME: rng.standard_normal((num_samples, _C_IN, _H, W)).astype(np.float32)} + + +def _assert_ln_over_conv(scores: dict[str, float]) -> None: + """Directional invariant: LayerNormalization must rank strictly above Conv.""" + assert "LayerNormalization" in scores, f"LayerNorm missing from scores: {scores}" + assert "Conv" in scores, f"Conv missing from scores: {scores}" + assert scores["LayerNormalization"] > scores["Conv"], ( + f"Expected LayerNormalization > Conv, got {scores}" + ) + + +@pytest.mark.parametrize("metric", ["kl_div", "mse", "cos"]) +def test_synthetic_deterministic_ln_highest(tmp_path, metric): + """Tier 1: synthetic graph + deterministic real inputs -- LN scores highest of all ops.""" + onnx_path = str(tmp_path / "sens_synth.onnx") + _build_conv_mm_ln_onnx(onnx_path) + calib = _deterministic_calibration() + + result = score( + onnx_path, + calibration_data=calib, + metric=metric, + target_precision="int8", + granularity="op_type", + calibration_eps=("cpu",), + op_types_scope=_SYNTHETIC_OP_SCOPE, + ) + assert result["calibration_source"] == "real" + assert result["num_calibration_samples"] == 8 + scores = result["scores"] + assert scores, "No scores produced for synthetic graph." + # Highest-scoring op should be LayerNormalization. + top_op = max(scores.items(), key=lambda kv: kv[1])[0] + assert top_op == "LayerNormalization", ( + f"Expected LayerNormalization to be the top-ranked op, got '{top_op}' from {scores}" + ) + _assert_ln_over_conv(scores) + + +def test_synthetic_random_calibration_directional(tmp_path): + """Tier 4: with ``calibration_data=None``, LN > Conv invariant still holds directionally.""" + onnx_path = str(tmp_path / "sens_synth.onnx") + _build_conv_mm_ln_onnx(onnx_path) + + result = score( + onnx_path, + calibration_data=None, + num_synthetic_samples=8, + metric="kl_div", + target_precision="int8", + granularity="op_type", + calibration_eps=("cpu",), + op_types_scope=_SYNTHETIC_OP_SCOPE, + ) + assert result["calibration_source"] == "synthetic" + _assert_ln_over_conv(result["scores"]) + + +def _require_fixture(name: str) -> str: + """Return a fixture path or ``pytest.skip`` if it isn't staged on this host.""" + path = os.path.join(_FIXTURE_DIR, name) + if not os.path.exists(path): + pytest.skip(f"Sensitivity fixture missing: {path}") + return path + + +@pytest.mark.slow +def test_coatnet_op_type_matches_manual_groundtruth(): + """Tier 2: CoAtNet-0 op-type ranking must surface the ops that ``--op_types_to_quantize + Conv`` implicitly avoids. + + Empirical ranking on CoAtNet-0 with 500-sample ImageNet calibration and ``kl_div``: + + Add 2.848 <-- highest impact + Mul 1.890 + LayerNormalization 1.653 + ReduceMean 1.570 + BatchNormalization 0.355 + Conv 0.181 + AveragePool 0.057 + Sigmoid 0.039 + MatMul 0.015 + Relu ~0 + Softmax ~0 + GlobalAveragePool ~0 + Gemm 0 + + Top-4 = Add / Mul / LayerNormalization / ReduceMean are the load-bearing failures + (residual paths, SE gating + softmax scale, norm boundaries). Conv sits ~10x below + the top-4 and quantizes cleanly, matching the manual "Conv-only wins 82% top-1" + ground truth read as a quantization policy. + + Wall-clock ~14 min on H100 with 500 samples / 13 probes (~60s per probe). + + Fixtures (override root via ``MODELOPT_SENSITIVITY_FIXTURES``): + * ``coatnet-0_rw_inpsize_1x3x224x224_opsetv_17_simplified.onnx`` -- baseline ONNX. + * ``imagenet_calib_500.npz`` -- 500-sample ImageNet calibration dict. + """ + onnx_path = _require_fixture( + "coatnet-0_rw_inpsize_1x3x224x224_opsetv_17_simplified.onnx" + ) + calib_path = _require_fixture("imagenet_calib_500.npz") + + result = score( + onnx_path, + calibration_data=calib_path, + metric="kl_div", + target_precision="int8", + granularity="op_type", + calibration_eps=("cuda:0", "cpu"), + ) + assert result["calibration_source"] == "real" + scores = result["scores"] + ranked = sorted(scores.items(), key=lambda kv: kv[1], reverse=True) + + top4 = {name for name, _ in ranked[:4]} + assert {"Add", "Mul", "LayerNormalization", "ReduceMean"}.issubset(top4), ( + f"Top-4 sensitive ops should include Add / Mul / LayerNormalization / " + f"ReduceMean (all > 1.5 KL), got {ranked}" + ) + # Conv sits ~10x below the top-4 -- justifies the Conv-only quantization policy. + assert scores["Conv"] < 0.5, ( + f"Conv score {scores['Conv']:.3f} unexpectedly high (top-4 are all > 1.5)" + ) + # These cluster at ~0 -- primitive won't recommend excluding them because there's + # nothing to exclude. + for op in ("Softmax", "Gemm", "GlobalAveragePool"): + assert scores.get(op, 0.0) < 0.001, ( + f"{op} score {scores.get(op, 0.0):.3g} should be ~0" + ) + + +@pytest.mark.slow_gpu +def test_coatnet_per_node_matches_manual_groundtruth(): + """Tier 3: CoAtNet-0 per-node ranking (LN/MHA nodes top, Conv nodes bottom). + + Wall clock ~30-60 min; gated behind ``@pytest.mark.slow_gpu`` so default CI stays fast. + """ + onnx_path = _require_fixture( + "coatnet-0_rw_inpsize_1x3x224x224_opsetv_17_simplified.onnx" + ) + calib_path = _require_fixture("imagenet_calib_500.npz") + + result = score( + onnx_path, + calibration_data=calib_path, + metric="kl_div", + target_precision="int8", + granularity="node", + calibration_eps=("cuda:0", "cpu"), + ) + assert result["calibration_source"] == "real" + ranked = sorted(result["scores"].items(), key=lambda kv: kv[1], reverse=True) + assert len(ranked) >= 20, "Per-node ranking is unexpectedly short." + top_k = 10 + bottom_k = 10 + top_names = [name for name, _ in ranked[:top_k]] + bottom_names = [name for name, _ in ranked[-bottom_k:]] + # LayerNorm / MHA subgraph nodes dominate the top of the ranking. + assert any("layernorm" in n.lower() or "attn" in n.lower() for n in top_names), ( + f"Expected LN or MHA nodes in top-{top_k}, got {top_names}" + ) + # Individual Conv nodes cluster at the bottom (Conv-only ground truth). + assert any("conv" in n.lower() for n in bottom_names), ( + f"Expected Conv nodes in bottom-{bottom_k}, got {bottom_names}" + ) diff --git a/tests/unit/onnx/quantization/test_nodes_to_quantize.py b/tests/unit/onnx/quantization/test_nodes_to_quantize.py new file mode 100644 index 00000000000..893b1133004 --- /dev/null +++ b/tests/unit/onnx/quantization/test_nodes_to_quantize.py @@ -0,0 +1,121 @@ +# 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 the ``nodes_to_quantize`` allow-list filter (symmetric with ``nodes_to_exclude``). + +Builds a small two-Conv ONNX graph and asserts that ``nodes_to_quantize=["conv_keep"]`` produces +Q/DQ around ``conv_keep`` only, leaving ``conv_skip`` in its original precision. This is the +primitive the ONNX sensitivity scanner relies on to isolate a single node for a per-target probe. +""" + +from __future__ import annotations + +import os + +import numpy as np +import onnx +import onnx_graphsurgeon as gs +from onnx import TensorProto, helper, numpy_helper + +import modelopt.onnx.quantization as moq + + +def _build_two_conv_onnx(path: str, opset: int = 17) -> None: + """Emit a 2-Conv ONNX with the node names the test filters on.""" + rng = np.random.default_rng(0) + w1 = rng.standard_normal((4, 3, 3, 3)).astype(np.float32) * 0.1 + b1 = np.zeros((4,), dtype=np.float32) + w2 = rng.standard_normal((4, 4, 3, 3)).astype(np.float32) * 0.1 + b2 = np.zeros((4,), dtype=np.float32) + + nodes = [ + helper.make_node( + "Conv", + ["input", "w1", "b1"], + ["conv_keep_out"], + name="conv_keep", + pads=[1, 1, 1, 1], + strides=[1, 1], + ), + helper.make_node( + "Conv", + ["conv_keep_out", "w2", "b2"], + ["output"], + name="conv_skip", + pads=[1, 1, 1, 1], + strides=[1, 1], + ), + ] + initializers = [ + numpy_helper.from_array(w1, "w1"), + numpy_helper.from_array(b1, "b1"), + numpy_helper.from_array(w2, "w2"), + numpy_helper.from_array(b2, "b2"), + ] + graph = helper.make_graph( + nodes=nodes, + name="nodes_to_quantize_test", + inputs=[helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 3, 8, 8])], + outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 4, 8, 8])], + initializer=initializers, + ) + onnx.save( + helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)], ir_version=8), + path, + ) + + +def _has_dq_predecessor(node: gs.Node, input_idx: int) -> bool: + """Return True when the input at ``input_idx`` of ``node`` is produced by DequantizeLinear.""" + inp = node.inputs[input_idx] + if not isinstance(inp, gs.Variable): + return False + producer = node.i(input_idx) + if producer and producer.op == "Cast": + producer = producer.i(0) + return bool(producer and producer.op == "DequantizeLinear") + + +def test_nodes_to_quantize_restricts_qdq_to_single_conv(tmp_path): + """`nodes_to_quantize=["conv_keep"]` inserts Q/DQ around conv_keep only.""" + onnx_path = str(tmp_path / "two_conv.onnx") + _build_two_conv_onnx(onnx_path) + calibration_data = {"input": np.random.default_rng(0).standard_normal((2, 3, 8, 8)).astype(np.float32)} + + moq.quantize( + onnx_path, + quantize_mode="int8", + calibration_data=calibration_data, + calibration_eps=["cpu"], + nodes_to_quantize=["^conv_keep$"], + high_precision_dtype="fp32", + ) + + quantized_path = onnx_path.replace(".onnx", ".quant.onnx") + assert os.path.isfile(quantized_path) + + graph = gs.import_onnx(onnx.load(quantized_path)) + keep_nodes = [n for n in graph.nodes if n.name == "conv_keep"] + skip_nodes = [n for n in graph.nodes if n.name == "conv_skip"] + assert len(keep_nodes) == 1, f"conv_keep not found in quantized graph: {[n.name for n in graph.nodes]}" + assert len(skip_nodes) == 1, f"conv_skip not found in quantized graph: {[n.name for n in graph.nodes]}" + + # conv_keep must have DQ on its activation input; conv_skip must not. + assert _has_dq_predecessor(keep_nodes[0], 0), ( + "conv_keep is not quantized despite nodes_to_quantize=['conv_keep']" + ) + assert not _has_dq_predecessor(skip_nodes[0], 0), ( + "conv_skip was quantized but nodes_to_quantize only listed conv_keep" + ) diff --git a/tests/unit/onnx/quantization/test_sensitivity_picker.py b/tests/unit/onnx/quantization/test_sensitivity_picker.py new file mode 100644 index 00000000000..a4d025372e1 --- /dev/null +++ b/tests/unit/onnx/quantization/test_sensitivity_picker.py @@ -0,0 +1,197 @@ +# 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. + +"""Unit tests for :mod:`modelopt.onnx.quantization.sensitivity.picker`.""" + +import pytest + +from modelopt.onnx.quantization.sensitivity.picker import ( + suggest_exclusion, + summarize_exclusion, +) + + +class TestCoverageMode: + """Tests the ``at most X%`` semantic: cumulative KL never exceeds target.""" + + def test_stops_before_crossing_target(self): + # Total = 10. coverage=0.5 -> target 5. top-1 is 4 (fits), top-2 would + # be 7 (crosses 5) -> stop at 1 node. + scores = {"a": 4.0, "b": 3.0, "c": 2.0, "d": 1.0} + assert suggest_exclusion(scores, coverage=0.5) == ["a"] + + def test_includes_second_when_it_fits(self): + # Total = 10. coverage=0.8 -> target 8. top-1 (4) + top-2 (7) both fit, + # top-3 would be 9 (crosses 8) -> stop at 2 nodes. + scores = {"a": 4.0, "b": 3.0, "c": 2.0, "d": 1.0} + assert suggest_exclusion(scores, coverage=0.8) == ["a", "b"] + + def test_full_coverage_returns_all_nodes(self): + # coverage=1.0 -> target = total, everything fits exactly. + scores = {"a": 1.0, "b": 2.0, "c": 3.0} + result = suggest_exclusion(scores, coverage=1.0) + assert set(result) == {"a", "b", "c"} + + def test_returns_sorted_by_kl_desc(self): + scores = {"low": 0.1, "high": 0.9, "mid": 0.5} + # coverage=1.0 -> everything fits, and result is sorted by KL desc. + assert suggest_exclusion(scores, coverage=1.0) == ["high", "mid", "low"] + + def test_top_node_alone_exceeds_target(self): + # Total = 10, coverage=0.2 -> target 2. Top node (5) alone exceeds + # target, so nothing is included. + scores = {"a": 5.0, "b": 3.0, "c": 2.0} + assert suggest_exclusion(scores, coverage=0.2) == [] + + def test_zero_target_returns_empty(self): + scores = {"a": 5.0, "b": 3.0} + assert suggest_exclusion(scores, coverage=0.0) == [] + + def test_zero_total_returns_empty(self): + assert suggest_exclusion({"a": 0.0, "b": 0.0}, coverage=0.9) == [] + + def test_empty_scores_returns_empty(self): + assert suggest_exclusion({}, coverage=0.9) == [] + + def test_max_nodes_caps_exclusion_set(self): + # 10 nodes at KL 10..1. Total = 55. coverage=1.0 would include all, + # but max_nodes=3 caps at 3. + scores = {chr(ord("a") + i): 10.0 - i for i in range(10)} + assert suggest_exclusion(scores, coverage=1.0, max_nodes=3) == ["a", "b", "c"] + + def test_min_score_floor_stops_before_low_nodes(self): + # Even at coverage=1.0, nodes below the floor are excluded from the set. + scores = {"hi_1": 5.0, "hi_2": 4.0, "trivial_1": 0.001, "trivial_2": 0.0001} + result = suggest_exclusion(scores, coverage=1.0, min_score_floor=0.01) + assert result == ["hi_1", "hi_2"] + + def test_vit_like_distribution_undershoots_cleanly(self): + # Mimics ViT-tiny's distribution: 15 nodes at KL ~3.8-6.7, then a + # sharp drop to ~3.05 for ranks 16-17, then a long tail. + big = {f"top_{i}": 6.7 - i * 0.2 for i in range(15)} # ranks 1-15, KL ~6.7 down to ~3.9 + borderline = {"rank_16": 3.057, "rank_17": 3.055} + tail = {f"tail_{i}": 0.5 - i * 0.02 for i in range(30)} + scores = {**big, **borderline, **tail} + total = sum(scores.values()) + result = suggest_exclusion(scores, coverage=0.90) + excluded_mass = sum(scores[n] for n in result) + # Actual coverage never exceeds requested. + assert excluded_mass <= 0.90 * total + # But should still capture most of the mass with fewer than the total. + assert len(result) < len(scores) + + +class TestThresholdMode: + """Tests the absolute-KL cutoff semantic: exclude all nodes above threshold.""" + + def test_picks_all_above_absolute_threshold(self): + scores = {"a": 5.0, "b": 3.0, "c": 1.0, "d": 0.5, "e": 0.05} + assert suggest_exclusion(scores, threshold=1.0) == ["a", "b"] + + def test_returns_sorted_by_kl_desc(self): + scores = {"low_hit": 0.6, "high_hit": 0.9, "mid_hit": 0.75, "miss": 0.1} + assert suggest_exclusion(scores, threshold=0.5) == ["high_hit", "mid_hit", "low_hit"] + + def test_boundary_score_is_excluded_from_set(self): + # A score exactly at the threshold does NOT get excluded (strict >). + scores = {"above": 0.11, "at": 0.10, "below": 0.09} + assert suggest_exclusion(scores, threshold=0.10) == ["above"] + + def test_no_nodes_above_threshold_returns_empty(self): + assert suggest_exclusion({"a": 0.01, "b": 0.005}, threshold=1.0) == [] + + def test_threshold_overrides_coverage(self): + scores = {"a": 5.0, "b": 3.0, "c": 2.98, "d": 2.0} + # coverage=0.99 would try to include most; threshold overrides. + assert suggest_exclusion(scores, coverage=0.99, threshold=2.5) == ["a", "b", "c"] + + def test_max_nodes_still_caps_threshold_mode(self): + # All 10 nodes have score > 5.0 but max_nodes=3 caps at 3. + scores = {chr(ord("a") + i): 10.0 - i * 0.1 for i in range(10)} + assert suggest_exclusion(scores, threshold=5.0, max_nodes=3) == ["a", "b", "c"] + + def test_min_score_floor_composes_with_threshold(self): + # threshold=0.1 would normally include all three, but min_score_floor=1.0 + # short-circuits after "a" (b=0.5 is below the floor). + scores = {"a": 5.0, "b": 0.5, "c": 0.3} + assert suggest_exclusion(scores, threshold=0.1, min_score_floor=1.0) == ["a"] + + +class TestNearTieWarning: + """Warning fires when the cut-off between included and excluded is a near-tie.""" + + def test_warning_fires_on_near_tied_cutoff(self, caplog): + # Ranks 16 and 17 are near-tied at KL 3.06 vs 3.05 (99.7% ratio); coverage=0.94 + # cuts between them. + scores = {f"node_{i:02d}": kl for i, kl in enumerate( + [6.7, 5.7, 4.6, 4.3, 4.1, 4.1, 4.0, 4.0, 3.8, 3.8, + 3.8, 3.8, 3.7, 3.7, 3.7, 3.06, 3.05, 0.8, 0.5, 0.1], 1)} + import logging + with caplog.at_level(logging.WARNING, logger="modelopt.onnx"): + suggest_exclusion(scores, coverage=0.94) + messages = [r.message for r in caplog.records] + assert any("near-tie at the exclusion cut-off" in m for m in messages) + + def test_no_warning_when_cut_is_not_a_near_tie(self, caplog): + # ViT-like distribution where coverage=0.75 cuts between very different KL values. + scores = {f"node_{i:02d}": kl for i, kl in enumerate( + [6.7, 5.7, 4.6, 4.3, 4.1, 0.5, 0.2, 0.1], 1)} + import logging + with caplog.at_level(logging.WARNING, logger="modelopt.onnx"): + suggest_exclusion(scores, coverage=0.75) + messages = [r.message for r in caplog.records] + assert not any("near-tie" in m for m in messages) + + def test_warning_disabled_by_none(self, caplog): + # Setting near_tie_ratio=None disables the warning entirely. + scores = {"a": 5.0, "b": 4.99, "c": 0.1} + import logging + with caplog.at_level(logging.WARNING, logger="modelopt.onnx"): + suggest_exclusion(scores, coverage=0.5, near_tie_ratio=None) + messages = [r.message for r in caplog.records] + assert not any("near-tie" in m for m in messages) + + def test_threshold_mode_also_warns_on_near_tie(self, caplog): + # threshold=3.056 cuts between KL 3.06 (above threshold) and 3.05 (below) -- near-tie. + scores = {"a": 6.7, "b": 5.7, "c": 3.06, "d": 3.05, "e": 0.1} + import logging + with caplog.at_level(logging.WARNING, logger="modelopt.onnx"): + suggest_exclusion(scores, threshold=3.056) + messages = [r.message for r in caplog.records] + assert any("near-tie" in m and "mode=threshold" in m for m in messages) + + +class TestSummarizeExclusion: + def test_reports_coverage_pct_and_counts(self): + scores = {"a": 4.0, "b": 3.0, "c": 2.0, "d": 1.0} + summary = summarize_exclusion(scores, ["a", "b"]) + assert summary["num_excluded"] == 2 + assert summary["num_previously_quantized"] == 4 + assert summary["num_remaining_quantized"] == 2 + assert summary["coverage_pct"] == pytest.approx(70.0) + assert summary["excluded_mass"] == pytest.approx(7.0) + assert summary["total_mass"] == pytest.approx(10.0) + + def test_empty_scores_zero_coverage(self): + summary = summarize_exclusion({}, []) + assert summary["coverage_pct"] == 0.0 + assert summary["num_excluded"] == 0 + + def test_missing_node_names_default_zero(self): + scores = {"a": 5.0, "b": 5.0} + summary = summarize_exclusion(scores, ["a", "unknown"]) + assert summary["excluded_mass"] == pytest.approx(5.0) + assert summary["coverage_pct"] == pytest.approx(50.0) + assert summary["num_excluded"] == 2 From f407d049d2b2e70030eda3e0f37d7883cf2fde90 Mon Sep 17 00:00:00 2001 From: gcunhase <4861122+gcunhase@users.noreply.github.com> Date: Fri, 21 Aug 2026 23:57:50 +0000 Subject: [PATCH 2/8] picker: add block-aware exclusion via optional blocks + block_agg args Extends ``suggest_exclusion`` with two new keyword arguments so operators can turn per-node sensitivity scores into a block-level exclusion set without reimplementing the coverage / threshold / near-tie logic themselves. Motivated by empirical validation on ViT-tiny where per-node picking hits a ~60% top-1 ceiling due to intra-block precision fragmentation; block-level exclusion recovers ~75% top-1 (within 1pp of native ``trtexec --int8 --fp16``). New API surface: * ``blocks: Mapping[str, Sequence[str | re.Pattern]] | None = None`` -- maps group name to a list of regex patterns matching node paths. Each node is assigned to at most one group (first-match wins across the ``blocks`` dict). Nodes matching no pattern become their own singleton group named after themselves, so architecturally-important standalone nodes (final ``LayerNormalization`` before the head, patch-embed ``Conv``, etc.) compete for exclusion on equal footing with multi-node blocks. When ``blocks`` is ``None`` (default) the picker behaves exactly as before -- fully backward-compatible. * ``block_agg: Literal["sum", "max", "mean"] = "sum"`` -- aggregation used to compute a group's score from its members. Kept as ``Literal`` for IDE/mypy support, plus a runtime ``ValueError`` in ``suggest_exclusion`` for defensive validation. Semantics: * When ``blocks`` is set, the picker computes per-group aggregated scores and applies coverage / threshold / near-tie / ``max_nodes`` / ``min_score_floor`` semantics identically to the per-node path. The returned exclusion list is the union of member node names across the selected groups, ready to pass as ``nodes_to_exclude=`` to ``modelopt.onnx.quantization.quantize``. * Natural pairings between ``block_agg`` and picker mode -- documented in the docstring and RST guide: - ``block_agg="sum"`` with ``coverage`` (recommended default): identical "fraction of total KL mass" semantic as per-node coverage. Portable across per-node and per-block grouping on the same model. - ``block_agg="max"`` with ``threshold``: same units as per-node threshold (excludes any group whose peak-node score exceeds the cutoff). Preserves operator intuition when transferring per-node threshold guidance to the block level. - Other combinations remain valid but change what ``coverage`` and ``threshold`` mean in units; the docstring calls this out explicitly. Implementation notes: * The existing per-node core (coverage / threshold / near-tie logic) is extracted into a private ``_pick_from_scores`` helper. Both the per-node and per-block paths call it, so both share identical semantics for every future behavior change. Zero duplication. * Two additional private helpers: ``_assign_groups`` (regex-based first-match-wins assignment with singleton fallback) and ``_aggregate_group_scores`` (dispatches on ``block_agg``). * Backward compatibility: every existing ``suggest_exclusion`` call site behaves exactly as before because ``blocks`` defaults to ``None`` and the ``block_agg`` value is only inspected when ``blocks`` is set. Documentation: * ``docs/source/guides/_onnx_quantization.rst`` gains a new subsection at the end -- "Grouping per-node scores into architectural blocks" -- with: - A ``vit_tiny_patch16_224`` (timm) worked example at ``coverage=0.95`` with ``block_agg="sum"`` that selects blocks 8, 10, 9, 11, 7 (~100 nodes across 5 whole transformer blocks), which recovers ~75% top-1 on ImageNet-1k versus ~60% for the best per-node picking. - A depth-2 example showing how to split each transformer block into ``blocks.N.attn`` and ``blocks.N.mlp`` sub-groups. - Guidance on the ``block_agg`` / picker-mode pairings. - A "when block-level grouping doesn't help" note calling out Conv-heavy architectures (MobileNet, ResNet families) where diffuse per-node sensitivity means the per-node picker still wins. Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com> Co-Authored-By: Claude Opus 4.7 --- docs/source/guides/_onnx_quantization.rst | 140 ++++++++++++++++ .../onnx/quantization/sensitivity/picker.py | 156 +++++++++++++++++- 2 files changed, 289 insertions(+), 7 deletions(-) diff --git a/docs/source/guides/_onnx_quantization.rst b/docs/source/guides/_onnx_quantization.rst index b62b768fbb5..758f5af824e 100644 --- a/docs/source/guides/_onnx_quantization.rst +++ b/docs/source/guides/_onnx_quantization.rst @@ -332,3 +332,143 @@ Python API -- threshold mode: can produce intra-group precision fragmentation. The warning suggests a slightly larger ``coverage`` (or smaller ``threshold``) to include the near-tied node. Set ``near_tie_ratio=None`` to disable the warning entirely. + +Grouping per-node scores into architectural blocks +-------------------------------------------------- + +For attention-heavy transformer architectures (ViT, DeiT, Swin, and hybrids +like CoAtNet's attention stages), per-node picking can miss the actual +accuracy-driving pattern: the top-KL nodes are selected, but excluding them +one by one leaves each affected transformer block with fragmented precision +-- some FP16 nodes, some INT8 nodes -- and softmax numerics degrade +catastrophically. Making the *transformer block* the atomic exclusion unit +avoids the fragmentation entirely. + +Pass a ``blocks`` mapping to :func:`suggest_exclusion` to switch the picker +from per-node to per-block ranking. Every node in the score dict is assigned +to at most one group (first-match wins across ``blocks``); unmatched nodes +automatically become their own singleton group. The picker aggregates +per-node scores into per-group scores via ``block_agg`` (default ``"sum"``), +applies the same coverage / threshold / near-tie / ``max_nodes`` semantics +to the *group* ranking, and returns the expanded node list ready for +``modelopt.onnx.quantization.quantize``. + +Example: ``vit_tiny_patch16_224`` from timm +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The following recipe runs a per-node sensitivity scan on ViT-tiny +(``timm.create_model("vit_tiny_patch16_224", pretrained=True)`` exported via +``torch.onnx.export``), then uses ``suggest_exclusion`` with block-level +grouping at ``coverage=0.95``: + +.. code-block:: python + + from modelopt.onnx.quantization import quantize + from modelopt.onnx.quantization.sensitivity import ( + score, suggest_exclusion, summarize_exclusion, + ) + + result = score( + onnx_path="vit_tiny_patch16_224.onnx", + calibration_data="imagenet_calib_500.npz", + granularity="node", + metric="kl_div", + target_precision="int8", + ) + + # One regex per transformer block: 12 depth-1 groups covering the whole + # block (norm1 + attn + norm2 + mlp + residual Adds). Nodes not matching + # any regex -- e.g. the final /norm/LayerNormalization before the head -- + # automatically become singleton groups and compete for exclusion on equal + # footing with the multi-node blocks. + blocks = {f"blocks.{n}": [rf"^/blocks/blocks\.{n}/"] for n in range(12)} + + excluded = suggest_exclusion( + result["scores"], + coverage=0.95, # capture 95% of total KL mass at the block level + blocks=blocks, + block_agg="sum", # preserves the per-node coverage semantic + ) + + print(summarize_exclusion(result["scores"], excluded)) + + quantize( + onnx_path="vit_tiny_patch16_224.onnx", + output_path="vit_tiny_patch16_224.block_excluded.onnx", + calibration_data="imagenet_calib_500.npz", + nodes_to_exclude=excluded, + quantize_mode="int8", + ) + +Empirically on a 500-image ImageNet-1k validation subset, block-level +exclusion at ``coverage=0.95`` recovers ~75% top-1 versus ~60% for the best +per-node picking (top-K or coverage) -- closing the ViT-tiny parity gap to +native ``trtexec --int8 --fp16``. + +Choosing a grouping depth +~~~~~~~~~~~~~~~~~~~~~~~~~ + +The ``blocks`` argument gives full control over grouping granularity. The +example above uses *depth-1* -- one group per transformer block, each +covering ~20 nodes. For finer control, split each block into its attention +and MLP residual branches (*depth-2*): + +.. code-block:: python + + blocks_depth2 = {} + for n in range(12): + blocks_depth2[f"blocks.{n}.attn"] = [ + rf"^/blocks/blocks\.{n}/norm1", + rf"^/blocks/blocks\.{n}/attn/", + rf"^/blocks/blocks\.{n}/Add$", # residual sum after attention + ] + blocks_depth2[f"blocks.{n}.mlp"] = [ + rf"^/blocks/blocks\.{n}/norm2", + rf"^/blocks/blocks\.{n}/mlp/", + rf"^/blocks/blocks\.{n}/Add_1$", # residual sum after MLP + ] + +Depth-2 is useful when only one branch of a transformer block is sensitive +and you want to keep the other branch at INT8 for latency. For hybrid +architectures like CoAtNet (``/stages/stages.N/blocks/blocks.M/``) or CNNs +like ResNet (``/layerN/M/``), the same principle applies with the +architecture's own path prefixes. + +Mixed depth within one dict is supported -- first-match ordering decides +group assignment when patterns overlap -- so you can use depth-2 for the +sensitivity hot region and depth-1 for the rest of the graph. + +Natural pairings between ``block_agg`` and picker mode +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Under ``blocks``, the picker's ``coverage`` and ``threshold`` semantics +operate on the *aggregated group score*, not on individual node scores. Two +combinations preserve intuition: + +* **``block_agg="sum"`` with ``coverage``** (recommended default): identical + "fraction of total KL mass" semantic as per-node coverage, because summing + group sums equals summing all node scores. Same ``coverage`` value + produces proportionally-sized exclusion sets across per-node and per-block + picking on the same model. +* **``block_agg="max"`` with ``threshold``**: same units as per-node + threshold (excludes any group whose peak-node score exceeds the cutoff). + Preserves operator intuition when transferring per-node threshold values + to the block level. + +Other combinations are valid but change what ``coverage`` and ``threshold`` +mean in units. Under ``block_agg="max"`` coverage counts fraction-of-total- +group-max-scores, not fraction-of-total-KL-mass. Under ``block_agg="sum"`` +threshold operates in summed-KL units per group, so per-node threshold +values must be scaled up by roughly the average block size to select a +comparable number of groups. + +When per-block picking doesn't help +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Block-level grouping is architecture-specific. For Conv-heavy models where +sensitivity is diffuse across many small MBConv or Bottleneck contributors +(MobileNet, ResNet families), per-node ``coverage`` or ``threshold`` +picking typically outperforms block grouping -- either by finding smaller +exclusion sets at equivalent accuracy or by finding higher accuracy at the +same latency. Reach for ``blocks`` first on transformer / attention-heavy +architectures; keep the per-node picker for other cases. diff --git a/modelopt/onnx/quantization/sensitivity/picker.py b/modelopt/onnx/quantization/sensitivity/picker.py index 361d5994c08..84690731a81 100644 --- a/modelopt/onnx/quantization/sensitivity/picker.py +++ b/modelopt/onnx/quantization/sensitivity/picker.py @@ -27,11 +27,21 @@ score exceeds an absolute cutoff. Simpler and more predictable when the operator already knows what per-target sensitivity score magnitude they consider "too sensitive to quantize" for a given model. + +The picker also supports **block-aware grouping** via the ``blocks`` argument: +per-node scores can be aggregated into user-defined architectural groups +(transformer blocks, residual blocks, MBConv stages, ...) and the picker's +coverage / threshold semantics apply to the group ranking rather than +individual nodes. This is useful for transformer / attention-heavy +architectures where per-node picking leaves precision boundaries scrambled +inside the affected blocks and softmax numerics degrade. """ from __future__ import annotations -from collections.abc import Mapping +import re +from collections.abc import Mapping, Sequence +from typing import Literal from modelopt.onnx.logging_config import logger @@ -41,6 +51,8 @@ def suggest_exclusion( coverage: float = 0.90, *, threshold: float | None = None, + blocks: Mapping[str, Sequence[str | re.Pattern]] | None = None, + block_agg: Literal["sum", "max", "mean"] = "sum", max_nodes: int | None = None, min_score_floor: float = 0.0, near_tie_ratio: float | None = 0.99, @@ -97,11 +109,46 @@ def suggest_exclusion( the load-bearing targets; on CoAtNet-0 or larger models ``0.05 - 0.5`` is a similar magnitude in relative terms. Use coverage mode if you need portability across models. - max_nodes: Optional cap on the exclusion set size. Prevents - long-tail-heavy distributions from producing very large - exclusion sets that fragment the graph and hurt latency. - Applied in both modes; whichever limit triggers first stops - the accumulation. + blocks: Optional mapping from group name to a list of regex patterns + (either compiled ``re.Pattern`` objects or plain regex strings) + that match node paths. When provided, the picker ranks *groups* + rather than individual nodes: each node in ``scores`` is + assigned to at most one group (first-match wins across the + ``blocks`` dict); nodes matching no pattern become their own + singleton group named after themselves. Group scores are + computed via ``block_agg``, and coverage / threshold / + ``max_nodes`` / ``near_tie_ratio`` apply identically to the + group ranking. The returned exclusion list is the union of + member node names across the selected groups. Default ``None`` + -- every node is its own singleton group, equivalent to + per-node picking. + block_agg: Aggregation function used to compute a group's score + from its members' individual scores when ``blocks`` is set. + One of ``"sum"``, ``"max"``, ``"mean"``. Natural pairings with + the two policy modes: + + * ``block_agg="sum"`` with **coverage** (recommended default): + identical "fraction of total KL mass" semantic as per-node + coverage, because summing group sums equals summing all node + scores. Portable across granularity choices. + * ``block_agg="max"`` with **threshold**: same units as + per-node threshold (excludes any group whose peak-node score + exceeds the cutoff). Preserves operator intuition when + transferring per-node threshold guidance to the block level. + * Other combinations are valid but change what ``coverage`` + and ``threshold`` mean in units. Under ``block_agg="max"`` + coverage counts fraction-of-total-group-max-scores (not + fraction-of-total-KL-mass). Under ``block_agg="sum"`` + threshold operates in summed-KL units per group, so + per-node threshold values must be scaled up to be + meaningful. Ignored when ``blocks`` is ``None``. + max_nodes: Optional cap on the exclusion set size. When ``blocks`` + is set, this caps the number of *groups* included in the + aggregate ranking before expansion; when ``blocks`` is ``None``, + it caps the number of individual targets. Prevents long-tail- + heavy distributions from producing very large exclusion sets + that fragment the graph and hurt latency. Applied in both + modes; whichever limit triggers first stops the accumulation. min_score_floor: Targets with individual score below this value are never included, even if the coverage target has not been reached (coverage mode) or the target exceeds ``threshold`` @@ -121,7 +168,53 @@ def suggest_exclusion( ``modelopt.onnx.quantization.quantize(..., nodes_to_exclude=...)`` if ``scores`` came from per-node granularity, or to ``modelopt.onnx.quantization.quantize(..., op_types_to_exclude=...)`` - if it came from per-op-type granularity. + if it came from per-op-type granularity. When ``blocks`` is set the + returned list is always suitable for ``nodes_to_exclude=`` because + it is the union of member node names across the selected groups. + """ + if blocks is not None: + if block_agg not in {"sum", "max", "mean"}: + raise ValueError( + f"block_agg must be 'sum', 'max', or 'mean' (got {block_agg!r})" + ) + groups = _assign_groups(scores, blocks) + group_scores = _aggregate_group_scores(scores, groups, block_agg) + selected_groups = _pick_from_scores( + group_scores, + coverage=coverage, + threshold=threshold, + max_nodes=max_nodes, + min_score_floor=min_score_floor, + near_tie_ratio=near_tie_ratio, + ) + return [n for g in selected_groups for n in groups[g]] + + return _pick_from_scores( + scores, + coverage=coverage, + threshold=threshold, + max_nodes=max_nodes, + min_score_floor=min_score_floor, + near_tie_ratio=near_tie_ratio, + ) + + +def _pick_from_scores( + scores: Mapping[str, float], + *, + coverage: float, + threshold: float | None, + max_nodes: int | None, + min_score_floor: float, + near_tie_ratio: float | None, +) -> list[str]: + """Core picker: coverage / threshold selection on any ``{name: score}`` dict. + + Called for per-node picking (from :func:`suggest_exclusion` with + ``blocks=None``) and for per-group picking (from :func:`suggest_exclusion` + with ``blocks`` set, after aggregating per-node scores into per-group + scores). Extracted so both paths share identical coverage / threshold / + near-tie / ``max_nodes`` / ``min_score_floor`` semantics. """ ranked = sorted(scores.items(), key=lambda kv: -kv[1]) if not ranked: @@ -166,6 +259,55 @@ def suggest_exclusion( return excluded +def _assign_groups( + scores: Mapping[str, float], + blocks: Mapping[str, Sequence[str | re.Pattern]], +) -> dict[str, list[str]]: + """Assign each node in ``scores`` to at most one group. + + Rules: + + * A node matching any regex in ``blocks[name]`` joins group ``name``. + * First-match wins across the iteration order of ``blocks`` -- callers + that need mixed-depth grouping should list more-specific groups + earlier. + * Nodes matching no pattern become their own singleton group named + after themselves, so architecturally-important standalone nodes + compete for exclusion on equal footing with multi-node blocks. + """ + compiled = { + gname: [re.compile(p) if isinstance(p, str) else p for p in patterns] + for gname, patterns in blocks.items() + } + groups: dict[str, list[str]] = {} + for node_name in scores: + matched: str | None = None + for gname, pats in compiled.items(): + if any(pat.match(node_name) for pat in pats): + matched = gname + break + key = matched if matched is not None else node_name + groups.setdefault(key, []).append(node_name) + return groups + + +def _aggregate_group_scores( + scores: Mapping[str, float], + groups: Mapping[str, Sequence[str]], + block_agg: Literal["sum", "max", "mean"], +) -> dict[str, float]: + """Aggregate per-node scores into per-group scores using ``block_agg``.""" + if block_agg == "sum": + return {g: sum(scores[n] for n in members) for g, members in groups.items()} + if block_agg == "max": + return {g: max(scores[n] for n in members) for g, members in groups.items()} + # mean + return { + g: (sum(scores[n] for n in members) / len(members)) if members else 0.0 + for g, members in groups.items() + } + + def _warn_near_tie( ranked: list[tuple[str, float]], excluded: list[str], From 449e556e1e5b84d73e1bdefde5f2429c402d0e2a Mon Sep 17 00:00:00 2001 From: gcunhase <4861122+gcunhase@users.noreply.github.com> Date: Mon, 24 Aug 2026 17:51:42 +0000 Subject: [PATCH 3/8] picker/docs: simplify docstrings; swap ViT block-picker example to max + threshold Post-review cleanup after the block-picker feature landed. Three files touched, all documentation / comment simplification -- zero API changes, zero behavior changes. picker.py ========= Compressed verbose bulleted guidance in ``suggest_exclusion``'s docstring without dropping semantic content: * Module docstring reduced from 20 to 5 lines (single-paragraph summary). * ``coverage`` guidance: three bulleted sub-ranges collapsed to a one-line paragraph (``0.85-0.90`` balances, ``0.95-0.99`` favors accuracy, ``0.70-0.80`` favors latency). * ``threshold`` guidance: bulleted per-architecture magnitudes collapsed to a one-line paragraph. * ``blocks`` docstring: kept all rules (first-match wins, singleton fallback, group ranking replaces node ranking, expansion at return) but stripped repeated framing. * ``block_agg`` docstring: kept the two natural-pairing bullets and the paragraph explaining unit shifts under off-diagonal combinations, but removed a sentence-level restatement of each bullet. * ``max_nodes``, ``min_score_floor``, ``near_tie_ratio``: one-paragraph descriptions. Net -170 lines added / +79 lines removed. The extracted helpers (``_pick_from_scores``, ``_assign_groups``, ``_aggregate_group_scores``, ``_warn_near_tie``) are functionally identical to what landed in a3584c81. __main__.py =========== * Removed "Mirrors the flag style of ``python -m modelopt.onnx.quantization.autotune``." sentence from the module docstring. * Shrunk the ``CalibrationSource`` assert comment from two lines to one: "Sanity-check the JSON schema; score() already emits the enum's string value." _onnx_quantization.rst ====================== Revisions after empirical validation of the block-picker recipe against the tested 101-node ViT-tiny hot region: * Trimmed the "primitive reuses ``quantize`` internally, so scales are properly calibrated (not autotune's placement-only descriptors)" clause to just "The primitive reuses ``quantize`` internally for each per-target probe." The autotune-contrast note was a maintainer-facing detail that did not belong in the user guide. * Corrected the ``calibration_method`` bullet from "``entropy`` (default), ``max``, ``mse``, ``percentile``, etc." to the honest "``entropy`` (default) or ``max``". The ONNX quantize path in ``int8.py`` / ``fp8.py`` only dispatches on ``entropy`` vs falls-through-to-MinMax; ``mse`` and ``percentile`` are silently degraded to MinMax with no error. ``PercentileCalibrater`` exists in ``ort_patching.py`` but is not reachable from the public ``calibration_method`` argument. * Tightened ``granularity`` and other bullet-list descriptions. * Removed a redundant "``calibration_source`` field of the output JSON records which mode was used" sentence. * Rewrote the ViT-tiny block-picker example to use the natural ``max`` + ``threshold`` pairing recommended in the picker's docstring instead of ``sum`` + ``max_nodes``. Validated empirically: ``suggest_exclusion(scores, threshold=0.1, blocks=blocks, block_agg="max")`` produces the same 101-node exclusion set as the hand-curated regex-union (checked node-by-node against the vit_tiny_hotregion_blocks_7_11_plus_norm exclusion) and recovers 74.80% top-1 (within 1pp variance of the earlier 75.20% measurement). * Consolidated the two rendered rankings (``max_agg`` and ``sum_agg``) into one side-by-side table so the reader sees both aggregations of the same data at once. Added a note explaining that either ``max + threshold=0.1`` or ``sum + coverage=1.0, max_nodes=6`` picks the same six groups, with a small internal-ordering difference on blocks.9 vs blocks.11 that does not affect the final selection. Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com> Co-Authored-By: Claude Opus 4.7 --- docs/source/guides/_onnx_quantization.rst | 182 ++++++------- .../onnx/quantization/sensitivity/__main__.py | 5 +- .../onnx/quantization/sensitivity/picker.py | 249 ++++++------------ 3 files changed, 174 insertions(+), 262 deletions(-) diff --git a/docs/source/guides/_onnx_quantization.rst b/docs/source/guides/_onnx_quantization.rst index 758f5af824e..c232682bf7c 100644 --- a/docs/source/guides/_onnx_quantization.rst +++ b/docs/source/guides/_onnx_quantization.rst @@ -135,22 +135,21 @@ downstream picker can decide which ops to keep at higher precision. Works across Transformer, and hybrid architectures alike -- the ranking reflects each model's own precision-sensitive pathways (residual paths, normalization boundaries, SE gating, attention projections, etc.) without any architecture-specific configuration. The primitive reuses -:func:`modelopt.onnx.quantization.quantize` internally for each per-target probe, so scales are -properly calibrated (not autotune's placement-only descriptors). +:func:`modelopt.onnx.quantization.quantize` internally for each per-target probe. .. _sensitivity-supported-options: Supported options ----------------- -- ``granularity``: ``op_type`` (default; probes each quantizable op type once, ~10-15 probes) or - ``node`` (probes each ONNX node individually, N_nodes probes; slower but per-instance). +- ``granularity``: ``op_type`` (default; probes each quantizable op type once) or + ``node`` (probes each ONNX node individually; slower). - ``metric``: ``kl_div`` (default; softmax-normalized KL divergence — recommended), ``mse`` (raw mean squared error; cheaper but scale-sensitive) or ``cos`` (``1 - cosine_similarity``; scale-invariant, robust to activation magnitude variance). - ``target_precision``: ``int8`` (default) or ``fp8``. -- ``calibration_method``: passed through to the underlying quantize call — ``entropy`` (default), - ``max``, ``mse``, ``percentile``, etc. +- ``calibration_method``: passed through to the underlying quantize call — ``entropy`` (default) + or ``max``. - ``calibration_data``: sequence of input-dicts, path to real data (``.npy`` / ``.npz`` / directory), or ``None`` to fall back to synthetic random tensors (directional-only; see note below). @@ -170,8 +169,8 @@ Python API: result = score( onnx_path="coatnet-0.onnx", calibration_data="imagenet_calib_500.npz", - granularity="op_type", # or "node" - metric="kl_div", # or "mse" or "cos" + granularity="op_type", # choices = {"op_type", "node"} + metric="kl_div", # choices = {"kl_div", "mse", "cos"} target_precision="int8", ) # result["scores"] is a dict {op_type_or_node_name: metric_value}, higher = more sensitive. @@ -252,8 +251,7 @@ Rendered ranking (CoAtNet-0, real 500-sample ImageNet calibration):: Omitting ``--calibration_data_path`` falls back to synthetic random inputs. Absolute scores are then directional-only and must not be paired with absolute thresholds -- attention-heavy models are the highest-risk degradation case because random Q times K^T produces near-uniform softmax - that hides real-input MHA quantization pathology. The ``calibration_source`` field of the - output JSON records which mode was used. + that hides real-input MHA quantization pathology. In per-node granularity the scanner iterates over every quantizable node in the graph and runs one probe per node; each probe uses the existing ``--nodes_to_quantize `` flag on the main @@ -270,7 +268,7 @@ function :func:`sensitivity.suggest_exclusion` turns that dictionary into an act :func:`modelopt.onnx.quantization.quantize`, and :func:`sensitivity.summarize_exclusion` reports what the exclusion set covers. -In the rest of this documentation, we'll assume ``per-node`` granularity for simplicity, +In the rest of this documentation, we'll cover ``per-node`` granularity for simplicity, but the same logic goes for ``per-op-type`` granularity. Two policy modes are supported: @@ -327,39 +325,37 @@ Python API -- threshold mode: The picker emits a ``logger.warning`` when the boundary between included and excluded nodes is a near-tie -- specifically, if the first-excluded node's sensitivity score is at least 99% of the last-included node's sensitivity score. In that case two - nodes with nearly - equivalent sensitivity end up in different precisions (one FP16, one INT8), which - can produce intra-group precision fragmentation. The warning suggests a slightly - larger ``coverage`` (or smaller ``threshold``) to include the near-tied node. Set - ``near_tie_ratio=None`` to disable the warning entirely. + nodes with nearly equivalent sensitivity end up in different precisions (one FP16, + one INT8), which can produce intra-group precision fragmentation. The warning suggests + a slightly larger ``coverage`` (or smaller ``threshold``) to include the near-tied node. + Set ``near_tie_ratio=None`` to disable the warning entirely. Grouping per-node scores into architectural blocks -------------------------------------------------- -For attention-heavy transformer architectures (ViT, DeiT, Swin, and hybrids -like CoAtNet's attention stages), per-node picking can miss the actual -accuracy-driving pattern: the top-KL nodes are selected, but excluding them -one by one leaves each affected transformer block with fragmented precision --- some FP16 nodes, some INT8 nodes -- and softmax numerics degrade -catastrophically. Making the *transformer block* the atomic exclusion unit -avoids the fragmentation entirely. +On attention-heavy transformer architectures (ViT, DeiT, Swin, CoAtNet's +attention stages), per-node picking can leave affected transformer blocks with +fragmented precision -- some FP16 nodes, some INT8 nodes -- and softmax +numerics degrade catastrophically. Making the *transformer block* the atomic +exclusion unit avoids the fragmentation. Pass a ``blocks`` mapping to :func:`suggest_exclusion` to switch the picker -from per-node to per-block ranking. Every node in the score dict is assigned +from per-node to per-block ranking. Each node in the score dict is assigned to at most one group (first-match wins across ``blocks``); unmatched nodes -automatically become their own singleton group. The picker aggregates -per-node scores into per-group scores via ``block_agg`` (default ``"sum"``), -applies the same coverage / threshold / near-tie / ``max_nodes`` semantics -to the *group* ranking, and returns the expanded node list ready for -``modelopt.onnx.quantization.quantize``. +become their own singleton group. Coverage / threshold / near-tie / +``max_nodes`` semantics apply to the *group* ranking, and the returned +exclusion list is the union of member nodes across the selected groups. See +:func:`suggest_exclusion`'s docstring for the ``block_agg`` / picker-mode +pairings (``sum`` + ``coverage`` and ``max`` + ``threshold`` preserve +per-node units). Example: ``vit_tiny_patch16_224`` from timm ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -The following recipe runs a per-node sensitivity scan on ViT-tiny -(``timm.create_model("vit_tiny_patch16_224", pretrained=True)`` exported via -``torch.onnx.export``), then uses ``suggest_exclusion`` with block-level -grouping at ``coverage=0.95``: +Per-node sensitivity scan on ViT-tiny (``timm.create_model( +"vit_tiny_patch16_224", pretrained=True)`` exported via +``torch.onnx.export``), then block-level exclusion at ``threshold=0.1`` on +group max-KL: .. code-block:: python @@ -376,20 +372,18 @@ grouping at ``coverage=0.95``: target_precision="int8", ) - # One regex per transformer block: 12 depth-1 groups covering the whole - # block (norm1 + attn + norm2 + mlp + residual Adds). Nodes not matching - # any regex -- e.g. the final /norm/LayerNormalization before the head -- - # automatically become singleton groups and compete for exclusion on equal - # footing with the multi-node blocks. + # 12 depth-1 groups, one per transformer block. Standalone nodes not + # matching any regex (e.g. the final /norm/LayerNormalization before the + # head) become singleton groups automatically. blocks = {f"blocks.{n}": [rf"^/blocks/blocks\.{n}/"] for n in range(12)} + # Exclude blocks with threshold above 0.1 KL. On ViT-tiny that cleanly + # captures blocks 7-11 and the final /norm/LayerNormalization singleton + # (see ranking below) while leaving blocks 0-6 in INT8. excluded = suggest_exclusion( result["scores"], - coverage=0.95, # capture 95% of total KL mass at the block level - blocks=blocks, - block_agg="sum", # preserves the per-node coverage semantic + threshold=0.1, blocks=blocks, block_agg="max", ) - print(summarize_exclusion(result["scores"], excluded)) quantize( @@ -400,18 +394,58 @@ grouping at ``coverage=0.95``: quantize_mode="int8", ) -Empirically on a 500-image ImageNet-1k validation subset, block-level -exclusion at ``coverage=0.95`` recovers ~75% top-1 versus ~60% for the best -per-node picking (top-K or coverage) -- closing the ViT-tiny parity gap to -native ``trtexec --int8 --fp16``. +Block-level ranking (ViT-tiny, real 500-sample ImageNet calibration). Both +aggregations shown side-by-side; rows sorted by ``max``:: + + Block ranking (kl_div, sorted by max_agg): + Group max_agg sum_agg + blocks.8 6.737 24.97 <-- highest impact + blocks.10 4.632 17.25 + blocks.11 4.296 14.70 + blocks.9 4.139 15.91 + /norm/LayerNormalization 4.105 4.11 + blocks.7 0.857 1.85 <-- last included at threshold=0.1 + blocks.0 0.011 0.05 + /Add 0.008 0.01 + blocks.6 0.006 ~0.01 + blocks.4 0.005 ~0.01 + blocks.1 0.004 ~0.01 + blocks.2 0.003 ~0.01 + blocks.3 0.003 ~0.01 + blocks.5 0.003 ~0.01 + /patch_embed/proj/Conv ~0 ~0 + /head/Gemm 0 0 <-- lowest impact + + summarize_exclusion: + coverage_pct 99.86 + num_excluded 101 (5 whole transformer blocks + 1 singleton) + num_previously_quantized 244 + num_remaining_quantized 143 + +Both aggregations pick the same top-6 groups (only their internal ordering +of the four hottest blocks differs: ``max`` orders them 8 > 10 > 11 > 9, +while ``sum`` orders 8 > 10 > 9 > 11 because blocks.9 has a slightly heavier +tail than blocks.11), so any of the following expressions produces the same +101-node exclusion: + +.. code-block:: python + + # max + threshold (recommended natural pairing, used in the example above) + suggest_exclusion(scores, threshold=0.1, blocks=blocks, block_agg="max") + + # sum + max_nodes (equivalent -- top 6 groups by cumulative KL mass) + suggest_exclusion(scores, coverage=1.0, max_nodes=6, blocks=blocks, block_agg="sum") + +Empirically on a 500-image ImageNet-1k validation subset, this 101-node +block-level exclusion recovers ~75% top-1 versus ~60% for the best per-node +picking -- closing the ViT-tiny parity gap to implicit quantization. Choosing a grouping depth ~~~~~~~~~~~~~~~~~~~~~~~~~ -The ``blocks`` argument gives full control over grouping granularity. The -example above uses *depth-1* -- one group per transformer block, each -covering ~20 nodes. For finer control, split each block into its attention -and MLP residual branches (*depth-2*): +The example above is *depth-1* (one group per transformer block). For finer +control, split each block into its attention and MLP residual branches +(*depth-2*): .. code-block:: python @@ -428,47 +462,17 @@ and MLP residual branches (*depth-2*): rf"^/blocks/blocks\.{n}/Add_1$", # residual sum after MLP ] -Depth-2 is useful when only one branch of a transformer block is sensitive -and you want to keep the other branch at INT8 for latency. For hybrid -architectures like CoAtNet (``/stages/stages.N/blocks/blocks.M/``) or CNNs -like ResNet (``/layerN/M/``), the same principle applies with the -architecture's own path prefixes. - -Mixed depth within one dict is supported -- first-match ordering decides -group assignment when patterns overlap -- so you can use depth-2 for the -sensitivity hot region and depth-1 for the rest of the graph. - -Natural pairings between ``block_agg`` and picker mode -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -Under ``blocks``, the picker's ``coverage`` and ``threshold`` semantics -operate on the *aggregated group score*, not on individual node scores. Two -combinations preserve intuition: - -* **``block_agg="sum"`` with ``coverage``** (recommended default): identical - "fraction of total KL mass" semantic as per-node coverage, because summing - group sums equals summing all node scores. Same ``coverage`` value - produces proportionally-sized exclusion sets across per-node and per-block - picking on the same model. -* **``block_agg="max"`` with ``threshold``**: same units as per-node - threshold (excludes any group whose peak-node score exceeds the cutoff). - Preserves operator intuition when transferring per-node threshold values - to the block level. - -Other combinations are valid but change what ``coverage`` and ``threshold`` -mean in units. Under ``block_agg="max"`` coverage counts fraction-of-total- -group-max-scores, not fraction-of-total-KL-mass. Under ``block_agg="sum"`` -threshold operates in summed-KL units per group, so per-node threshold -values must be scaled up by roughly the average block size to select a -comparable number of groups. +Use depth-2 to keep one branch of a transformer block at INT8 while +excluding the other. The same principle transfers to hybrids like CoAtNet +(``/stages/stages.N/blocks/blocks.M/``) or CNNs like ResNet (``/layerN/M/``) +with the architecture's own path prefixes. Mixed depth in one dict works +too -- first-match ordering decides assignment when patterns overlap. When per-block picking doesn't help ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -Block-level grouping is architecture-specific. For Conv-heavy models where +Block-level grouping is architecture-specific. On Conv-heavy models where sensitivity is diffuse across many small MBConv or Bottleneck contributors -(MobileNet, ResNet families), per-node ``coverage`` or ``threshold`` -picking typically outperforms block grouping -- either by finding smaller -exclusion sets at equivalent accuracy or by finding higher accuracy at the -same latency. Reach for ``blocks`` first on transformer / attention-heavy -architectures; keep the per-node picker for other cases. +(MobileNet, ResNet families), per-node ``coverage`` or ``threshold`` picking +typically outperforms block grouping. Reach for ``blocks`` first on +transformer / attention-heavy architectures. diff --git a/modelopt/onnx/quantization/sensitivity/__main__.py b/modelopt/onnx/quantization/sensitivity/__main__.py index ba8ecb0f87d..0a33e3aadef 100644 --- a/modelopt/onnx/quantization/sensitivity/__main__.py +++ b/modelopt/onnx/quantization/sensitivity/__main__.py @@ -16,7 +16,7 @@ """Command-line entrypoint for the ONNX quantization sensitivity scan. Runs :func:`modelopt.onnx.quantization.sensitivity.score` and renders the ranked results to stderr -and to a JSON file. Mirrors the flag style of ``python -m modelopt.onnx.quantization.autotune``. +and to a JSON file. """ from __future__ import annotations @@ -235,8 +235,7 @@ def main(argv: list[str] | None = None) -> int: calibration_eps=args.calibration_eps, op_types_scope=args.op_types_scope, ) - # Round-trip through str(CalibrationSource(...)) is unnecessary -- score() already emits a plain - # string. Assert here for documentation of the expected schema. + # Sanity-check the JSON schema; score() already emits the enum's string value. assert result["calibration_source"] in {c.value for c in CalibrationSource} payload = {"onnx_path": os.path.abspath(args.onnx_path), **result} diff --git a/modelopt/onnx/quantization/sensitivity/picker.py b/modelopt/onnx/quantization/sensitivity/picker.py index 84690731a81..20b6a99d3ca 100644 --- a/modelopt/onnx/quantization/sensitivity/picker.py +++ b/modelopt/onnx/quantization/sensitivity/picker.py @@ -15,26 +15,11 @@ """Exclusion picker for the sensitivity primitive. -Turns a per-node or per-op-type score dictionary produced by :func:`sensitivity.score` -into an actionable ``--nodes_to_exclude`` or ``--op_types_to_exclude`` list (depending on -granularity) for :func:`modelopt.onnx.quantization.quantize`. Supports two policy modes: - -* **Coverage mode** (default): pick the largest target set whose cumulative - sensitivity score stays at or below ``coverage * total_mass``. Portable - across architectures because the target is a fraction, not an absolute - number. -* **Threshold mode**: exclude every target whose individual sensitivity - score exceeds an absolute cutoff. Simpler and more predictable when the - operator already knows what per-target sensitivity score magnitude they - consider "too sensitive to quantize" for a given model. - -The picker also supports **block-aware grouping** via the ``blocks`` argument: -per-node scores can be aggregated into user-defined architectural groups -(transformer blocks, residual blocks, MBConv stages, ...) and the picker's -coverage / threshold semantics apply to the group ranking rather than -individual nodes. This is useful for transformer / attention-heavy -architectures where per-node picking leaves precision boundaries scrambled -inside the affected blocks and softmax numerics degrade. +Turns a per-node or per-op-type score dictionary produced by :func:`sensitivity.score` into an +actionable ``--nodes_to_exclude`` or ``--op_types_to_exclude`` list. Supports coverage mode (pick +the largest set whose cumulative score stays at or below ``coverage * total_mass``) and threshold +mode (exclude every target whose individual score exceeds an absolute cutoff), and can optionally +aggregate per-node scores into user-defined architectural groups via the ``blocks`` argument. """ from __future__ import annotations @@ -60,117 +45,57 @@ def suggest_exclusion( """Return an exclusion list from a per-target sensitivity score dictionary. Two policy modes are supported: - - * **Coverage mode** (the default): return the largest target set whose - cumulative sensitivity score stays at or below ``coverage * total_mass``. - Used when ``threshold`` is ``None``. The actual coverage will be less - than or equal to the requested value -- adding the next target in the - ranking would exceed the requested value, so the picker stops before - crossing it. - * **Threshold mode**: return every target whose sensitivity score - exceeds ``threshold``. Used when ``threshold`` is a float; - ``coverage`` is ignored in this mode. - - Coverage mode is architecture-portable (the target is a fraction of the - model's total mass, so the same ``coverage`` value produces - proportionally-sized exclusion sets on different models). Threshold mode - is simpler and more predictable when the operator already knows the - sensitivity score magnitude they consider "too sensitive to quantize" - for the specific model. + - **Coverage mode** (default) returns the largest target set whose cumulative sensitivity score stays + at or below ``coverage * total_mass`` -- the picker stops *before* crossing the target, so the + actual coverage is always <= requested. + - **Threshold mode** (used when ``threshold`` is set; ``coverage`` is ignored) returns every target + whose individual score strictly exceeds ``threshold``. + + Coverage is architecture-portable because it is a fraction of the model's own total mass; threshold + is model-specific but simpler when the operator already knows a reasonable per-target cutoff. Args: - scores: Per-target (node or op-type) sensitivity scores from - :func:`sensitivity.score` output. - coverage: Fraction of total sensitivity score mass to leave unquantized (coverage - mode only). Guidance: - - * ``0.85 - 0.90`` (default): balanced exploration. Recovers the - majority of the accuracy gap between default QDQ and the FP16 - reference while keeping the exclusion set small enough to - preserve most of the INT8 latency benefit. For architectures - with concentrated sensitivity distributions (e.g., - ResNet-family with sensitivity clustered in the first - bottleneck), ``0.80 - 0.85`` may produce equivalent accuracy - with a smaller exclusion set. - * ``0.95 - 0.99``: accuracy-critical deployments. Larger - exclusion set, approaches the FP16 accuracy ceiling, at the - cost of more Cast boundaries and reduced INT8 latency benefit. - * ``0.70 - 0.80``: performance-critical deployments. Smaller - exclusion set, maximizes INT8 coverage for latency at the - cost of a wider accuracy gap versus the FP16 reference. - - threshold: Absolute sensitivity score cutoff (threshold mode). When - set, every target with individual sensitivity score strictly - greater than ``threshold`` is excluded from quantization; - ``coverage`` is ignored. Set to ``None`` (default) to use - coverage mode. Guidance is model-dependent because per-target - sensitivity score magnitudes scale with model complexity: on - ResNet-50 a value of ``0.005 - 0.02`` picks up - the load-bearing targets; on CoAtNet-0 or larger models - ``0.05 - 0.5`` is a similar magnitude in relative terms. Use - coverage mode if you need portability across models. - blocks: Optional mapping from group name to a list of regex patterns - (either compiled ``re.Pattern`` objects or plain regex strings) - that match node paths. When provided, the picker ranks *groups* - rather than individual nodes: each node in ``scores`` is - assigned to at most one group (first-match wins across the - ``blocks`` dict); nodes matching no pattern become their own - singleton group named after themselves. Group scores are - computed via ``block_agg``, and coverage / threshold / - ``max_nodes`` / ``near_tie_ratio`` apply identically to the - group ranking. The returned exclusion list is the union of - member node names across the selected groups. Default ``None`` - -- every node is its own singleton group, equivalent to - per-node picking. - block_agg: Aggregation function used to compute a group's score - from its members' individual scores when ``blocks`` is set. - One of ``"sum"``, ``"max"``, ``"mean"``. Natural pairings with - the two policy modes: - - * ``block_agg="sum"`` with **coverage** (recommended default): - identical "fraction of total KL mass" semantic as per-node - coverage, because summing group sums equals summing all node - scores. Portable across granularity choices. - * ``block_agg="max"`` with **threshold**: same units as - per-node threshold (excludes any group whose peak-node score - exceeds the cutoff). Preserves operator intuition when - transferring per-node threshold guidance to the block level. - * Other combinations are valid but change what ``coverage`` - and ``threshold`` mean in units. Under ``block_agg="max"`` - coverage counts fraction-of-total-group-max-scores (not - fraction-of-total-KL-mass). Under ``block_agg="sum"`` - threshold operates in summed-KL units per group, so - per-node threshold values must be scaled up to be - meaningful. Ignored when ``blocks`` is ``None``. - max_nodes: Optional cap on the exclusion set size. When ``blocks`` - is set, this caps the number of *groups* included in the - aggregate ranking before expansion; when ``blocks`` is ``None``, - it caps the number of individual targets. Prevents long-tail- - heavy distributions from producing very large exclusion sets - that fragment the graph and hurt latency. Applied in both - modes; whichever limit triggers first stops the accumulation. - min_score_floor: Targets with individual score below this value are - never included, even if the coverage target has not been - reached (coverage mode) or the target exceeds ``threshold`` - (threshold mode -- a defensive check). - near_tie_ratio: If the first-excluded target's sensitivity score is - at least this fraction of the last-included target's sensitivity - score, a warning is emitted via ``logger.warning`` recommending - the operator consider a slightly larger coverage / smaller - threshold to avoid intra-group precision fragmentation. Default - 0.99 (warn when the first-excluded target's sensitivity score is - within 1% of the last-included's). Set to ``None`` to disable - the warning entirely. + scores: Per-target (node or op-type) sensitivity scores from :func:`sensitivity.score`. + coverage: Fraction of total sensitivity score mass to leave unquantized (portable across models). + ``0.85-0.90`` (default) balances accuracy and INT8 latency benefit; ``0.95-0.99`` + favors accuracy; ``0.70-0.80`` favors latency. Portable across models. + threshold: Absolute score cutoff. Every target with score strictly greater than + ``threshold`` is excluded. Magnitudes are model-dependent. + blocks: Optional mapping from group name to a list of regex patterns that match node paths + (i.e, ``{group_name: [regex, ...]}``). When set, the picker ranks *groups* rather than + individual nodes: each node is assigned to at most one group (first-match wins across + the dict), unmatched nodes become their own singleton group, and coverage / threshold / + ``max_nodes`` / ``near_tie_ratio`` semantics apply to the group ranking. The returned + exclusion list is the union of member nodes across the selected groups. + Default ``None`` = per-node picking. + block_agg: Aggregation function used to compute a group's score from its members' individual + scores when ``blocks`` is set. One of ``"sum"``, ``"max"``, ``"mean"``. Natural pairings + with the two policy modes: + - ``block_agg="sum"`` with **coverage** (recommended default): preserves the coverage + semantic regardless of granularity choices (group sums equals to summing all node scores) + - ``block_agg="max"`` with **threshold**: preserves per-node threshold units and operator + intuition when transferring per-node threshold guidance to the block level. + + Other combinations are valid but change what ``coverage`` and ``threshold`` mean in + units. Under ``block_agg="max"`` coverage counts fraction-of-total-group-max-scores + (not fraction-of-total-score-mass). Under ``block_agg="sum"`` threshold operates in + summed-score units per group, so per-node threshold values must be scaled up to be + meaningful. Ignored when ``blocks`` is ``None``; ``"mean"`` is provided for completeness. + max_nodes: Optional cap on the number of selected items -- individual targets when + ``blocks`` is ``None``, or groups when ``blocks`` is set. Prevents long-tail-heavy + distributions from producing very large exclusion sets that fragment the graph and + hurt latency. + min_score_floor: Targets with individual score below this value are never included, even + if the coverage target has not been reached or the target exceeds ``threshold``. + near_tie_ratio: If the first-excluded target's score is at least this fraction of the + last-included target's score, a warning is emitted recommending a slightly larger + coverage / smaller threshold to avoid intra-group precision fragmentation. Set to + ``None`` to disable. Default 0.99. Returns: - List of target names (from ``scores`` keys), sorted from highest to - lowest sensitivity score. Pass to - ``modelopt.onnx.quantization.quantize(..., nodes_to_exclude=...)`` if - ``scores`` came from per-node granularity, or to - ``modelopt.onnx.quantization.quantize(..., op_types_to_exclude=...)`` - if it came from per-op-type granularity. When ``blocks`` is set the - returned list is always suitable for ``nodes_to_exclude=`` because - it is the union of member node names across the selected groups. + List of target names sorted highest-to-lowest score. Pass to ``nodes_to_exclude=`` for + per-node scores (or when ``blocks`` is set) and to ``op_types_to_exclude=`` for + per-op-type scores. """ if blocks is not None: if block_agg not in {"sum", "max", "mean"}: @@ -208,21 +133,17 @@ def _pick_from_scores( min_score_floor: float, near_tie_ratio: float | None, ) -> list[str]: - """Core picker: coverage / threshold selection on any ``{name: score}`` dict. + """Coverage / threshold selection on any ``{name: score}`` dict. - Called for per-node picking (from :func:`suggest_exclusion` with - ``blocks=None``) and for per-group picking (from :func:`suggest_exclusion` - with ``blocks`` set, after aggregating per-node scores into per-group - scores). Extracted so both paths share identical coverage / threshold / - near-tie / ``max_nodes`` / ``min_score_floor`` semantics. + Shared between per-node picking and per-group picking (which aggregates per-node scores into + per-group scores first) so both paths use identical selection semantics. """ ranked = sorted(scores.items(), key=lambda kv: -kv[1]) if not ranked: return [] + # Threshold mode if threshold is not None: - # Threshold mode: pick every target whose sensitivity score strictly - # exceeds ``threshold``. Iteration order is highest-to-lowest score. excluded: list[str] = [] for name, score in ranked: if score <= threshold or score < min_score_floor: @@ -233,10 +154,7 @@ def _pick_from_scores( _warn_near_tie(ranked, excluded, near_tie_ratio, mode="threshold") return excluded - # Coverage mode: pick the largest target set whose cumulative sensitivity - # score stays at or below ``coverage * total_mass``. Stops BEFORE crossing - # the requested value, so the actual coverage is <= requested. Guarantees - # the operator never gets more exclusion than they asked for. + # Coverage mode total = sum(scores.values()) if total <= 0.0: return [] @@ -248,7 +166,6 @@ def _pick_from_scores( if score < min_score_floor: break if cumulative + score > target: - # Adding this target would exceed the requested coverage; stop. break excluded.append(name) cumulative += score @@ -266,14 +183,11 @@ def _assign_groups( """Assign each node in ``scores`` to at most one group. Rules: - - * A node matching any regex in ``blocks[name]`` joins group ``name``. - * First-match wins across the iteration order of ``blocks`` -- callers - that need mixed-depth grouping should list more-specific groups - earlier. - * Nodes matching no pattern become their own singleton group named - after themselves, so architecturally-important standalone nodes - compete for exclusion on equal footing with multi-node blocks. + - A node matching any regex in ``blocks[name]`` joins group ``name``. + - First-match wins across the iteration order of ``blocks``, so callers that mix depths list + more-specific groups earlier. + - Nodes matching no pattern become their own singleton group named after themselves so + architecturally-important standalone nodes compete on equal footing with multi-node blocks. """ compiled = { gname: [re.compile(p) if isinstance(p, str) else p for p in patterns] @@ -296,7 +210,7 @@ def _aggregate_group_scores( groups: Mapping[str, Sequence[str]], block_agg: Literal["sum", "max", "mean"], ) -> dict[str, float]: - """Aggregate per-node scores into per-group scores using ``block_agg``.""" + """Aggregate per-node scores into per-group scores.""" if block_agg == "sum": return {g: sum(scores[n] for n in members) for g, members in groups.items()} if block_agg == "max": @@ -314,15 +228,12 @@ def _warn_near_tie( near_tie_ratio: float | None, mode: str, ) -> None: - """Emit a logger warning if the cut-off between included and excluded is a near-tie. + """Warn if the last-included and first-excluded scores are within ``near_tie_ratio``. - A near-tie means the first-excluded target's sensitivity score is at - least ``near_tie_ratio`` of the last-included target's sensitivity score. - In that case, the two targets carry nearly equivalent sensitivity signal - but end up in different precisions (one FP16, one INT8), which can - produce intra-group fragmentation and unnecessary Cast overhead. The - operator can widen the coverage or lower the threshold to bring the - near-tied target into the exclusion set. + When the two boundary targets carry nearly equivalent sensitivity but land in different + precisions (one FP16, one INT8), the resulting Cast boundary tends to produce intra-group + fragmentation. This warning helps guiding the user into adjusting coverage or threshold + to include the near-tied target. """ if near_tie_ratio is None: return @@ -350,30 +261,28 @@ def summarize_exclusion( scores: Mapping[str, float], excluded: list[str], ) -> dict: - """Return a summary dictionary describing an exclusion set. + """Return a summary dict describing an exclusion set. - Useful for logging or reporting the effect of :func:`suggest_exclusion` - before feeding the result into ``modelopt.onnx.quantization.quantize``. + Useful for logging the effect of :func:`suggest_exclusion` before feeding the result into + :func:`modelopt.onnx.quantization.quantize`. Args: - scores: The full per-target (node or op-type) sensitivity scores. - excluded: The list of target names that will be excluded from - quantization. + scores: The full per-target sensitivity scores. + excluded: The list of target names that will be excluded from quantization. Returns: Dict with: - - * ``coverage_pct``: Percentage of total sensitivity score mass + - ``coverage_pct``: Percentage of total sensitivity score mass captured by the exclusion set. - * ``num_excluded``: Number of targets to exclude from quantization. - * ``num_previously_quantized``: Total number of quantizable targets + - ``num_excluded``: Number of targets to exclude from quantization. + - ``num_previously_quantized``: Total number of quantizable targets the primitive probed (i.e., what would have been quantized without the exclusion set). - * ``num_remaining_quantized``: How many targets will still be + - ``num_remaining_quantized``: How many targets will still be quantized after the exclusion set is applied. - * ``excluded_mass``: Absolute cumulative sensitivity score + - ``excluded_mass``: Absolute cumulative sensitivity score captured by the exclusion set. - * ``total_mass``: Sum of sensitivity scores across every probed target. + - ``total_mass``: Sum of sensitivity scores across every probed target. """ total_mass = sum(scores.values()) excluded_mass = sum(float(scores.get(name, 0.0)) for name in excluded) From 2e3bb2a3b3699036bdb51ed2ab44b0f2d8274a5c Mon Sep 17 00:00:00 2001 From: gcunhase <4861122+gcunhase@users.noreply.github.com> Date: Mon, 24 Aug 2026 18:10:37 +0000 Subject: [PATCH 4/8] style: apply ruff auto-fix + format across sensitivity module Ran `ruff check --fix --unsafe-fixes` and `ruff format` from the repo root against the sensitivity package and its tests. All changes are mechanical: - `TYPE_CHECKING` guard on `collections.abc` imports (TC003). - Trailing-whitespace stripping in multi-line docstrings. - One-line reformat of a short `ValueError(...)` that fits on 100 cols. - Multi-line list-literal expansion in test fixtures (ruff format). - Import combining onto a single line where it fits. No semantic changes. `mypy --config-file pyproject.toml` still passes clean across the 5 sensitivity source files, `quantize.py`, `op_types.py`, and `utils.py`. All three pygrep-hooks RST patterns (rst-backticks, rst-directive-colons, rst-inline-touching-normal) return zero hits on the docs update. Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com> Co-Authored-By: modelopt-fix-agent-bot (Opus 4.7) --- .../onnx/quantization/sensitivity/__init__.py | 5 +- .../onnx/quantization/sensitivity/__main__.py | 5 +- .../onnx/quantization/sensitivity/metrics.py | 5 +- .../onnx/quantization/sensitivity/picker.py | 46 +++++++++---------- .../onnx/quantization/sensitivity/score.py | 20 ++++---- .../gpu/onnx/quantization/test_sensitivity.py | 34 ++++---------- .../quantization/test_nodes_to_quantize.py | 12 +++-- .../quantization/test_sensitivity_picker.py | 45 ++++++++++++++---- 8 files changed, 92 insertions(+), 80 deletions(-) diff --git a/modelopt/onnx/quantization/sensitivity/__init__.py b/modelopt/onnx/quantization/sensitivity/__init__.py index 8c1a64d433c..da620b5568a 100644 --- a/modelopt/onnx/quantization/sensitivity/__init__.py +++ b/modelopt/onnx/quantization/sensitivity/__init__.py @@ -22,10 +22,7 @@ target so a downstream picker can decide which ops or nodes to keep at higher precision. """ -from modelopt.onnx.quantization.sensitivity.picker import ( - suggest_exclusion, - summarize_exclusion, -) +from modelopt.onnx.quantization.sensitivity.picker import suggest_exclusion, summarize_exclusion from modelopt.onnx.quantization.sensitivity.score import ( CalibrationSource, Granularity, diff --git a/modelopt/onnx/quantization/sensitivity/__main__.py b/modelopt/onnx/quantization/sensitivity/__main__.py index 0a33e3aadef..0aefcc213a3 100644 --- a/modelopt/onnx/quantization/sensitivity/__main__.py +++ b/modelopt/onnx/quantization/sensitivity/__main__.py @@ -96,7 +96,10 @@ def _render_ranked_table(result: dict, show_zero_scores: bool = False) -> str: ranked = visible if not ranked: - return header + f"\n (all {hidden} target(s) scored 0.0 -- pass --show_zero_scores to see them)" + return ( + header + + f"\n (all {hidden} target(s) scored 0.0 -- pass --show_zero_scores to see them)" + ) name_width = max(len(name) for name, _ in ranked) lines = [header] diff --git a/modelopt/onnx/quantization/sensitivity/metrics.py b/modelopt/onnx/quantization/sensitivity/metrics.py index 4b7bb135a66..c439b696d83 100644 --- a/modelopt/onnx/quantization/sensitivity/metrics.py +++ b/modelopt/onnx/quantization/sensitivity/metrics.py @@ -86,9 +86,8 @@ def mse(fp16_act: np.ndarray, quant_act: np.ndarray) -> float: Returns: Mean squared error across all elements, as a Python float. """ - diff = ( - _flatten_per_sample(fp16_act).astype(np.float64) - - _flatten_per_sample(quant_act).astype(np.float64) + diff = _flatten_per_sample(fp16_act).astype(np.float64) - _flatten_per_sample(quant_act).astype( + np.float64 ) return float(np.mean(diff * diff)) diff --git a/modelopt/onnx/quantization/sensitivity/picker.py b/modelopt/onnx/quantization/sensitivity/picker.py index 20b6a99d3ca..5098cd9e116 100644 --- a/modelopt/onnx/quantization/sensitivity/picker.py +++ b/modelopt/onnx/quantization/sensitivity/picker.py @@ -25,11 +25,13 @@ from __future__ import annotations import re -from collections.abc import Mapping, Sequence -from typing import Literal +from typing import TYPE_CHECKING, Literal from modelopt.onnx.logging_config import logger +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + def suggest_exclusion( scores: Mapping[str, float], @@ -48,37 +50,37 @@ def suggest_exclusion( - **Coverage mode** (default) returns the largest target set whose cumulative sensitivity score stays at or below ``coverage * total_mass`` -- the picker stops *before* crossing the target, so the actual coverage is always <= requested. - - **Threshold mode** (used when ``threshold`` is set; ``coverage`` is ignored) returns every target + - **Threshold mode** (used when ``threshold`` is set; ``coverage`` is ignored) returns every target whose individual score strictly exceeds ``threshold``. - - Coverage is architecture-portable because it is a fraction of the model's own total mass; threshold + + Coverage is architecture-portable because it is a fraction of the model's own total mass; threshold is model-specific but simpler when the operator already knows a reasonable per-target cutoff. Args: scores: Per-target (node or op-type) sensitivity scores from :func:`sensitivity.score`. coverage: Fraction of total sensitivity score mass to leave unquantized (portable across models). - ``0.85-0.90`` (default) balances accuracy and INT8 latency benefit; ``0.95-0.99`` + ``0.85-0.90`` (default) balances accuracy and INT8 latency benefit; ``0.95-0.99`` favors accuracy; ``0.70-0.80`` favors latency. Portable across models. threshold: Absolute score cutoff. Every target with score strictly greater than ``threshold`` is excluded. Magnitudes are model-dependent. blocks: Optional mapping from group name to a list of regex patterns that match node paths - (i.e, ``{group_name: [regex, ...]}``). When set, the picker ranks *groups* rather than - individual nodes: each node is assigned to at most one group (first-match wins across - the dict), unmatched nodes become their own singleton group, and coverage / threshold / - ``max_nodes`` / ``near_tie_ratio`` semantics apply to the group ranking. The returned - exclusion list is the union of member nodes across the selected groups. + (i.e, ``{group_name: [regex, ...]}``). When set, the picker ranks *groups* rather than + individual nodes: each node is assigned to at most one group (first-match wins across + the dict), unmatched nodes become their own singleton group, and coverage / threshold / + ``max_nodes`` / ``near_tie_ratio`` semantics apply to the group ranking. The returned + exclusion list is the union of member nodes across the selected groups. Default ``None`` = per-node picking. - block_agg: Aggregation function used to compute a group's score from its members' individual - scores when ``blocks`` is set. One of ``"sum"``, ``"max"``, ``"mean"``. Natural pairings + block_agg: Aggregation function used to compute a group's score from its members' individual + scores when ``blocks`` is set. One of ``"sum"``, ``"max"``, ``"mean"``. Natural pairings with the two policy modes: - - ``block_agg="sum"`` with **coverage** (recommended default): preserves the coverage + - ``block_agg="sum"`` with **coverage** (recommended default): preserves the coverage semantic regardless of granularity choices (group sums equals to summing all node scores) - ``block_agg="max"`` with **threshold**: preserves per-node threshold units and operator intuition when transferring per-node threshold guidance to the block level. - - Other combinations are valid but change what ``coverage`` and ``threshold`` mean in - units. Under ``block_agg="max"`` coverage counts fraction-of-total-group-max-scores - (not fraction-of-total-score-mass). Under ``block_agg="sum"`` threshold operates in + + Other combinations are valid but change what ``coverage`` and ``threshold`` mean in + units. Under ``block_agg="max"`` coverage counts fraction-of-total-group-max-scores + (not fraction-of-total-score-mass). Under ``block_agg="sum"`` threshold operates in summed-score units per group, so per-node threshold values must be scaled up to be meaningful. Ignored when ``blocks`` is ``None``; ``"mean"`` is provided for completeness. max_nodes: Optional cap on the number of selected items -- individual targets when @@ -99,9 +101,7 @@ def suggest_exclusion( """ if blocks is not None: if block_agg not in {"sum", "max", "mean"}: - raise ValueError( - f"block_agg must be 'sum', 'max', or 'mean' (got {block_agg!r})" - ) + raise ValueError(f"block_agg must be 'sum', 'max', or 'mean' (got {block_agg!r})") groups = _assign_groups(scores, blocks) group_scores = _aggregate_group_scores(scores, groups, block_agg) selected_groups = _pick_from_scores( @@ -184,7 +184,7 @@ def _assign_groups( Rules: - A node matching any regex in ``blocks[name]`` joins group ``name``. - - First-match wins across the iteration order of ``blocks``, so callers that mix depths list + - First-match wins across the iteration order of ``blocks``, so callers that mix depths list more-specific groups earlier. - Nodes matching no pattern become their own singleton group named after themselves so architecturally-important standalone nodes compete on equal footing with multi-node blocks. @@ -232,7 +232,7 @@ def _warn_near_tie( When the two boundary targets carry nearly equivalent sensitivity but land in different precisions (one FP16, one INT8), the resulting Cast boundary tends to produce intra-group - fragmentation. This warning helps guiding the user into adjusting coverage or threshold + fragmentation. This warning helps guiding the user into adjusting coverage or threshold to include the near-tied target. """ if near_tie_ratio is None: diff --git a/modelopt/onnx/quantization/sensitivity/score.py b/modelopt/onnx/quantization/sensitivity/score.py index 18e4251d71c..9dd4ed4a5d0 100644 --- a/modelopt/onnx/quantization/sensitivity/score.py +++ b/modelopt/onnx/quantization/sensitivity/score.py @@ -29,8 +29,8 @@ import re import tempfile import time -from collections.abc import Callable, Sequence from enum import Enum +from typing import TYPE_CHECKING import numpy as np import onnx @@ -48,6 +48,9 @@ from modelopt.onnx.quantization.sensitivity.metrics import cos_dist, kl_div, mse from modelopt.onnx.utils import gen_random_inputs, get_input_names, get_op_types_in_graph +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + __all__ = ["CalibrationSource", "Granularity", "Metric", "score"] @@ -104,7 +107,8 @@ def _default_op_types_scope(onnx_model: onnx.ModelProto) -> set[str]: """ activation_ops = get_activation_ops() return { - op for op in get_op_types_in_graph(onnx_model) + op + for op in get_op_types_in_graph(onnx_model) if ( is_default_quantizable_op_by_ort(op) or op in activation_ops @@ -191,9 +195,7 @@ def score( f"Unknown metric '{metric}'. Expected one of {list(_METRIC_FUNCS.keys())}." ) if granularity not in (Granularity.OP_TYPE.value, Granularity.NODE.value): - raise ValueError( - f"Unknown granularity '{granularity}'. Expected 'op_type' or 'node'." - ) + raise ValueError(f"Unknown granularity '{granularity}'. Expected 'op_type' or 'node'.") if target_precision not in ("int8", "fp8"): raise ValueError( f"Unsupported target_precision '{target_precision}'. Expected 'int8' or 'fp8'." @@ -210,9 +212,7 @@ def score( f"target_precision={target_precision}" ) - quantizable_ops = ( - set(op_types_scope) if op_types_scope else _default_op_types_scope(onnx_model) - ) + quantizable_ops = set(op_types_scope) if op_types_scope else _default_op_types_scope(onnx_model) if granularity == Granularity.OP_TYPE.value: targets = _enumerate_op_type_targets(onnx_model, quantizable_ops) else: @@ -317,9 +317,7 @@ def _resolve_calibration_data( return _stack_sample_list(list(calibration_data)), CalibrationSource.REAL -def _load_calibration_from_path( - path: str, input_names: list[str] -) -> dict[str, np.ndarray]: +def _load_calibration_from_path(path: str, input_names: list[str]) -> dict[str, np.ndarray]: """Load real calibration data from ``.npy``, ``.npz``, or a directory of ``.npz`` files. Args: diff --git a/tests/gpu/onnx/quantization/test_sensitivity.py b/tests/gpu/onnx/quantization/test_sensitivity.py index ad726f3d6a0..36f2d36b009 100644 --- a/tests/gpu/onnx/quantization/test_sensitivity.py +++ b/tests/gpu/onnx/quantization/test_sensitivity.py @@ -90,12 +90,8 @@ def _build_conv_mm_ln_onnx(path: str, opset: int = 17) -> None: pads=[1, 1, 1, 1], strides=[1, 1], ), - helper.make_node( - "Flatten", ["conv2_out"], ["flat_out"], name="flatten_1", axis=1 - ), - helper.make_node( - "MatMul", ["flat_out", "mm_w"], ["mm_out"], name="matmul_1" - ), + helper.make_node("Flatten", ["conv2_out"], ["flat_out"], name="flatten_1", axis=1), + helper.make_node("MatMul", ["flat_out", "mm_w"], ["mm_out"], name="matmul_1"), helper.make_node( "LayerNormalization", ["mm_out", "ln_scale", "ln_bias"], @@ -109,19 +105,11 @@ def _build_conv_mm_ln_onnx(path: str, opset: int = 17) -> None: graph = helper.make_graph( nodes=nodes, name="sens_test_graph", - inputs=[ - helper.make_tensor_value_info( - _INPUT_NAME, TensorProto.FLOAT, [1, _C_IN, _H, W] - ) - ], - outputs=[ - helper.make_tensor_value_info(_OUTPUT_NAME, TensorProto.FLOAT, [1, _LOGITS]) - ], + inputs=[helper.make_tensor_value_info(_INPUT_NAME, TensorProto.FLOAT, [1, _C_IN, _H, W])], + outputs=[helper.make_tensor_value_info(_OUTPUT_NAME, TensorProto.FLOAT, [1, _LOGITS])], initializer=initializers, ) - model = helper.make_model( - graph, opset_imports=[helper.make_opsetid("", opset)], ir_version=8 - ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)], ir_version=8) onnx.save(model, path) @@ -227,9 +215,7 @@ def test_coatnet_op_type_matches_manual_groundtruth(): * ``coatnet-0_rw_inpsize_1x3x224x224_opsetv_17_simplified.onnx`` -- baseline ONNX. * ``imagenet_calib_500.npz`` -- 500-sample ImageNet calibration dict. """ - onnx_path = _require_fixture( - "coatnet-0_rw_inpsize_1x3x224x224_opsetv_17_simplified.onnx" - ) + onnx_path = _require_fixture("coatnet-0_rw_inpsize_1x3x224x224_opsetv_17_simplified.onnx") calib_path = _require_fixture("imagenet_calib_500.npz") result = score( @@ -256,9 +242,7 @@ def test_coatnet_op_type_matches_manual_groundtruth(): # These cluster at ~0 -- primitive won't recommend excluding them because there's # nothing to exclude. for op in ("Softmax", "Gemm", "GlobalAveragePool"): - assert scores.get(op, 0.0) < 0.001, ( - f"{op} score {scores.get(op, 0.0):.3g} should be ~0" - ) + assert scores.get(op, 0.0) < 0.001, f"{op} score {scores.get(op, 0.0):.3g} should be ~0" @pytest.mark.slow_gpu @@ -267,9 +251,7 @@ def test_coatnet_per_node_matches_manual_groundtruth(): Wall clock ~30-60 min; gated behind ``@pytest.mark.slow_gpu`` so default CI stays fast. """ - onnx_path = _require_fixture( - "coatnet-0_rw_inpsize_1x3x224x224_opsetv_17_simplified.onnx" - ) + onnx_path = _require_fixture("coatnet-0_rw_inpsize_1x3x224x224_opsetv_17_simplified.onnx") calib_path = _require_fixture("imagenet_calib_500.npz") result = score( diff --git a/tests/unit/onnx/quantization/test_nodes_to_quantize.py b/tests/unit/onnx/quantization/test_nodes_to_quantize.py index 893b1133004..36b5545034e 100644 --- a/tests/unit/onnx/quantization/test_nodes_to_quantize.py +++ b/tests/unit/onnx/quantization/test_nodes_to_quantize.py @@ -92,7 +92,9 @@ def test_nodes_to_quantize_restricts_qdq_to_single_conv(tmp_path): """`nodes_to_quantize=["conv_keep"]` inserts Q/DQ around conv_keep only.""" onnx_path = str(tmp_path / "two_conv.onnx") _build_two_conv_onnx(onnx_path) - calibration_data = {"input": np.random.default_rng(0).standard_normal((2, 3, 8, 8)).astype(np.float32)} + calibration_data = { + "input": np.random.default_rng(0).standard_normal((2, 3, 8, 8)).astype(np.float32) + } moq.quantize( onnx_path, @@ -109,8 +111,12 @@ def test_nodes_to_quantize_restricts_qdq_to_single_conv(tmp_path): graph = gs.import_onnx(onnx.load(quantized_path)) keep_nodes = [n for n in graph.nodes if n.name == "conv_keep"] skip_nodes = [n for n in graph.nodes if n.name == "conv_skip"] - assert len(keep_nodes) == 1, f"conv_keep not found in quantized graph: {[n.name for n in graph.nodes]}" - assert len(skip_nodes) == 1, f"conv_skip not found in quantized graph: {[n.name for n in graph.nodes]}" + assert len(keep_nodes) == 1, ( + f"conv_keep not found in quantized graph: {[n.name for n in graph.nodes]}" + ) + assert len(skip_nodes) == 1, ( + f"conv_skip not found in quantized graph: {[n.name for n in graph.nodes]}" + ) # conv_keep must have DQ on its activation input; conv_skip must not. assert _has_dq_predecessor(keep_nodes[0], 0), ( diff --git a/tests/unit/onnx/quantization/test_sensitivity_picker.py b/tests/unit/onnx/quantization/test_sensitivity_picker.py index a4d025372e1..c472f3325b1 100644 --- a/tests/unit/onnx/quantization/test_sensitivity_picker.py +++ b/tests/unit/onnx/quantization/test_sensitivity_picker.py @@ -17,10 +17,7 @@ import pytest -from modelopt.onnx.quantization.sensitivity.picker import ( - suggest_exclusion, - summarize_exclusion, -) +from modelopt.onnx.quantization.sensitivity.picker import suggest_exclusion, summarize_exclusion class TestCoverageMode: @@ -135,10 +132,36 @@ class TestNearTieWarning: def test_warning_fires_on_near_tied_cutoff(self, caplog): # Ranks 16 and 17 are near-tied at KL 3.06 vs 3.05 (99.7% ratio); coverage=0.94 # cuts between them. - scores = {f"node_{i:02d}": kl for i, kl in enumerate( - [6.7, 5.7, 4.6, 4.3, 4.1, 4.1, 4.0, 4.0, 3.8, 3.8, - 3.8, 3.8, 3.7, 3.7, 3.7, 3.06, 3.05, 0.8, 0.5, 0.1], 1)} + scores = { + f"node_{i:02d}": kl + for i, kl in enumerate( + [ + 6.7, + 5.7, + 4.6, + 4.3, + 4.1, + 4.1, + 4.0, + 4.0, + 3.8, + 3.8, + 3.8, + 3.8, + 3.7, + 3.7, + 3.7, + 3.06, + 3.05, + 0.8, + 0.5, + 0.1, + ], + 1, + ) + } import logging + with caplog.at_level(logging.WARNING, logger="modelopt.onnx"): suggest_exclusion(scores, coverage=0.94) messages = [r.message for r in caplog.records] @@ -146,9 +169,11 @@ def test_warning_fires_on_near_tied_cutoff(self, caplog): def test_no_warning_when_cut_is_not_a_near_tie(self, caplog): # ViT-like distribution where coverage=0.75 cuts between very different KL values. - scores = {f"node_{i:02d}": kl for i, kl in enumerate( - [6.7, 5.7, 4.6, 4.3, 4.1, 0.5, 0.2, 0.1], 1)} + scores = { + f"node_{i:02d}": kl for i, kl in enumerate([6.7, 5.7, 4.6, 4.3, 4.1, 0.5, 0.2, 0.1], 1) + } import logging + with caplog.at_level(logging.WARNING, logger="modelopt.onnx"): suggest_exclusion(scores, coverage=0.75) messages = [r.message for r in caplog.records] @@ -158,6 +183,7 @@ def test_warning_disabled_by_none(self, caplog): # Setting near_tie_ratio=None disables the warning entirely. scores = {"a": 5.0, "b": 4.99, "c": 0.1} import logging + with caplog.at_level(logging.WARNING, logger="modelopt.onnx"): suggest_exclusion(scores, coverage=0.5, near_tie_ratio=None) messages = [r.message for r in caplog.records] @@ -167,6 +193,7 @@ def test_threshold_mode_also_warns_on_near_tie(self, caplog): # threshold=3.056 cuts between KL 3.06 (above threshold) and 3.05 (below) -- near-tie. scores = {"a": 6.7, "b": 5.7, "c": 3.06, "d": 3.05, "e": 0.1} import logging + with caplog.at_level(logging.WARNING, logger="modelopt.onnx"): suggest_exclusion(scores, threshold=3.056) messages = [r.message for r in caplog.records] From 64335e08f6ca5737057403d89586edd345665020 Mon Sep 17 00:00:00 2001 From: gcunhase <4861122+gcunhase@users.noreply.github.com> Date: Mon, 24 Aug 2026 19:26:39 +0000 Subject: [PATCH 5/8] docs/picker/tests: post-review simplifications across the sensitivity module Review-driven cleanup after the block-picker + ruff-format commits. Zero API changes, zero behavior changes; documentation, docstring, and test-scaffolding trims only. CHANGELOG.rst ============= Move the ``modelopt.onnx.quantization.sensitivity`` entry from the Megatron subsection to Quantization (where it belongs) and mention the optional block-level aggregation (``blocks=`` / ``block_agg=``) so the line covers the block-picker feature that landed alongside the primitive. docs/source/guides/_onnx_quantization.rst ========================================= * Opening (128-138): compressed 11 lines of marketing-toned framing ("runs into the same friction ... hand-crafted exclusion policies until they find one that works", "Works across CNN, Transformer, and hybrid architectures alike") down to a 4-line spec of what ``sensitivity.score`` does. * Supported-options bullets (142-161): trimmed the parenthetical restatements from ``metric`` / ``calibration_method`` / ``calibration_data`` / ``op_types_scope`` -- the details are in the ``score`` docstring. * Synthetic-calibration note: dropped the "random Q times K^T produces near-uniform softmax that hides real-input MHA quantization pathology" overexplanation; the "directional-only; attention-heavy is the highest-risk case" punchline is enough. * Per-node granularity paragraph: deleted (the ``granularity`` bullet above already covers it). * Meta-narration: dropped "In the rest of this documentation, we'll cover per-node granularity for simplicity, but the same logic goes for per-op-type granularity." * Coverage / threshold bullets (274-286): 13-line bulleted restatement of ``suggest_exclusion``'s docstring cut to 6 lines pointing the reader at the docstring for full argument reference. * Block-picker intro (333-350): dropped "and softmax numerics degrade catastrophically" hedge + ``block_agg`` docstring parenthetical. * Near-tie note: 7-line bulleted restatement of the picker docstring compressed to 3 lines with the actual guidance in one sentence. * Empirically-recovers phrasing: replaced "closing the ViT-tiny parity gap to implicit quantization" (jargon) with the direct number comparison "~75% top-1 versus ~60% for the best per-node picking". * When per-block picking doesn't help: dropped "typically" and "Reach for ``blocks`` first" chatty wording. * CoAtNet timm-export prep: dropped the "Use the analogous timm handle for any other model family" filler sentence. modelopt/onnx/quantization/sensitivity/picker.py ================================================ * Module docstring: 8 lines -> 1 sentence ("Turn a sensitivity score dictionary into an exclusion list, with optional block-level aggregation."). * ``suggest_exclusion`` docstring: 55 lines -> 32 lines. Consolidated the coverage/threshold trade-off restatement (previously in the mode summary + the ``coverage`` arg + a free-standing paragraph); the standalone paragraph moved to the RST guide. ``block_agg`` collapsed from a 12-line bullet forest to a 4-line paragraph -- the natural pairing table (``sum`` with ``coverage``, ``max`` with ``threshold``) is preserved. * ``_aggregate_group_scores``: dropped ``if members else 0.0`` defensive ternary. ``groups`` is built by ``setdefault().append()`` so every key is guaranteed at least one member. * ``_warn_near_tie``: renamed local variables from ``last_included_kl`` / ``first_excluded_kl`` to ``last_included_score`` / ``first_excluded_score`` -- this module supports MSE and cos too, so the ``_kl`` suffix was metric-specific and misleading. Rewrote the awkward "helps guiding the user into adjusting coverage or threshold" docstring line as "The warning prompts widening ``coverage`` or narrowing ``threshold``". * ``_pick_from_scores``: dropped ``# Threshold mode`` / ``# Coverage mode`` block comments -- the ``if threshold is not None:`` branch is self-labeling. * ``summarize_exclusion``: dropped ``float(...)`` cast on ``scores.get(...)`` -- ``scores`` is already ``Mapping[str, float]``. Collapsed the 12-line bulleted Returns block to a single paragraph. modelopt/onnx/quantization/sensitivity/score.py =============================================== * Module docstring: 7 lines -> 4 lines ("Core sensitivity primitive: rank quantizable targets by per-target Q/DQ drift."). * ``_default_op_types_scope`` docstring: 11 lines -> 5 lines. Kept the load-bearing rationale (copy ops excluded because TRT never produces INT8 kernels for them) but dropped the narrative "clutters the output with 'don't do this anyway' entries". * ``op_types_scope`` argument in ``score()``'s docstring: removed the CLI-specific "hides those from the pretty-printed table by default" aside -- that behavior is documented in ``__main__.py``. modelopt/onnx/quantization/sensitivity/__init__.py ================================================== Module docstring: 8 lines -> 1 sentence ("ONNX quantization sensitivity: rank quantizable targets by per-target Q/DQ drift."). modelopt/onnx/quantization/sensitivity/__main__.py ================================================== * Removed the defensive ``assert result["calibration_source"] in {c.value for c in CalibrationSource}`` along with its "score() already emits the enum's string value" comment -- ``score()`` sets that field from the enum's own ``.value`` so the assert was checking that ``score()`` doesn't lie about its own contract. ``CalibrationSource`` is no longer imported here (still ships in the public API via ``__init__.py``). * ``show_zero_scores`` docstring: 4-line "graph plumbing" explanation trimmed to a one-line statement of what the flag does. * ``--show_zero_scores`` CLI help: same treatment. * ``_load_calibration`` docstring: dropped "matches what the main quantize CLI does" background parenthetical. modelopt/onnx/quantization/sensitivity/metrics.py ================================================= * Module docstring: 8 lines -> 1 sentence ("Proxy metrics between reference and quantized activations. Higher = more distortion."). * ``kl_div`` / ``mse`` / ``cos_dist`` docstrings: property-first. Each metric now leads with its scale-sensitivity property ("Robust to activation magnitude scale.", "Sensitive to activation magnitude scale.", "Scale-invariant.") instead of comparative narration ("recommended default because ...", "a target whose output happens to be large in absolute value will look more sensitive under MSE than under KL / cosine", etc.). tests/gpu/onnx/quantization/test_sensitivity.py ================================================ * Added a module-scoped ``coatnet_fixtures`` fixture that ``pytest.skip``\ s cleanly when the pre-staged CoAtNet-0 ONNX + ``imagenet_calib_500.npz`` are absent. Both integration tests now receive the tuple instead of duplicating the ``_require_fixture`` calls. * Trimmed ``test_coatnet_op_type_matches_manual_groundtruth``'s docstring: the 13-row op-type ranking table is already documented in the RST guide, so the test docstring just states the top-4 assertion invariant. * Renumbered the tiers **by cost**: Tier 1 is now the fast synthetic-random regression guard, Tier 2 is the fast synthetic-real deterministic test, Tier 3 is the CoAtNet op-type integration, and Tier 4 is the CoAtNet per-node integration. The previous numbering interleaved a fast fallback test as Tier 4 after two slow real-model tests, which was harder to scan. The two synthetic tests are also reordered in the source file to match. All docstring ``Tier N:`` labels, the module-header tier list, and the ``coatnet_fixtures`` fixture docstring reference ("for tier 3 / 4 tests") were updated in the same pass. * Fixed a stale env-var reference in the same trimmed docstring: ``MODELOPT_SENSITIVITY_FIXTURES`` never existed; the test reads ``MODELOPT_ONNX_ACCURACY_MODELS_DIR`` (see line 48). tests/unit/onnx/quantization/test_sensitivity_picker.py ======================================================== Hoisted ``import logging`` to module scope -- the four ``TestNearTieWarning`` tests each imported it inside their body. Net -3 lines / +1 line, and the ``caplog.at_level(logging.WARNING, ...)`` calls now reference the module-level module correctly. Validation ========== ``ruff check`` + ``ruff format`` clean on all touched files. ``mypy`` clean on the five source files. ``pytest`` still runs on the fork's own env (Computelab); locally we only sanity-check ruff + mypy. Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com> Co-Authored-By: modelopt-fix-agent-bot (Opus 4.7) --- CHANGELOG.rst | 2 +- docs/source/guides/_onnx_quantization.rst | 114 ++++++----------- .../onnx/quantization/sensitivity/__init__.py | 9 +- .../onnx/quantization/sensitivity/__main__.py | 24 +--- .../onnx/quantization/sensitivity/metrics.py | 25 +--- .../onnx/quantization/sensitivity/picker.py | 120 ++++++------------ .../onnx/quantization/sensitivity/score.py | 27 ++-- .../gpu/onnx/quantization/test_sensitivity.py | 108 +++++++--------- .../quantization/test_sensitivity_picker.py | 10 +- 9 files changed, 148 insertions(+), 291 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 44dfbbcec56..77f234e426a 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -9,6 +9,7 @@ Changelog *Quantization* - 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 ``modelopt.onnx.quantization.sensitivity`` — per-op-type or per-node accuracy sensitivity ranking for ONNX PTQ, plus a coverage or threshold-based exclusion picker (with optional block-level aggregation) that turns the ranking into an actionable ``--nodes_to_exclude`` or ``--op_types_to_exclude`` list. - 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. *Megatron Framework (M-LM / M-Bridge)* @@ -16,7 +17,6 @@ Changelog - Add SFT-masked data support to ``examples/megatron_bridge/distill.py``: ``--sft --sft_dataset_root `` distills on raw prompt-completion JSONL (``{"input", "output"}`` records) with the loss masked to the response tokens, using Megatron-Bridge's ``FinetuningDatasetConfig`` and the model's own HuggingFace tokenizer instead of the pretraining ``GPTDataset`` and ``NullTokenizer``. - Add per-expert weight quantization for Transformer Engine ``TEGroupedMLP`` (fused MoE experts): each expert now has its own ``weight_quantizer`` (a ``GroupedQuantizer`` holding one ``TensorQuantizer`` per expert) with an independent ``amax``, instead of a single shared ``amax`` across all experts. Applies to ``mtq.quantize`` calibration, HF / Megatron export, and QAD. - Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager. -- Add ``modelopt.onnx.quantization.sensitivity`` — per-op-type or per-node accuracy sensitivity ranking for ONNX PTQ, plus coverage or threshold-based exclusion picker that turns the ranking into an actionable ``--nodes_to_exclude`` or ``--op_types_to_exclude`` list, depending on granularity. *Misc* diff --git a/docs/source/guides/_onnx_quantization.rst b/docs/source/guides/_onnx_quantization.rst index c232682bf7c..e4738081a29 100644 --- a/docs/source/guides/_onnx_quantization.rst +++ b/docs/source/guides/_onnx_quantization.rst @@ -125,16 +125,9 @@ The following command will build the engine using fp16 precision. After building Quantization Sensitivity Scan ============================= -Post-training quantization of any ONNX model often runs into the same friction: it is unclear -which ops or nodes destroy accuracy at INT8/FP8, and practitioners iterate through hand-crafted -exclusion policies until they find one that works. The -:func:`modelopt.onnx.quantization.sensitivity.score` primitive automates that investigation for -any ONNX model with a calibration dataset. It ranks quantizable targets (op types or individual -nodes) by a proxy metric between the reference and per-target quantized activations, so a -downstream picker can decide which ops to keep at higher precision. Works across CNN, -Transformer, and hybrid architectures alike -- the ranking reflects each model's own -precision-sensitive pathways (residual paths, normalization boundaries, SE gating, attention -projections, etc.) without any architecture-specific configuration. The primitive reuses +:func:`modelopt.onnx.quantization.sensitivity.score` ranks each quantizable target (op type or +individual node) by a proxy metric between the reference and per-target quantized activations, +so a downstream picker can decide which targets to keep at higher precision. It reuses :func:`modelopt.onnx.quantization.quantize` internally for each per-target probe. .. _sensitivity-supported-options: @@ -144,21 +137,15 @@ Supported options - ``granularity``: ``op_type`` (default; probes each quantizable op type once) or ``node`` (probes each ONNX node individually; slower). -- ``metric``: ``kl_div`` (default; softmax-normalized KL divergence — recommended), ``mse`` - (raw mean squared error; cheaper but scale-sensitive) or ``cos`` (``1 - cosine_similarity``; - scale-invariant, robust to activation magnitude variance). +- ``metric``: ``kl_div`` (default), ``mse``, or ``cos`` (``1 - cosine_similarity``). - ``target_precision``: ``int8`` (default) or ``fp8``. -- ``calibration_method``: passed through to the underlying quantize call — ``entropy`` (default) - or ``max``. +- ``calibration_method``: ``entropy`` (default) or ``max``. - ``calibration_data``: sequence of input-dicts, path to real data (``.npy`` / ``.npz`` / - directory), or ``None`` to fall back to synthetic random tensors (directional-only; see note - below). -- ``op_types_scope``: optional whitelist of op types to probe. If omitted, defaults to the - intersection of ops actually present in the graph and the union of ORT's default quantizable - set, activation ops, normalization ops, and fusible reduction ops. Graph plumbing (``Cast`` / - ``Constant`` / ``Shape`` / ...) is skipped so wall-clock is not wasted on zero-drift probes. - Any ops that slip past the filter but still produce zero drift are hidden from the CLI table - by default (pass ``--show_zero_scores`` to see them; they always appear in the JSON). + directory), or ``None`` for synthetic random tensors (directional-only; see note below). +- ``op_types_scope``: optional whitelist of op types to probe. If omitted, defaults to ops + present in the graph intersected with the union of ORT's default quantizable set, activation + ops, normalization ops, and fusible reduction ops (graph plumbing like ``Cast`` / + ``Constant`` / ``Shape`` is skipped). Python API: @@ -169,8 +156,8 @@ Python API: result = score( onnx_path="coatnet-0.onnx", calibration_data="imagenet_calib_500.npz", - granularity="op_type", # choices = {"op_type", "node"} - metric="kl_div", # choices = {"kl_div", "mse", "cos"} + granularity="op_type", + metric="kl_div", target_precision="int8", ) # result["scores"] is a dict {op_type_or_node_name: metric_value}, higher = more sensitive. @@ -205,10 +192,6 @@ from timm's ``coatnet_0_rw_224.sw_in1k`` (``pretrained=True``), the code looks l np.savez("imagenet_calib_500.npz", **{input_name: np.stack(samples).astype(np.float32)}) -Use the analogous timm handle for any other model family (``resnet50``, ``mobilenetv3_large_100``, -``vit_base_patch16_224``, ...); the ``resolve_model_data_config`` +``create_transform`` pair keeps -preprocessing consistent with the exported ONNX regardless of architecture. - Command line: .. code-block:: bash @@ -248,15 +231,9 @@ Rendered ranking (CoAtNet-0, real 500-sample ImageNet calibration):: .. note:: - Omitting ``--calibration_data_path`` falls back to synthetic random inputs. Absolute scores are - then directional-only and must not be paired with absolute thresholds -- attention-heavy models - are the highest-risk degradation case because random Q times K^T produces near-uniform softmax - that hides real-input MHA quantization pathology. - -In per-node granularity the scanner iterates over every quantizable node in the graph and runs -one probe per node; each probe uses the existing ``--nodes_to_quantize `` flag on the main -quantize CLI to quantize that node alone (everything else stays FP16) so the resulting output -drift attributes to that specific node. + Omitting ``--calibration_data_path`` falls back to synthetic random inputs; scores are + directional-only and must not be paired with absolute thresholds. Attention-heavy models + are the highest-risk degradation case. Turning scores into an exclusion list ------------------------------------- @@ -268,22 +245,16 @@ function :func:`sensitivity.suggest_exclusion` turns that dictionary into an act :func:`modelopt.onnx.quantization.quantize`, and :func:`sensitivity.summarize_exclusion` reports what the exclusion set covers. -In the rest of this documentation, we'll cover ``per-node`` granularity for simplicity, -but the same logic goes for ``per-op-type`` granularity. - Two policy modes are supported: -- **Coverage mode** (default): return the largest node set whose cumulative sensitivity - score stays at or below ``coverage * total_mass``. The actual coverage is always less - than or equal to the requested value ("at most X%"). Architecture-portable because the - target is a fraction, not an absolute number -- ``coverage=0.90`` means the same thing - on any model regardless of sensitivity score magnitudes. -- **Threshold mode**: return every node whose individual sensitivity score exceeds - ``threshold`` (no cumulative-mass logic). Simpler and more predictable when the - operator already knows the per-node sensitivity score magnitude that separates - "quantize safely" from "keep at higher precision" for a specific model. Per-node - sensitivity score magnitudes are not portable across models. When ``threshold`` is - set, ``coverage`` is ignored. +- **Coverage mode** (default): exclude the largest node set whose cumulative sensitivity score + stays at or below ``coverage * total_mass``. Architecture-portable -- ``coverage=0.90`` means + the same thing on any model. +- **Threshold mode**: exclude every node whose individual score exceeds ``threshold``. Simpler + when the operator already knows a per-node cutoff for a specific model. Setting ``threshold`` + ignores ``coverage``. + +See :func:`suggest_exclusion` for the full argument reference. Python API -- coverage mode: @@ -322,32 +293,24 @@ Python API -- threshold mode: .. note:: - The picker emits a ``logger.warning`` when the boundary between included and - excluded nodes is a near-tie -- specifically, if the first-excluded node's sensitivity - score is at least 99% of the last-included node's sensitivity score. In that case two - nodes with nearly equivalent sensitivity end up in different precisions (one FP16, - one INT8), which can produce intra-group precision fragmentation. The warning suggests - a slightly larger ``coverage`` (or smaller ``threshold``) to include the near-tied node. - Set ``near_tie_ratio=None`` to disable the warning entirely. + The picker warns when the exclusion boundary is a near-tie (default: + first-excluded score >= 99% of last-included). Widen ``coverage`` or narrow + ``threshold`` to absorb the near-tied target, or set ``near_tie_ratio=None`` to silence. Grouping per-node scores into architectural blocks -------------------------------------------------- On attention-heavy transformer architectures (ViT, DeiT, Swin, CoAtNet's -attention stages), per-node picking can leave affected transformer blocks with -fragmented precision -- some FP16 nodes, some INT8 nodes -- and softmax -numerics degrade catastrophically. Making the *transformer block* the atomic -exclusion unit avoids the fragmentation. +attention stages), per-node picking can leave transformer blocks with +fragmented precision -- some FP16 nodes, some INT8 nodes. Making the +*transformer block* the atomic exclusion unit avoids the fragmentation. Pass a ``blocks`` mapping to :func:`suggest_exclusion` to switch the picker -from per-node to per-block ranking. Each node in the score dict is assigned -to at most one group (first-match wins across ``blocks``); unmatched nodes -become their own singleton group. Coverage / threshold / near-tie / -``max_nodes`` semantics apply to the *group* ranking, and the returned -exclusion list is the union of member nodes across the selected groups. See -:func:`suggest_exclusion`'s docstring for the ``block_agg`` / picker-mode -pairings (``sum`` + ``coverage`` and ``max`` + ``threshold`` preserve -per-node units). +from per-node to per-block ranking. Each node is assigned to at most one +group (first-match wins across ``blocks``); unmatched nodes become their +own singleton group. Coverage / threshold / near-tie / ``max_nodes`` +semantics apply to the *group* ranking, and the returned exclusion list is +the union of member nodes across the selected groups. Example: ``vit_tiny_patch16_224`` from timm ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -436,9 +399,8 @@ tail than blocks.11), so any of the following expressions produces the same # sum + max_nodes (equivalent -- top 6 groups by cumulative KL mass) suggest_exclusion(scores, coverage=1.0, max_nodes=6, blocks=blocks, block_agg="sum") -Empirically on a 500-image ImageNet-1k validation subset, this 101-node -block-level exclusion recovers ~75% top-1 versus ~60% for the best per-node -picking -- closing the ViT-tiny parity gap to implicit quantization. +On a 500-image ImageNet-1k validation subset, this 101-node block-level +exclusion recovers ~75% top-1 versus ~60% for the best per-node picking. Choosing a grouping depth ~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -474,5 +436,5 @@ When per-block picking doesn't help Block-level grouping is architecture-specific. On Conv-heavy models where sensitivity is diffuse across many small MBConv or Bottleneck contributors (MobileNet, ResNet families), per-node ``coverage`` or ``threshold`` picking -typically outperforms block grouping. Reach for ``blocks`` first on -transformer / attention-heavy architectures. +outperforms block grouping. Use ``blocks`` on transformer / attention-heavy +architectures. diff --git a/modelopt/onnx/quantization/sensitivity/__init__.py b/modelopt/onnx/quantization/sensitivity/__init__.py index da620b5568a..117ad9e5416 100644 --- a/modelopt/onnx/quantization/sensitivity/__init__.py +++ b/modelopt/onnx/quantization/sensitivity/__init__.py @@ -13,14 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""ONNX quantization sensitivity scan. - -Ranks quantization targets (op types or individual nodes) by the accuracy impact they would have if -quantized. The core primitive, :func:`score`, mutates the graph with a properly calibrated single- -target Q/DQ pass (via the standard :func:`modelopt.onnx.quantization.quantize` entry point), runs -both the FP16 reference and the quantized model through ONNXRuntime, and reports a proxy metric per -target so a downstream picker can decide which ops or nodes to keep at higher precision. -""" +"""ONNX quantization sensitivity: rank quantizable targets by per-target Q/DQ drift.""" from modelopt.onnx.quantization.sensitivity.picker import suggest_exclusion, summarize_exclusion from modelopt.onnx.quantization.sensitivity.score import ( diff --git a/modelopt/onnx/quantization/sensitivity/__main__.py b/modelopt/onnx/quantization/sensitivity/__main__.py index 0aefcc213a3..913aa1cd4c3 100644 --- a/modelopt/onnx/quantization/sensitivity/__main__.py +++ b/modelopt/onnx/quantization/sensitivity/__main__.py @@ -29,12 +29,7 @@ import numpy as np from modelopt.onnx.logging_config import logger -from modelopt.onnx.quantization.sensitivity.score import ( - CalibrationSource, - Granularity, - Metric, - score, -) +from modelopt.onnx.quantization.sensitivity.score import Granularity, Metric, score def _default_output_json(onnx_path: str) -> str: @@ -46,9 +41,9 @@ def _default_output_json(onnx_path: str) -> str: def _load_calibration(path: str | None) -> str | dict | None: """Return calibration input for :func:`score`. - If ``path`` is a ``.npz`` file, load it eagerly so the caller sees a proper ``dict`` (matches - what the main quantize CLI does). Directories and ``.npy`` files are passed through as strings so - :func:`score` uses its path-loader. + If ``path`` is a ``.npz`` file, load it eagerly so the caller sees a proper ``dict``. + Directories and ``.npy`` files are passed through as strings so :func:`score` uses its + path-loader. Args: path: Filesystem location or ``None`` for the synthetic-random fallback. @@ -72,9 +67,6 @@ def _render_ranked_table(result: dict, show_zero_scores: bool = False) -> str: Args: result: The return value of :func:`score`. show_zero_scores: If False (default), hide targets whose drift score is exactly ``0.0``. - Such targets typically indicate op types the underlying quantize call skipped (graph - plumbing like ``Cast`` or ``Reshape``); their zero score is legitimate but noisy in - the ranked table. All scores -- including zeros -- always appear in the JSON output. Returns: A newline-joined string with a header, one row per non-hidden target, and highest / lowest @@ -199,10 +191,7 @@ def get_parser() -> argparse.ArgumentParser: parser.add_argument( "--show_zero_scores", action="store_true", - help=( - "Include zero-drift targets (op types the underlying quantize call could not affect) " - "in the stderr ranked table. They always appear in the JSON regardless." - ), + help="Include zero-score targets in the stderr ranked table.", ) return parser @@ -238,9 +227,6 @@ def main(argv: list[str] | None = None) -> int: calibration_eps=args.calibration_eps, op_types_scope=args.op_types_scope, ) - # Sanity-check the JSON schema; score() already emits the enum's string value. - assert result["calibration_source"] in {c.value for c in CalibrationSource} - payload = {"onnx_path": os.path.abspath(args.onnx_path), **result} output_json = args.output_json or _default_output_json(args.onnx_path) os.makedirs(os.path.dirname(os.path.abspath(output_json)) or ".", exist_ok=True) diff --git a/modelopt/onnx/quantization/sensitivity/metrics.py b/modelopt/onnx/quantization/sensitivity/metrics.py index c439b696d83..0955fb87db6 100644 --- a/modelopt/onnx/quantization/sensitivity/metrics.py +++ b/modelopt/onnx/quantization/sensitivity/metrics.py @@ -13,14 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Proxy metrics for ONNX quantization sensitivity scoring. - -Each metric maps a pair ``(fp16_act, quant_act)`` of aligned activation tensors to a non-negative -scalar. Higher values mean the quantization target under test caused more distortion of the model's -output, so the caller ranks targets by increasing sensitivity to decide what to keep at higher -precision. Callers pass the raw activations exactly as ORT returned them; each metric normalizes -internally where relevant (e.g. softmax for KL) and averages across the leading batch dimension. -""" +"""Proxy metrics between reference and quantized activations. Higher = more distortion.""" import numpy as np @@ -55,10 +48,8 @@ def _softmax(logits: np.ndarray, axis: int = -1) -> np.ndarray: def kl_div(fp16_act: np.ndarray, quant_act: np.ndarray) -> float: """KL divergence between softmax-normalized FP16 and quantized activations. - Both tensors are flattened per-sample and passed through softmax to obtain probability - distributions, then the KL divergence ``sum(p * log(p / q))`` is computed per sample and - averaged. This is the recommended default metric because it matches the intuition "output - distribution should be similar" and is robust to activation magnitude scale. + Robust to activation magnitude scale. Both tensors are flattened per-sample, passed through + softmax, and ``sum(p * log(p / q))`` is averaged across samples. Args: fp16_act: FP16 reference activations, shape ``(num_samples, ...)``. @@ -74,10 +65,7 @@ def kl_div(fp16_act: np.ndarray, quant_act: np.ndarray) -> float: def mse(fp16_act: np.ndarray, quant_act: np.ndarray) -> float: - """Mean squared error on raw activation values. - - Cheap to compute but sensitive to activation magnitude scale; a target whose output happens to - be large in absolute value will look more sensitive under MSE than under KL / cosine. + """Mean squared error on raw activation values. Sensitive to activation magnitude scale. Args: fp16_act: FP16 reference activations. @@ -93,10 +81,7 @@ def mse(fp16_act: np.ndarray, quant_act: np.ndarray) -> float: def cos_dist(fp16_act: np.ndarray, quant_act: np.ndarray) -> float: - """Cosine distance ``1 - cos(fp16, quant)`` averaged across samples. - - Scale-invariant: robust to models with wide activation-magnitude variance where MSE would be - dominated by the largest tensors. + """Cosine distance ``1 - cos(fp16, quant)`` averaged across samples. Scale-invariant. Args: fp16_act: FP16 reference activations. diff --git a/modelopt/onnx/quantization/sensitivity/picker.py b/modelopt/onnx/quantization/sensitivity/picker.py index 5098cd9e116..c1e61a49f11 100644 --- a/modelopt/onnx/quantization/sensitivity/picker.py +++ b/modelopt/onnx/quantization/sensitivity/picker.py @@ -13,14 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Exclusion picker for the sensitivity primitive. - -Turns a per-node or per-op-type score dictionary produced by :func:`sensitivity.score` into an -actionable ``--nodes_to_exclude`` or ``--op_types_to_exclude`` list. Supports coverage mode (pick -the largest set whose cumulative score stays at or below ``coverage * total_mass``) and threshold -mode (exclude every target whose individual score exceeds an absolute cutoff), and can optionally -aggregate per-node scores into user-defined architectural groups via the ``blocks`` argument. -""" +"""Turn a sensitivity score dictionary into an exclusion list, with optional block-level aggregation.""" from __future__ import annotations @@ -46,56 +39,33 @@ def suggest_exclusion( ) -> list[str]: """Return an exclusion list from a per-target sensitivity score dictionary. - Two policy modes are supported: - - **Coverage mode** (default) returns the largest target set whose cumulative sensitivity score stays - at or below ``coverage * total_mass`` -- the picker stops *before* crossing the target, so the - actual coverage is always <= requested. - - **Threshold mode** (used when ``threshold`` is set; ``coverage`` is ignored) returns every target - whose individual score strictly exceeds ``threshold``. - - Coverage is architecture-portable because it is a fraction of the model's own total mass; threshold - is model-specific but simpler when the operator already knows a reasonable per-target cutoff. + Coverage mode (default) picks the largest target set whose cumulative score stays at or + below ``coverage * total_mass``. Threshold mode (when ``threshold`` is set; ``coverage`` + is then ignored) picks every target whose individual score exceeds ``threshold``. Args: - scores: Per-target (node or op-type) sensitivity scores from :func:`sensitivity.score`. - coverage: Fraction of total sensitivity score mass to leave unquantized (portable across models). - ``0.85-0.90`` (default) balances accuracy and INT8 latency benefit; ``0.95-0.99`` - favors accuracy; ``0.70-0.80`` favors latency. Portable across models. - threshold: Absolute score cutoff. Every target with score strictly greater than - ``threshold`` is excluded. Magnitudes are model-dependent. - blocks: Optional mapping from group name to a list of regex patterns that match node paths - (i.e, ``{group_name: [regex, ...]}``). When set, the picker ranks *groups* rather than - individual nodes: each node is assigned to at most one group (first-match wins across - the dict), unmatched nodes become their own singleton group, and coverage / threshold / - ``max_nodes`` / ``near_tie_ratio`` semantics apply to the group ranking. The returned - exclusion list is the union of member nodes across the selected groups. - Default ``None`` = per-node picking. - block_agg: Aggregation function used to compute a group's score from its members' individual - scores when ``blocks`` is set. One of ``"sum"``, ``"max"``, ``"mean"``. Natural pairings - with the two policy modes: - - ``block_agg="sum"`` with **coverage** (recommended default): preserves the coverage - semantic regardless of granularity choices (group sums equals to summing all node scores) - - ``block_agg="max"`` with **threshold**: preserves per-node threshold units and operator - intuition when transferring per-node threshold guidance to the block level. - - Other combinations are valid but change what ``coverage`` and ``threshold`` mean in - units. Under ``block_agg="max"`` coverage counts fraction-of-total-group-max-scores - (not fraction-of-total-score-mass). Under ``block_agg="sum"`` threshold operates in - summed-score units per group, so per-node threshold values must be scaled up to be - meaningful. Ignored when ``blocks`` is ``None``; ``"mean"`` is provided for completeness. - max_nodes: Optional cap on the number of selected items -- individual targets when - ``blocks`` is ``None``, or groups when ``blocks`` is set. Prevents long-tail-heavy - distributions from producing very large exclusion sets that fragment the graph and - hurt latency. - min_score_floor: Targets with individual score below this value are never included, even - if the coverage target has not been reached or the target exceeds ``threshold``. - near_tie_ratio: If the first-excluded target's score is at least this fraction of the - last-included target's score, a warning is emitted recommending a slightly larger - coverage / smaller threshold to avoid intra-group precision fragmentation. Set to - ``None`` to disable. Default 0.99. + scores: Per-target sensitivity scores from :func:`sensitivity.score`. + coverage: Fraction of total sensitivity score mass to leave unquantized. Portable + across models. ``0.85-0.90`` (default) balances accuracy and INT8 latency benefit; + ``0.95-0.99`` favors accuracy; ``0.70-0.80`` favors latency. + threshold: Absolute score cutoff. Model-dependent. + blocks: Optional ``{group_name: [regex, ...]}``. When set, ranks *groups* rather than + individual nodes: each node joins at most one group (first-match wins), unmatched + nodes become singleton groups, and all selection semantics apply to the group + ranking. The returned exclusion list is the union of member nodes across selected + groups. + block_agg: Aggregation for group scores when ``blocks`` is set: ``"sum"`` (default; + natural with ``coverage``), ``"max"`` (natural with ``threshold``; preserves + per-node units), or ``"mean"``. Off-diagonal combinations change what ``coverage`` + and ``threshold`` mean in units. + max_nodes: Optional cap on the number of selected items. Prevents long-tail + distributions from producing large exclusion sets that fragment the graph. + min_score_floor: Targets below this score are never included. + near_tie_ratio: Emit a warning when the first-excluded score is at least this fraction + of the last-included score (default 0.99). ``None`` disables it. Returns: - List of target names sorted highest-to-lowest score. Pass to ``nodes_to_exclude=`` for + Target names sorted highest-to-lowest score. Pass to ``nodes_to_exclude=`` for per-node scores (or when ``blocks`` is set) and to ``op_types_to_exclude=`` for per-op-type scores. """ @@ -142,7 +112,6 @@ def _pick_from_scores( if not ranked: return [] - # Threshold mode if threshold is not None: excluded: list[str] = [] for name, score in ranked: @@ -154,7 +123,6 @@ def _pick_from_scores( _warn_near_tie(ranked, excluded, near_tie_ratio, mode="threshold") return excluded - # Coverage mode total = sum(scores.values()) if total <= 0.0: return [] @@ -216,10 +184,7 @@ def _aggregate_group_scores( if block_agg == "max": return {g: max(scores[n] for n in members) for g, members in groups.items()} # mean - return { - g: (sum(scores[n] for n in members) / len(members)) if members else 0.0 - for g, members in groups.items() - } + return {g: sum(scores[n] for n in members) / len(members) for g, members in groups.items()} def _warn_near_tie( @@ -231,26 +196,25 @@ def _warn_near_tie( """Warn if the last-included and first-excluded scores are within ``near_tie_ratio``. When the two boundary targets carry nearly equivalent sensitivity but land in different - precisions (one FP16, one INT8), the resulting Cast boundary tends to produce intra-group - fragmentation. This warning helps guiding the user into adjusting coverage or threshold - to include the near-tied target. + precisions (one FP16, one INT8), the resulting Cast boundary produces intra-group + fragmentation. The warning prompts widening ``coverage`` or narrowing ``threshold``. """ if near_tie_ratio is None: return if not excluded or len(excluded) >= len(ranked): return - last_included_kl = ranked[len(excluded) - 1][1] - if last_included_kl <= 0.0: + last_included_score = ranked[len(excluded) - 1][1] + if last_included_score <= 0.0: return - first_excluded_name, first_excluded_kl = ranked[len(excluded)] - ratio = first_excluded_kl / last_included_kl + first_excluded_name, first_excluded_score = ranked[len(excluded)] + ratio = first_excluded_score / last_included_score if ratio < near_tie_ratio: return last_included_name = ranked[len(excluded) - 1][0] logger.warning( f"suggest_exclusion (mode={mode}): near-tie at the exclusion cut-off. " - f"Last included target '{last_included_name}' has score={last_included_kl:.5f}, " - f"first excluded target '{first_excluded_name}' has score={first_excluded_kl:.5f} " + f"Last included target '{last_included_name}' has score={last_included_score:.5f}, " + f"first excluded target '{first_excluded_name}' has score={first_excluded_score:.5f} " f"({100.0 * ratio:.2f}% of last-included). " f"Consider a slightly larger coverage / smaller threshold to include the " f"near-tied target and avoid intra-group precision fragmentation." @@ -271,21 +235,13 @@ def summarize_exclusion( excluded: The list of target names that will be excluded from quantization. Returns: - Dict with: - - ``coverage_pct``: Percentage of total sensitivity score mass - captured by the exclusion set. - - ``num_excluded``: Number of targets to exclude from quantization. - - ``num_previously_quantized``: Total number of quantizable targets - the primitive probed (i.e., what would have been quantized - without the exclusion set). - - ``num_remaining_quantized``: How many targets will still be - quantized after the exclusion set is applied. - - ``excluded_mass``: Absolute cumulative sensitivity score - captured by the exclusion set. - - ``total_mass``: Sum of sensitivity scores across every probed target. + Dict with ``coverage_pct`` (percentage of total mass captured by the exclusion set), + ``num_excluded``, ``num_previously_quantized``, ``num_remaining_quantized``, + ``excluded_mass`` (absolute cumulative score), and ``total_mass`` (sum across all + probed targets). """ total_mass = sum(scores.values()) - excluded_mass = sum(float(scores.get(name, 0.0)) for name in excluded) + excluded_mass = sum(scores.get(name, 0.0) for name in excluded) coverage_pct = 100.0 * excluded_mass / total_mass if total_mass > 0.0 else 0.0 return { "coverage_pct": coverage_pct, diff --git a/modelopt/onnx/quantization/sensitivity/score.py b/modelopt/onnx/quantization/sensitivity/score.py index 9dd4ed4a5d0..0946754ba11 100644 --- a/modelopt/onnx/quantization/sensitivity/score.py +++ b/modelopt/onnx/quantization/sensitivity/score.py @@ -13,13 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Core ONNX quantization sensitivity primitive: :func:`score`. +"""Core sensitivity primitive: rank quantizable targets by per-target Q/DQ drift. -For every quantization target (an op type or a single node), inserts calibrated Q/DQ nodes on just -that target via the standard :func:`modelopt.onnx.quantization.quantize` entry point, runs the -resulting ONNX and the unquantized reference through ONNXRuntime on the same calibration inputs, -and computes a proxy metric between the two graph-output activation sets. Higher score means the -target degrades the model more if quantized -- so callers keep high scores at higher precision. +For every op type or node, :func:`score` inserts calibrated Q/DQ on that target only via +:func:`modelopt.onnx.quantization.quantize`, runs both the reference and quantized graphs through +ONNXRuntime, and computes a proxy metric between their outputs. Higher score means the target +degrades the model more if quantized. """ from __future__ import annotations @@ -88,16 +87,11 @@ class CalibrationSource(str, Enum): def _default_op_types_scope(onnx_model: onnx.ModelProto) -> set[str]: - """Return op types worth probing by default: present in the graph AND known-quantizable. + """Return op types worth probing by default. - Intersects the set of op types actually present in the graph with the union of ORT's default - quantizable ops, activation ops, normalization ops, and fusible reduction ops. Layout / copy - ops (``Transpose`` / ``Reshape`` / ``Concat`` / ...) are then excluded via - :func:`is_copy_op` -- they show up in ORT's default quantizable set but their sensitivity - signal reflects Q/DQ insertion at data-movement boundaries rather than any INT8-kernel - trade-off, and TensorRT never actually produces INT8 kernels for them, so ranking them - clutters the output with "don't do this anyway" entries. Graph plumbing (``Cast`` / - ``Constant`` / ``Shape`` / ...) not on any of the above lists is also skipped. + Intersects ops present in the graph with ORT's default quantizable set / activation / + normalization / fusible-reduction ops, minus copy ops (such as Transpose and Reshape) + because TensorRT never produces INT8 kernels for them. Args: onnx_model: Loaded ONNX model to enumerate. @@ -173,8 +167,7 @@ def score( ``Shape`` / ...) is skipped by default because it produces zero-drift probes. Ops that slip past the filter but that the underlying :func:`modelopt.onnx.quantization.quantize` still cannot quantize are reported - with score ``0.0`` -- the CLI hides those from the pretty-printed table by - default but they always appear in the JSON output. + with score ``0.0``. work_dir: Directory to place intermediate per-target quantized ONNX files. Defaults to a fresh temporary directory that is removed after the call returns. diff --git a/tests/gpu/onnx/quantization/test_sensitivity.py b/tests/gpu/onnx/quantization/test_sensitivity.py index 36f2d36b009..d32ff5fafe5 100644 --- a/tests/gpu/onnx/quantization/test_sensitivity.py +++ b/tests/gpu/onnx/quantization/test_sensitivity.py @@ -15,14 +15,14 @@ """Tests for the ONNX quantization sensitivity primitive. -Tiers: +Tiers, from lightest to heaviest: -1. Synthetic-graph unit test with real deterministic inputs -- LayerNorm scores highest. -2. CoAtNet-0 op-type integration (``@pytest.mark.slow`` + real ImageNet calibration). -3. CoAtNet-0 per-node integration (``@pytest.mark.slow_gpu`` + real ImageNet calibration). -4. Synthetic-random calibration regression guard -- LayerNorm still > Conv directionally. +1. Synthetic-random calibration smoke test -- LayerNorm > Conv directionally. +2. Synthetic graph + deterministic real inputs -- LayerNorm scores highest. +3. CoAtNet-0 op-type integration (``@pytest.mark.slow`` + real ImageNet calibration). +4. CoAtNet-0 per-node integration (``@pytest.mark.slow_gpu`` + real ImageNet calibration). -Tiers 2 and 3 read a pre-staged CoAtNet-0 ONNX + calibration ``.npz`` from a fixtures directory +Tiers 3 and 4 read a pre-staged CoAtNet-0 ONNX + calibration ``.npz`` from a fixtures directory resolved via ``MODELOPT_ONNX_ACCURACY_MODELS_DIR`` (default ``/tmp``). Missing fixtures ``pytest.skip`` cleanly. """ @@ -128,9 +128,28 @@ def _assert_ln_over_conv(scores: dict[str, float]) -> None: ) +def test_synthetic_random_calibration_directional(tmp_path): + """Tier 1: with ``calibration_data=None`` -- LN > Conv holds.""" + onnx_path = str(tmp_path / "sens_synth.onnx") + _build_conv_mm_ln_onnx(onnx_path) + + result = score( + onnx_path, + calibration_data=None, + num_synthetic_samples=8, + metric="kl_div", + target_precision="int8", + granularity="op_type", + calibration_eps=("cpu",), + op_types_scope=_SYNTHETIC_OP_SCOPE, + ) + assert result["calibration_source"] == "synthetic" + _assert_ln_over_conv(result["scores"]) + + @pytest.mark.parametrize("metric", ["kl_div", "mse", "cos"]) def test_synthetic_deterministic_ln_highest(tmp_path, metric): - """Tier 1: synthetic graph + deterministic real inputs -- LN scores highest of all ops.""" + """Tier 2: synthetic graph + deterministic real inputs -- LN scores highest of all ops.""" onnx_path = str(tmp_path / "sens_synth.onnx") _build_conv_mm_ln_onnx(onnx_path) calib = _deterministic_calibration() @@ -156,25 +175,6 @@ def test_synthetic_deterministic_ln_highest(tmp_path, metric): _assert_ln_over_conv(scores) -def test_synthetic_random_calibration_directional(tmp_path): - """Tier 4: with ``calibration_data=None``, LN > Conv invariant still holds directionally.""" - onnx_path = str(tmp_path / "sens_synth.onnx") - _build_conv_mm_ln_onnx(onnx_path) - - result = score( - onnx_path, - calibration_data=None, - num_synthetic_samples=8, - metric="kl_div", - target_precision="int8", - granularity="op_type", - calibration_eps=("cpu",), - op_types_scope=_SYNTHETIC_OP_SCOPE, - ) - assert result["calibration_source"] == "synthetic" - _assert_ln_over_conv(result["scores"]) - - def _require_fixture(name: str) -> str: """Return a fixture path or ``pytest.skip`` if it isn't staged on this host.""" path = os.path.join(_FIXTURE_DIR, name) @@ -183,40 +183,29 @@ def _require_fixture(name: str) -> str: return path +@pytest.fixture(scope="module") +def coatnet_fixtures() -> tuple[str, str]: + """CoAtNet-0 baseline ONNX + 500-sample ImageNet calibration for tier 3 / 4 tests.""" + return ( + _require_fixture("coatnet-0_rw_inpsize_1x3x224x224_opsetv_17_simplified.onnx"), + _require_fixture("imagenet_calib_500.npz"), + ) + + @pytest.mark.slow -def test_coatnet_op_type_matches_manual_groundtruth(): - """Tier 2: CoAtNet-0 op-type ranking must surface the ops that ``--op_types_to_quantize +def test_coatnet_op_type_matches_manual_groundtruth(coatnet_fixtures): + """Tier 3: CoAtNet-0 op-type ranking must surface the ops that ``--op_types_to_quantize Conv`` implicitly avoids. - Empirical ranking on CoAtNet-0 with 500-sample ImageNet calibration and ``kl_div``: - - Add 2.848 <-- highest impact - Mul 1.890 - LayerNormalization 1.653 - ReduceMean 1.570 - BatchNormalization 0.355 - Conv 0.181 - AveragePool 0.057 - Sigmoid 0.039 - MatMul 0.015 - Relu ~0 - Softmax ~0 - GlobalAveragePool ~0 - Gemm 0 - - Top-4 = Add / Mul / LayerNormalization / ReduceMean are the load-bearing failures - (residual paths, SE gating + softmax scale, norm boundaries). Conv sits ~10x below - the top-4 and quantizes cleanly, matching the manual "Conv-only wins 82% top-1" - ground truth read as a quantization policy. - - Wall-clock ~14 min on H100 with 500 samples / 13 probes (~60s per probe). - - Fixtures (override root via ``MODELOPT_SENSITIVITY_FIXTURES``): - * ``coatnet-0_rw_inpsize_1x3x224x224_opsetv_17_simplified.onnx`` -- baseline ONNX. - * ``imagenet_calib_500.npz`` -- 500-sample ImageNet calibration dict. + On CoAtNet-0 with 500-sample ImageNet calibration and ``kl_div``, the top-4 are + ``Add`` / ``Mul`` / ``LayerNormalization`` / ``ReduceMean`` (all > 1.5 KL) and + ``Conv`` sits ~10x below, matching the manual "Conv-only wins 82% top-1" ground truth. + Full ranking is documented in :doc:`_onnx_quantization`. + + Wall-clock ~14 min on H100. Fixtures (override root via ``MODELOPT_ONNX_ACCURACY_MODELS_DIR``): + ``coatnet-0_rw_inpsize_1x3x224x224_opsetv_17_simplified.onnx`` and ``imagenet_calib_500.npz``. """ - onnx_path = _require_fixture("coatnet-0_rw_inpsize_1x3x224x224_opsetv_17_simplified.onnx") - calib_path = _require_fixture("imagenet_calib_500.npz") + onnx_path, calib_path = coatnet_fixtures result = score( onnx_path, @@ -246,13 +235,12 @@ def test_coatnet_op_type_matches_manual_groundtruth(): @pytest.mark.slow_gpu -def test_coatnet_per_node_matches_manual_groundtruth(): - """Tier 3: CoAtNet-0 per-node ranking (LN/MHA nodes top, Conv nodes bottom). +def test_coatnet_per_node_matches_manual_groundtruth(coatnet_fixtures): + """Tier 4: CoAtNet-0 per-node ranking (LN/MHA nodes top, Conv nodes bottom). Wall clock ~30-60 min; gated behind ``@pytest.mark.slow_gpu`` so default CI stays fast. """ - onnx_path = _require_fixture("coatnet-0_rw_inpsize_1x3x224x224_opsetv_17_simplified.onnx") - calib_path = _require_fixture("imagenet_calib_500.npz") + onnx_path, calib_path = coatnet_fixtures result = score( onnx_path, diff --git a/tests/unit/onnx/quantization/test_sensitivity_picker.py b/tests/unit/onnx/quantization/test_sensitivity_picker.py index c472f3325b1..085b46683c0 100644 --- a/tests/unit/onnx/quantization/test_sensitivity_picker.py +++ b/tests/unit/onnx/quantization/test_sensitivity_picker.py @@ -15,6 +15,8 @@ """Unit tests for :mod:`modelopt.onnx.quantization.sensitivity.picker`.""" +import logging + import pytest from modelopt.onnx.quantization.sensitivity.picker import suggest_exclusion, summarize_exclusion @@ -160,8 +162,6 @@ def test_warning_fires_on_near_tied_cutoff(self, caplog): 1, ) } - import logging - with caplog.at_level(logging.WARNING, logger="modelopt.onnx"): suggest_exclusion(scores, coverage=0.94) messages = [r.message for r in caplog.records] @@ -172,8 +172,6 @@ def test_no_warning_when_cut_is_not_a_near_tie(self, caplog): scores = { f"node_{i:02d}": kl for i, kl in enumerate([6.7, 5.7, 4.6, 4.3, 4.1, 0.5, 0.2, 0.1], 1) } - import logging - with caplog.at_level(logging.WARNING, logger="modelopt.onnx"): suggest_exclusion(scores, coverage=0.75) messages = [r.message for r in caplog.records] @@ -182,8 +180,6 @@ def test_no_warning_when_cut_is_not_a_near_tie(self, caplog): def test_warning_disabled_by_none(self, caplog): # Setting near_tie_ratio=None disables the warning entirely. scores = {"a": 5.0, "b": 4.99, "c": 0.1} - import logging - with caplog.at_level(logging.WARNING, logger="modelopt.onnx"): suggest_exclusion(scores, coverage=0.5, near_tie_ratio=None) messages = [r.message for r in caplog.records] @@ -192,8 +188,6 @@ def test_warning_disabled_by_none(self, caplog): def test_threshold_mode_also_warns_on_near_tie(self, caplog): # threshold=3.056 cuts between KL 3.06 (above threshold) and 3.05 (below) -- near-tie. scores = {"a": 6.7, "b": 5.7, "c": 3.06, "d": 3.05, "e": 0.1} - import logging - with caplog.at_level(logging.WARNING, logger="modelopt.onnx"): suggest_exclusion(scores, threshold=3.056) messages = [r.message for r in caplog.records] From 56565051d5135169a23a9091e86824929ee83c5f Mon Sep 17 00:00:00 2001 From: gcunhase <4861122+gcunhase@users.noreply.github.com> Date: Mon, 24 Aug 2026 19:46:06 +0000 Subject: [PATCH 6/8] changelog: move sensitivity entry to end of the Quantization section The sensitivity entry was sitting between ``mtq.temporarily_fold_weights`` and ``nvfp4_act_headroom``; convention here is that new-in-release entries append at the end of their subsection so the last-added line is always the newest. Moving the bullet down two positions so its ordering matches the section's convention. No content change to the bullet itself. Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com> Co-Authored-By: modelopt-fix-agent-bot (Opus 4.7) --- CHANGELOG.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 77f234e426a..2690967f23e 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -9,8 +9,8 @@ Changelog *Quantization* - 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 ``modelopt.onnx.quantization.sensitivity`` — per-op-type or per-node accuracy sensitivity ranking for ONNX PTQ, plus a coverage or threshold-based exclusion picker (with optional block-level aggregation) that turns the ranking into an actionable ``--nodes_to_exclude`` or ``--op_types_to_exclude`` list. - 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 ``modelopt.onnx.quantization.sensitivity`` — per-op-type or per-node accuracy sensitivity ranking for ONNX PTQ, plus a coverage or threshold-based exclusion picker (with optional block-level aggregation) that turns the ranking into an actionable ``--nodes_to_exclude`` or ``--op_types_to_exclude`` list. *Megatron Framework (M-LM / M-Bridge)* From a63acde641f2924bd12b423ef8b31794694309c0 Mon Sep 17 00:00:00 2001 From: gcunhase <4861122+gcunhase@users.noreply.github.com> Date: Mon, 24 Aug 2026 20:17:32 +0000 Subject: [PATCH 7/8] test: fix IndexError in test_nodes_to_quantize (graph-input probe + brittle helper) Two-part fix for `test_nodes_to_quantize_restricts_qdq_to_single_conv` crashing with `IndexError: list index out of range` instead of asserting cleanly. 1. Graph shape. The synthetic ONNX put `conv_keep` at the very first Conv, so its input was the graph input tensor with no producer node. Add a leading `Relu` so `conv_keep` sits on an interior tensor -- the shape sensitivity's per-node probe actually hits when it isolates an interior Conv, and the setup ModelOpt's Q/DQ insertion is designed around. 2. `_has_dq_predecessor` helper. `gs.Node.i(input_idx)` calls `self.inputs[input_idx].inputs[0]` under the hood; when the input tensor has an empty producer list (e.g., a graph input) that raises `IndexError`. Rewrite the helper to walk `inp.inputs` explicitly and return `False` when any step of the chain is missing. This makes an unquantized target fail the test with the intended "conv_keep is not quantized" assertion message rather than an opaque crash inside the helper. No production-code change; unit test only. Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com> Co-Authored-By: modelopt-fix-agent-bot (Opus 4.7) --- .../quantization/test_nodes_to_quantize.py | 29 ++++++++++++++----- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/tests/unit/onnx/quantization/test_nodes_to_quantize.py b/tests/unit/onnx/quantization/test_nodes_to_quantize.py index 36b5545034e..f66273babd6 100644 --- a/tests/unit/onnx/quantization/test_nodes_to_quantize.py +++ b/tests/unit/onnx/quantization/test_nodes_to_quantize.py @@ -33,7 +33,12 @@ def _build_two_conv_onnx(path: str, opset: int = 17) -> None: - """Emit a 2-Conv ONNX with the node names the test filters on.""" + """Emit a Relu + 2-Conv ONNX with the node names the test filters on. + + The leading ``Relu`` shifts ``conv_keep`` off the graph-input tensor so its input has a real + producer node (a common shape in real models) and mirrors what sensitivity's per-node probe + hits when it isolates an interior Conv. + """ rng = np.random.default_rng(0) w1 = rng.standard_normal((4, 3, 3, 3)).astype(np.float32) * 0.1 b1 = np.zeros((4,), dtype=np.float32) @@ -41,9 +46,10 @@ def _build_two_conv_onnx(path: str, opset: int = 17) -> None: b2 = np.zeros((4,), dtype=np.float32) nodes = [ + helper.make_node("Relu", ["input"], ["relu_out"], name="pre_relu"), helper.make_node( "Conv", - ["input", "w1", "b1"], + ["relu_out", "w1", "b1"], ["conv_keep_out"], name="conv_keep", pads=[1, 1, 1, 1], @@ -78,14 +84,21 @@ def _build_two_conv_onnx(path: str, opset: int = 17) -> None: def _has_dq_predecessor(node: gs.Node, input_idx: int) -> bool: - """Return True when the input at ``input_idx`` of ``node`` is produced by DequantizeLinear.""" + """Return True when the input at ``input_idx`` of ``node`` is produced by DequantizeLinear. + + Returns False (rather than raising) when the input has no producer (graph input) or when + the producer chain is shorter than expected. + """ inp = node.inputs[input_idx] - if not isinstance(inp, gs.Variable): + if not isinstance(inp, gs.Variable) or not inp.inputs: return False - producer = node.i(input_idx) - if producer and producer.op == "Cast": - producer = producer.i(0) - return bool(producer and producer.op == "DequantizeLinear") + producer = inp.inputs[0] + if producer.op == "Cast": + cast_inp = producer.inputs[0] if producer.inputs else None + if not isinstance(cast_inp, gs.Variable) or not cast_inp.inputs: + return False + producer = cast_inp.inputs[0] + return producer.op == "DequantizeLinear" def test_nodes_to_quantize_restricts_qdq_to_single_conv(tmp_path): From 956bd85b952902a7cb088047ac3985326d882dec Mon Sep 17 00:00:00 2001 From: gcunhase <4861122+gcunhase@users.noreply.github.com> Date: Mon, 24 Aug 2026 20:59:10 +0000 Subject: [PATCH 8/8] test: bump 2-Conv test weights to 16 channels to clear small-Conv auto-exclusion `find_nodes_from_convs_to_exclude` in `modelopt/onnx/quantization/graph_utils.py:1163` silently drops any Conv whose OC and IC are both < 16 (and fail the `%8` fallback rule) from the quantizable set. The previous test used `(4, 3, 3, 3)` and `(4, 4, 3, 3)` weights, so both Convs got appended to `nodes_to_exclude` before the `nodes_to_quantize` allowlist was even evaluated -- our `["^conv_keep$"]` was then filtered to empty at int8.py:247 and no QDQ was inserted, masking the plumbing entirely. Bump both Convs to `(16, 16, 3, 3)` weights, biases to `(16,)`, graph I/O to `[1, 16, 8, 8]`, and calibration data to `(2, 16, 8, 8)`. Convs now pass the size gate cleanly and `nodes_to_quantize=["^conv_keep$"]` inserts QDQ around `conv_keep` only. Also revert the interim `Relu` node that was added while diagnosing an earlier `IndexError` -- that crash was resolved by the previous commit's `_has_dq_predecessor` rewrite (guards empty producer lists), so `conv_keep` can sit directly against the graph input again without crashing. Restores the original two-Conv shape and the original `_build_two_conv_onnx` one-liner docstring. No production-code change; unit test only. Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com> Co-Authored-By: modelopt-fix-agent-bot (Opus 4.7) --- .../quantization/test_nodes_to_quantize.py | 24 +++++++------------ 1 file changed, 9 insertions(+), 15 deletions(-) diff --git a/tests/unit/onnx/quantization/test_nodes_to_quantize.py b/tests/unit/onnx/quantization/test_nodes_to_quantize.py index f66273babd6..3cb19f75208 100644 --- a/tests/unit/onnx/quantization/test_nodes_to_quantize.py +++ b/tests/unit/onnx/quantization/test_nodes_to_quantize.py @@ -33,23 +33,17 @@ def _build_two_conv_onnx(path: str, opset: int = 17) -> None: - """Emit a Relu + 2-Conv ONNX with the node names the test filters on. - - The leading ``Relu`` shifts ``conv_keep`` off the graph-input tensor so its input has a real - producer node (a common shape in real models) and mirrors what sensitivity's per-node probe - hits when it isolates an interior Conv. - """ + """Emit a 2-Conv ONNX with the node names the test filters on.""" rng = np.random.default_rng(0) - w1 = rng.standard_normal((4, 3, 3, 3)).astype(np.float32) * 0.1 - b1 = np.zeros((4,), dtype=np.float32) - w2 = rng.standard_normal((4, 4, 3, 3)).astype(np.float32) * 0.1 - b2 = np.zeros((4,), dtype=np.float32) + w1 = rng.standard_normal((16, 16, 3, 3)).astype(np.float32) * 0.1 + b1 = np.zeros((16,), dtype=np.float32) + w2 = rng.standard_normal((16, 16, 3, 3)).astype(np.float32) * 0.1 + b2 = np.zeros((16,), dtype=np.float32) nodes = [ - helper.make_node("Relu", ["input"], ["relu_out"], name="pre_relu"), helper.make_node( "Conv", - ["relu_out", "w1", "b1"], + ["input", "w1", "b1"], ["conv_keep_out"], name="conv_keep", pads=[1, 1, 1, 1], @@ -73,8 +67,8 @@ def _build_two_conv_onnx(path: str, opset: int = 17) -> None: graph = helper.make_graph( nodes=nodes, name="nodes_to_quantize_test", - inputs=[helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 3, 8, 8])], - outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 4, 8, 8])], + inputs=[helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 16, 8, 8])], + outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 16, 8, 8])], initializer=initializers, ) onnx.save( @@ -106,7 +100,7 @@ def test_nodes_to_quantize_restricts_qdq_to_single_conv(tmp_path): onnx_path = str(tmp_path / "two_conv.onnx") _build_two_conv_onnx(onnx_path) calibration_data = { - "input": np.random.default_rng(0).standard_normal((2, 3, 8, 8)).astype(np.float32) + "input": np.random.default_rng(0).standard_normal((2, 16, 8, 8)).astype(np.float32) } moq.quantize(