diff --git a/CHANGELOG.rst b/CHANGELOG.rst old mode 100755 new mode 100644 index 95e3f12ef58..8a15cadf871 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -13,6 +13,7 @@ Changelog - Add a calibration-free streaming Kimi-K3 converter and checkpoint-mirror recipe for NVFP4 routed experts with ``input_scale=1.0`` and 128x128 block-FP8 KDA/MLA attention weights. The converter operates shard-by-shard on the source checkpoint's packed MXFP4 experts instead of loading the 2.8T model through the in-memory ``hf_ptq.py`` path. - Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference and ``mtq.preserve_quantizer_attributes_context`` for restoring temporary quantizer property and type changes. Temporary folding snapshots affected fake-quant weights on a configurable device and restores them with their quantizer state; retained pre-quant scales are inactive, while shared weights, shared quantizers, and ``SequentialQuantizer`` weights are unsupported. - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Instead of setting the global scale from the largest per-block amax seen during calibration (plain ``max``, which leaves no room above it so any larger activation saturates), it anchors the scale to a low percentile of the per-block amax distribution, leaving the rest of the FP8 block-scale range as headroom: ``amax = max(rho * anchor, upper)``, where ``anchor`` and ``upper`` are the per-block amaxes at ``anchor_percentile`` (default 1) and ``upper_percentile`` (default 99.99; set to 100 to never clip calibration data), and ``rho`` (default 16384) is the headroom factor. Applies only to NVFP4 dynamic-block input quantizers; ``SequentialQuantizer`` activation quantizers raise. Weight scales are an orthogonal axis selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian``), so one recipe can combine a weight calibration with this activation policy in a single pass. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` with only the calibration algorithm swapped and exports a standard NVFP4 checkpoint. +- Add ``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)* diff --git a/docs/source/guides/_onnx_quantization.rst b/docs/source/guides/_onnx_quantization.rst index e4d0c2d93d6..80321333efa 100644 --- a/docs/source/guides/_onnx_quantization.rst +++ b/docs/source/guides/_onnx_quantization.rst @@ -121,3 +121,323 @@ 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 +============================= + +: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: + +Supported options +----------------- + +- ``granularity``: ``op_type`` (default; probes each quantizable op type once) or + ``node`` (probes each ONNX node individually; slower). +- ``metric``: ``kl_div`` (default), ``mse``, or ``cos`` (``1 - cosine_similarity``). +- ``target_precision``: ``int8`` (default) or ``fp8``. +- ``calibration_method``: ``entropy`` (default) or ``max``. +- ``calibration_data``: sequence of input-dicts, path to real data (``.npy`` / ``.npz`` / + 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: + +.. 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", + metric="kl_div", + 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 + + from itertools import islice + + 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 ex in islice(ds, 500)] + np.savez("imagenet_calib_500.npz", + **{input_name: np.stack(samples).astype(np.float32)}) + +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; 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 +------------------------------------- + +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. + +Two policy modes are supported: + +- **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: + +.. 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 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 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 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 +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +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 + + 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", + ) + + # 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"], + threshold=0.1, blocks=blocks, block_agg="max", + ) + 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", + ) + +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 + + scores = result["scores"] # from the score() call above + + # 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") + +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 +~~~~~~~~~~~~~~~~~~~~~~~~~ + +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 + + 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 + ] + +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. On Conv-heavy models where +sensitivity is diffuse across many small MBConv or Bottleneck contributors +(MobileNet, ResNet families), per-node ``coverage`` or ``threshold`` picking +outperforms block grouping. Use ``blocks`` on transformer / attention-heavy +architectures. diff --git a/examples/onnx_ptq/README.md b/examples/onnx_ptq/README.md index 3cee5535a84..c7c26915340 100644 --- a/examples/onnx_ptq/README.md +++ b/examples/onnx_ptq/README.md @@ -229,6 +229,49 @@ python -m modelopt.onnx.quantization \ For more fine-tuned Autotune flags, please refer to the [API guide](https://nvidia.github.io/Model-Optimizer/guides/_onnx_quantization.html) and the [Autotune guide](https://nvidia.github.io/Model-Optimizer/guides/9_autotune.html). +### Recover accuracy with sensitivity-driven node exclusion + +Post-training quantization of ONNX models can result in accuracy degradation, and it is often unclear which ops or nodes are more sensitive to precision lowering. To aid in this debugging, we propose using a sensitivity score function to rank each quantizable target (op type or individual node) by its impact on model output and then using a downstream picker to decide which targets to keep in higher precision. See the [Quantization Sensitivity Scan guide](https://nvidia.github.io/Model-Optimizer/guides/_onnx_quantization.html#quantization-sensitivity-scan) for more details. + +End-to-end workflow: + +```python +from modelopt.onnx.quantization import quantize +from modelopt.onnx.quantization.sensitivity import ( + score, + suggest_exclusion, + summarize_exclusion, +) + +# 1. Rank the quantizable targets by their impact on model output. +result = score( + onnx_path=".onnx", + calibration_data=".npy", + granularity="node", # or "op_type" + metric="kl_div", # or "mse", "cos" + target_precision="int8", +) + +# 2. Turn the ranking into an exclusion list. Coverage mode (default) leaves the +# largest set whose cumulative sensitivity mass stays at or below the requested +# fraction. Threshold mode (`threshold=`) excludes every target whose +# individual score exceeds an absolute cutoff. +excluded = suggest_exclusion(result["scores"], coverage=0.90) +print(summarize_exclusion(result["scores"], excluded)) + +# 3. Quantize with the exclusion applied. Use ``nodes_to_exclude=`` for per-node +# and ``op_types_to_exclude=`` for op-type granularity. +quantize( + onnx_path=".onnx", + quantize_mode="int8", + calibration_data=".npy", + nodes_to_exclude=excluded, + output_path=".sens_excluded.quant.onnx", +) +``` + +An optional `blocks=` / `block_agg=` argument to `suggest_exclusion` ranks entire blocks instead of individual nodes. See the [guide](https://nvidia.github.io/Model-Optimizer/guides/_onnx_quantization.html#grouping-per-node-scores-into-architectural-blocks) for more details. + ## Resources - 📅 [Roadmap](https://github.com/NVIDIA/Model-Optimizer/issues/1699) 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..8c11ea12eaa --- /dev/null +++ b/modelopt/onnx/quantization/sensitivity/__init__.py @@ -0,0 +1,29 @@ +# 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: rank quantizable targets by per-target Q/DQ drift.""" + +# ruff: noqa: F405 +from .picker import * +from .score import * + +__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..732db0e7ecb --- /dev/null +++ b/modelopt/onnx/quantization/sensitivity/__main__.py @@ -0,0 +1,271 @@ +# 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. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys + +from modelopt.onnx.logging_config import logger +from modelopt.onnx.quantization.__main__ import validate_file_size +from modelopt.onnx.quantization.sensitivity.score import Granularity, Metric, score + +# 2 GiB matches the ``--onnx_path`` guard in ``modelopt.onnx.quantization.__main__``. +_ONNX_MAX_SIZE_BYTES = 2 * (1024**3) +# 4 GiB accommodates ImageNet-scale calibration NPZ files. +_CALIB_MAX_SIZE_BYTES = 4 * (1024**3) +# 16 GiB aggregate cap for a directory of .npz shards. +_CALIB_DIR_MAX_TOTAL_BYTES = 16 * (1024**3) + + +def _validate_calibration_dir(path: str) -> None: + """Enforce per-file and aggregate size limits on a directory of ``.npz`` calibration shards. + + The directory loader in :func:`score` concatenates every ``.npz`` in the directory without + bounds, so a directory containing many large shards can exhaust process memory during load. + Cap each shard at ``_CALIB_MAX_SIZE_BYTES`` and the aggregate at + ``_CALIB_DIR_MAX_TOTAL_BYTES``. + + Args: + path: Directory expected to contain one or more ``.npz`` calibration shards. + + Raises: + FileNotFoundError: If ``path`` contains no ``.npz`` files. + ValueError: If any shard or the aggregate exceeds the limit. + """ + import glob + + files = sorted(glob.glob(os.path.join(path, "*.npz"))) + if not files: + raise FileNotFoundError(f"No .npz files found under calibration directory: {path}") + total = 0 + for f in files: + validate_file_size(f, _CALIB_MAX_SIZE_BYTES) + total += os.path.getsize(f) + if total > _CALIB_DIR_MAX_TOTAL_BYTES: + raise ValueError( + f"Aggregate calibration directory size {total} bytes exceeds " + f"{_CALIB_DIR_MAX_TOTAL_BYTES} bytes ({len(files)} shards under {path})." + ) + + +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 _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``. + + 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 and, when applicable, + the number of unprobed / failed targets. + """ + scores = result["scores"] + failed = result.get("failed", []) + header = ( + f"Sensitivity scan ({result['target_precision']} / " + f"{result['metric']} / {result['granularity']}):" + ) + if not scores: + if failed: + return ( + header + f"\n (no scores produced; {len(failed)} target(s) failed to probe -- " + f"see calibration_source / failed in the JSON)" + ) + 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: + footer = f"\n (all {hidden} target(s) scored 0.0 -- pass --show_zero_scores to see them)" + if failed: + footer += f"\n ({len(failed)} additional target(s) failed to probe)" + return header + footer + + 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)" + ) + if failed: + lines.append(f" ({len(failed)} target(s) failed to probe -- see failed in 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 graph outputs.", + ) + 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=["cpu", "cuda:0", "trt"], + 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-score targets in the stderr ranked table.", + ) + 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) + + # Boundary validation on user-supplied paths -- mirrors modelopt.onnx.quantization.__main__. + validate_file_size(args.onnx_path, _ONNX_MAX_SIZE_BYTES) + if args.calibration_data_path is not None: + if os.path.isdir(args.calibration_data_path): + _validate_calibration_dir(args.calibration_data_path) + else: + validate_file_size(args.calibration_data_path, _CALIB_MAX_SIZE_BYTES) + + 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." + ) + + result = score( + onnx_path=args.onnx_path, + calibration_data=args.calibration_data_path, # path -> score() delegates to its loader + 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, + ) + 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..29e0415ad64 --- /dev/null +++ b/modelopt/onnx/quantization/sensitivity/metrics.py @@ -0,0 +1,101 @@ +# 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 between reference-graph and per-target-quantized graph outputs. Higher = more distortion.""" + +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. + + 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, ...)``. + 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. Sensitive to activation magnitude scale. + + 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. + + A check is added in the cases of both the reference and quantized outputs being 0 to + prevent unchanged zero-output probes being perceived as maximally sensitive. + + 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 = np.where(norm > 0, dot / (norm + _EPS), 1.0) + 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..fce3726ae64 --- /dev/null +++ b/modelopt/onnx/quantization/sensitivity/picker.py @@ -0,0 +1,258 @@ +# 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. + +"""Turn a sensitivity score dictionary into an exclusion list, with optional block-level aggregation.""" + +from __future__ import annotations + +import re +from typing import TYPE_CHECKING, Literal + +from modelopt.onnx.logging_config import logger + +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + +__all__ = ["suggest_exclusion", "summarize_exclusion"] + + +def suggest_exclusion( + scores: Mapping[str, float], + 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, +) -> list[str]: + """Return an exclusion list from a per-target sensitivity score dictionary. + + Coverage mode (default) walks targets in descending score order, accumulating them until + the next one would push the cumulative score above ``coverage * total_mass``, at which + point it stops (rank-prefix). Threshold mode (when ``threshold`` is set; ``coverage`` + is then ignored) picks every target whose individual score exceeds ``threshold``. + + Args: + 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: + 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"}: + 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]: + """Coverage / threshold selection on any ``{name: score}`` dict. + + 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 [] + + if threshold is not None: + 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 + + 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: + 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 _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``, 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] + 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.""" + 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) for g, members in groups.items()} + + +def _warn_near_tie( + ranked: list[tuple[str, float]], + excluded: list[str], + near_tie_ratio: float | None, + mode: str, +) -> None: + """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 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_score = ranked[len(excluded) - 1][1] + if last_included_score <= 0.0: + return + 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_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." + ) + + +def summarize_exclusion( + scores: Mapping[str, float], + excluded: list[str], +) -> dict: + """Return a summary dict describing an exclusion set. + + 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 sensitivity scores. + excluded: The list of target names that will be excluded from quantization. + + Returns: + 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). + """ + effective_excluded = {name for name in excluded if name in scores} + total_mass = sum(scores.values()) + excluded_mass = sum(scores[name] for name in effective_excluded) + coverage_pct = 100.0 * excluded_mass / total_mass if total_mass > 0.0 else 0.0 + num_excluded = len(effective_excluded) + return { + "coverage_pct": coverage_pct, + "num_excluded": num_excluded, + "num_previously_quantized": len(scores), + "num_remaining_quantized": len(scores) - num_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..9c2e6f1a889 --- /dev/null +++ b/modelopt/onnx/quantization/sensitivity/score.py @@ -0,0 +1,468 @@ +# 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 sensitivity primitive: rank quantizable targets by per-target Q/DQ drift. + +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 + +import glob +import os +import re +import tempfile +import time +from enum import Enum +from typing import TYPE_CHECKING + +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 + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + +__all__ = ["CalibrationSource", "Granularity", "Metric", "score"] + + +class Metric(str, Enum): + """Proxy metrics between reference-graph and per-target-quantized graph outputs.""" + + 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. + + 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. + + 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: list[str] = ["cpu", "cuda:0", "trt"], + 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 appended + to the returned ``failed`` list so callers can tell "unprobed" from + "quantizing this target is free." + 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. + * ``failed``: list of targets whose probe was NOT recorded in ``scores``. Populated when + :func:`quantize` raised or when the probe ran successfully but inserted zero Q/DQ + nodes (i.e. the underlying quantize path silently declined to quantize this target). + Distinguishing this from ``scores == 0.0`` matters because ``0.0`` means "quantizing + this target is free" while ``failed`` means "we don't know." + * ``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: {calibration_source.value} calibration, {num_samples} samples, " + f"granularity={granularity}, metric={metric}, 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] = {} + failed: list[str] = [] + 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}' " + f"({type(e).__name__}); recording as unprobed." + ) + logger.debug(f"quantize() failure detail for '{target_name}':", exc_info=True) + failed.append(target_name) + continue + + # Distinguish "probe inserted no QDQ" (unprobed) from "probe inserted QDQ and drift + # was zero" (safe to quantize) -- otherwise both look identical as ``scores == 0.0``. + if _count_qdq_nodes(probe_path) == 0: + logger.warning( + f"[{idx}/{len(targets)}] quantize() inserted no Q/DQ nodes for target " + f"'{target_name}' -- recording as unprobed instead of a 0.0 drift score." + ) + failed.append(target_name) + 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, + "failed": failed, + "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" + + +def _count_qdq_nodes(onnx_path: str) -> int: + """Return the number of QuantizeLinear + DequantizeLinear nodes in ``onnx_path``. + + Used to distinguish a probe that ran successfully but inserted zero Q/DQ (silently no-op + because ORT's registry dropped the target op type) from one that inserted real Q/DQ and + happened to produce zero drift. + """ + model = onnx.load(onnx_path, load_external_data=False) + return sum( + 1 for node in model.graph.node if node.op_type in {"QuantizeLinear", "DequantizeLinear"} + ) 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/_test_utils/onnx/quantization/sensitivity/models.py b/tests/_test_utils/onnx/quantization/sensitivity/models.py new file mode 100644 index 00000000000..f60215090e1 --- /dev/null +++ b/tests/_test_utils/onnx/quantization/sensitivity/models.py @@ -0,0 +1,131 @@ +# 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. + +"""Shared fixtures and synthetic-graph builders for the sensitivity test suite.""" + +from __future__ import annotations + +import os + +import numpy as np +import onnx +import pytest +from onnx import TensorProto, helper, numpy_helper + +_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 by tests to constrain +# the scoring scope to the ops actually present in this small graph. +SYNTHETIC_OP_SCOPE = ["Conv", "MatMul", "LayerNormalization"] + + +def build_conv_mm_ln_onnx(path: str, opset: int = 17) -> None: + """Build a 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}" + ) + + +def require_fixture(name: str) -> str: + """Return a fixture path under ``MODELOPT_ONNX_ACCURACY_MODELS_DIR`` or ``pytest.skip``.""" + path = os.path.join(_FIXTURE_DIR, name) + if not os.path.exists(path): + pytest.skip(f"Sensitivity fixture missing: {path}") + return path + + +def get_coatnet_paths() -> tuple[str, str]: + """CoAtNet-0 baseline ONNX + 500-sample ImageNet calibration; ``pytest.skip`` if missing.""" + return ( + require_fixture("coatnet-0_rw_inpsize_1x3x224x224_opsetv_17_simplified.onnx"), + require_fixture("imagenet_calib_500.npz"), + ) diff --git a/tests/gpu/onnx/quantization/sensitivity/test_score.py b/tests/gpu/onnx/quantization/sensitivity/test_score.py new file mode 100644 index 00000000000..79d4bddfb6c --- /dev/null +++ b/tests/gpu/onnx/quantization/sensitivity/test_score.py @@ -0,0 +1,194 @@ +# 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. + + +from __future__ import annotations + +import pytest + +from modelopt.onnx.quantization.sensitivity import score +from tests._test_utils.onnx.quantization.sensitivity.models import ( + SYNTHETIC_OP_SCOPE, + build_conv_mm_ln_onnx, + deterministic_calibration, + get_coatnet_paths, +) + + +@pytest.fixture(scope="module") +def synthetic_onnx_path(tmp_path_factory): + """Build the synthetic 2-Conv + 1-MatMul + 1-LayerNorm graph once per test module.""" + path = str(tmp_path_factory.mktemp("sens_synth") / "sens_synth.onnx") + build_conv_mm_ln_onnx(path) + return path + + +def test_synthetic_random_calibration_directional(synthetic_onnx_path): + """With ``calibration_data=None``, ``LN > Conv`` invariant holds directionally.""" + result = score( + synthetic_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" + scores = result["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(synthetic_onnx_path, metric): + """Synthetic graph + deterministic real inputs -- ``LayerNormalization`` scores highest of all ops.""" + result = score( + synthetic_onnx_path, + calibration_data=deterministic_calibration(), + 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 + top_op = max(result["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 {result['scores']}" + ) + + +def test_failed_probe_is_recorded(synthetic_onnx_path, monkeypatch): + """A probe that inserts no Q/DQ nodes is recorded in ``failed`` and absent from ``scores``.""" + import shutil + + # Patch quantize() in score() to copy the input as-is, so the probe path has no Q/DQ nodes + import modelopt.onnx.quantization.sensitivity.score as score_module + + def _fake_quantize(**kwargs): + shutil.copy(kwargs["onnx_path"], kwargs["output_path"]) + + monkeypatch.setattr(score_module, "quantize", _fake_quantize) + + # Calculate scores + result = score( + synthetic_onnx_path, + calibration_data=deterministic_calibration(), + metric="kl_div", + target_precision="int8", + granularity="op_type", + calibration_eps=["cpu"], + op_types_scope=SYNTHETIC_OP_SCOPE, + ) + assert result["failed"], "Expected failed probes to be surfaced, got empty list" + assert not result["scores"], ( + f"Expected empty scores when every probe fails, got {result['scores']}" + ) + + +def test_failed_probe_records_exceptions(synthetic_onnx_path, monkeypatch): + """A probe whose ``quantize()`` call raises is recorded in ``failed`` and absent from ``scores``.""" + # Patch quantize() in score() to raise an issue, so every probe hits the except branch that + # appends the target to ``failed``. + import modelopt.onnx.quantization.sensitivity.score as score_module + + def _raising_quantize(**kwargs): + raise RuntimeError("Simulated quantize failure") + + monkeypatch.setattr(score_module, "quantize", _raising_quantize) + + result = score( + synthetic_onnx_path, + calibration_data=deterministic_calibration(), + metric="kl_div", + target_precision="int8", + granularity="op_type", + calibration_eps=["cpu"], + op_types_scope=SYNTHETIC_OP_SCOPE, + ) + assert result["failed"], "Expected failed probes to be surfaced, got empty list" + assert not result["scores"], ( + f"Expected empty scores when every probe raises, got {result['scores']}" + ) + + +@pytest.mark.manual(reason="CoAtNet-0 integration; ~14 min on H100, needs pre-staged fixtures") +def test_coatnet_op_type_matches_manual_groundtruth(): + """CoAtNet-0 op-type ranking surfaces the ops that ``--op_types_to_quantize Conv`` avoids. + + Top-4 = ``Add`` / ``Mul`` / ``LayerNormalization`` / ``ReduceMean`` (all > 1.5 KL); ``Conv`` + sits ~10x below. Matches the manual "Conv-only wins 82% top-1" ground truth. Wall-clock + ~14 min on H100. Opt-in via ``pytest --run-manual``. + """ + onnx_path, calib_path = get_coatnet_paths() + + 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}" + ) + assert scores["Conv"] < 0.5, ( + f"Conv score {scores['Conv']:.3f} unexpectedly high (top-4 are all > 1.5)" + ) + 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.manual( + reason="CoAtNet-0 per-node integration; ~30-60 min on H100, needs pre-staged fixtures" +) +def test_coatnet_per_node_matches_manual_groundtruth(): + """CoAtNet-0 per-node ranking: LN / MHA nodes in top-10, individual Conv nodes in bottom-10. + + Wall-clock ~30-60 min on H100. Opt-in via ``pytest --run-manual``. + """ + onnx_path, calib_path = get_coatnet_paths() + + 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) + k = 10 + assert len(ranked) >= 2 * k, "Per-node ranking is unexpectedly short." + top_names = [name for name, _ in ranked[:k]] + bottom_names = [name for name, _ in ranked[-k:]] + assert any("layernorm" in n.lower() or "attn" in n.lower() for n in top_names), ( + f"Expected LN or MHA nodes in top-{k}, got {top_names}" + ) + assert any("conv" in n.lower() for n in bottom_names), ( + f"Expected Conv nodes in bottom-{k}, got {bottom_names}" + ) diff --git a/tests/unit/onnx/quantization/sensitivity/test_metrics.py b/tests/unit/onnx/quantization/sensitivity/test_metrics.py new file mode 100644 index 00000000000..c3f01e21cb6 --- /dev/null +++ b/tests/unit/onnx/quantization/sensitivity/test_metrics.py @@ -0,0 +1,107 @@ +# 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.metrics`.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from modelopt.onnx.quantization.sensitivity.metrics import ( + _flatten_per_sample, + cos_dist, + kl_div, + mse, +) + + +class TestIdenticalInputs: + """All three metrics collapse to (near-)zero on identical inputs.""" + + @pytest.mark.parametrize("metric", [kl_div, mse, cos_dist]) + def test_identical_2d_inputs_score_zero(self, metric): + x = np.random.default_rng(0).standard_normal((4, 8)).astype(np.float32) + assert metric(x, x) == pytest.approx(0.0, abs=1e-6) + + @pytest.mark.parametrize("metric", [kl_div, mse, cos_dist]) + def test_identical_4d_inputs_score_zero(self, metric): + x = np.random.default_rng(0).standard_normal((2, 3, 4, 5)).astype(np.float32) + assert metric(x, x) == pytest.approx(0.0, abs=1e-6) + + +class TestCosDistOrthogonal: + """Orthogonal vectors produce ``cos_dist == 1`` (cosine similarity 0).""" + + def test_orthogonal_vectors(self): + p = np.array([[1.0, 0.0]], dtype=np.float32) + q = np.array([[0.0, 1.0]], dtype=np.float32) + assert cos_dist(p, q) == pytest.approx(1.0, abs=1e-6) + + def test_anti_parallel_vectors(self): + p = np.array([[1.0, 0.0]], dtype=np.float32) + q = np.array([[-1.0, 0.0]], dtype=np.float32) + # cos_sim = -1, so cos_dist = 1 - (-1) = 2. + assert cos_dist(p, q) == pytest.approx(2.0, abs=1e-6) + + def test_both_zero_vectors_return_zero_distance(self): + # A probe whose reference and quantized outputs are both all-zero (e.g. a hard-relu / + # masked-out branch) should score as identical, not maximally sensitive. + p = np.zeros((2, 4), dtype=np.float32) + q = np.zeros((2, 4), dtype=np.float32) + assert cos_dist(p, q) == pytest.approx(0.0, abs=1e-6) + + +class TestScaleSensitivity: + """``mse`` scales with input magnitude; ``cos_dist`` does not; ``kl_div`` is invariant on + softmax outputs regardless of scale.""" + + def test_mse_grows_with_magnitude(self): + rng = np.random.default_rng(0) + base = rng.standard_normal((4, 8)).astype(np.float32) + perturbed = base + 0.1 * rng.standard_normal((4, 8)).astype(np.float32) + small = mse(base, perturbed) + big = mse(10.0 * base, 10.0 * perturbed) + assert big > 50.0 * small, ( + f"MSE should scale ~100x on 10x-larger inputs; got small={small} big={big}" + ) + + def test_cos_dist_is_scale_invariant(self): + rng = np.random.default_rng(0) + base = rng.standard_normal((4, 8)).astype(np.float32) + perturbed = base + 0.1 * rng.standard_normal((4, 8)).astype(np.float32) + small = cos_dist(base, perturbed) + big = cos_dist(10.0 * base, 10.0 * perturbed) + assert small == pytest.approx(big, rel=1e-4) + + +class TestFlattenPerSample: + """``_flatten_per_sample`` reshapes any-rank tensor to ``(num_samples, num_features)``.""" + + def test_1d_treated_as_single_sample(self): + x = np.arange(4, dtype=np.float32) + assert _flatten_per_sample(x).shape == (4, 1) + + def test_2d_passes_through(self): + x = np.zeros((3, 5), dtype=np.float32) + assert _flatten_per_sample(x).shape == (3, 5) + + def test_4d_collapses_feature_dims(self): + x = np.zeros((2, 3, 4, 5), dtype=np.float32) + assert _flatten_per_sample(x).shape == (2, 60) + + def test_0d_becomes_1x1(self): + x = np.float32(3.14) + assert _flatten_per_sample(x).shape == (1, 1) diff --git a/tests/unit/onnx/quantization/sensitivity/test_picker.py b/tests/unit/onnx/quantization/sensitivity/test_picker.py new file mode 100644 index 00000000000..5e86c4a1d38 --- /dev/null +++ b/tests/unit/onnx/quantization/sensitivity/test_picker.py @@ -0,0 +1,264 @@ +# 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 logging + +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): + 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): + scores = {"a": 4.0, "b": 3.0, "c": 2.0, "d": 1.0} + assert suggest_exclusion(scores, coverage=0.8) == ["a", "b"] + + @pytest.mark.parametrize( + ("scores", "expected"), + [ + pytest.param({"a": 1.0, "b": 2.0, "c": 3.0}, ["c", "b", "a"], id="ascending_input"), + pytest.param( + {"low": 0.1, "high": 0.9, "mid": 0.5}, ["high", "mid", "low"], id="unsorted_input" + ), + ], + ) + def test_full_coverage_returns_all_sorted_by_score_desc(self, scores, expected): + # coverage=1.0 -> everything fits and result is sorted by KL desc. + assert suggest_exclusion(scores, coverage=1.0) == expected + + def test_top_node_alone_exceeds_target(self): + scores = {"a": 5.0, "b": 3.0, "c": 2.0} + assert suggest_exclusion(scores, coverage=0.2) == [] + + @pytest.mark.parametrize( + ("scores", "coverage"), + [ + pytest.param({"a": 5.0, "b": 3.0}, 0.0, id="zero_target"), + pytest.param({"a": 0.0, "b": 0.0}, 0.9, id="zero_total"), + pytest.param({}, 0.9, id="empty_scores"), + ], + ) + def test_returns_empty_for_boundary_cases(self, scores, coverage): + assert suggest_exclusion(scores, coverage=coverage) == [] + + def test_max_nodes_caps_exclusion_set(self): + 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): + scores = {"hi_1": 5.0, "hi_2": 4.0, "trivial_1": 0.001, "trivial_2": 0.0001} + assert suggest_exclusion(scores, coverage=1.0, min_score_floor=0.01) == ["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, sharp drop to ~3.05 + # at ranks 16-17, then a long tail. Regression witness for the "at most X%" semantic. + big = {f"top_{i}": 6.7 - i * 0.2 for i in range(15)} + 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) + assert sum(scores[n] for n in result) <= 0.90 * 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} + assert suggest_exclusion(scores, coverage=0.99, threshold=2.5) == ["a", "b", "c"] + + def test_max_nodes_still_caps_threshold_mode(self): + 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): + scores = {"a": 5.0, "b": 0.5, "c": 0.3} + assert suggest_exclusion(scores, threshold=0.1, min_score_floor=1.0) == ["a"] + + +class TestBlocks: + """Tests the ``blocks=`` / ``block_agg=`` block-aware picker.""" + + def test_first_match_wins_across_blocks_iteration_order(self): + # 3 nodes; both group "a" and group "b" would match node "shared" via regex, + # but "a" is listed first -> "shared" joins "a". + scores = {"n_a": 5.0, "shared": 4.0, "n_b": 3.0} + blocks = { + "a": [r"^n_a$", r"^shared$"], + "b": [r"^shared$", r"^n_b$"], + } + result = suggest_exclusion(scores, coverage=1.0, blocks=blocks, block_agg="sum") + # Group "a" carries {n_a, shared} = 9.0; group "b" carries {n_b} = 3.0. + # Both groups included at coverage=1.0. + assert set(result) == {"n_a", "shared", "n_b"} + + def test_unmatched_nodes_become_singleton_groups(self): + # "orphan" matches no group -> becomes its own singleton, ranked on its own score. + scores = {"n_a1": 5.0, "n_a2": 4.0, "orphan": 3.0} + blocks = {"a": [r"^n_a"]} + # Sum aggregation: group "a" = 9.0, "orphan" = 3.0. Coverage=0.75 -> target 9.0. + # Only group "a" fits (9.0 <= 9.0); orphan singleton (3.0) would push cumulative to 12. + excluded = suggest_exclusion(scores, coverage=0.75, blocks=blocks, block_agg="sum") + assert set(excluded) == {"n_a1", "n_a2"} + + def test_sum_aggregation(self): + scores = {"a1": 1.0, "a2": 2.0, "b1": 4.0} + blocks = {"a": [r"^a"], "b": [r"^b"]} + # Sum: group a = 3.0, group b = 4.0. threshold=3.5 -> only b (>3.5) excluded. + assert set(suggest_exclusion(scores, threshold=3.5, blocks=blocks, block_agg="sum")) == { + "b1" + } + + def test_max_aggregation(self): + scores = {"a1": 1.0, "a2": 2.0, "b1": 4.0} + blocks = {"a": [r"^a"], "b": [r"^b"]} + # Max: group a = 2.0, group b = 4.0. threshold=3.5 -> only b excluded. + assert set(suggest_exclusion(scores, threshold=3.5, blocks=blocks, block_agg="max")) == { + "b1" + } + + def test_mean_aggregation(self): + scores = {"a1": 1.0, "a2": 5.0, "b1": 4.0} + blocks = {"a": [r"^a"], "b": [r"^b"]} + # Mean: group a = 3.0, group b = 4.0. threshold=3.5 -> only b excluded. + assert set(suggest_exclusion(scores, threshold=3.5, blocks=blocks, block_agg="mean")) == { + "b1" + } + + def test_invalid_block_agg_raises(self): + with pytest.raises(ValueError, match="block_agg"): + suggest_exclusion( + {"a": 1.0}, + coverage=1.0, + blocks={"g": [r"^a$"]}, + block_agg="invalid", # type: ignore[arg-type] + ) + + def test_return_value_is_union_of_member_nodes(self): + # blocks= returns member nodes across selected groups, not group names. + scores = {"blk0_n1": 5.0, "blk0_n2": 4.0, "blk1_n1": 1.0} + blocks = {"blk0": [r"^blk0_"], "blk1": [r"^blk1_"]} + excluded = suggest_exclusion(scores, threshold=0.5, blocks=blocks, block_agg="max") + # Both groups have max > 0.5 -> both included -> all 3 member nodes in exclusion. + assert set(excluded) == {"blk0_n1", "blk0_n2", "blk1_n1"} + + def test_threshold_max_and_coverage_sum_pick_equivalent_set(self): + # Docs claim: threshold=0.1, block_agg="max" and coverage=1.0, max_nodes=, + # block_agg="sum" pick the same top-K groups on ViT-tiny-shaped data. + # Synthetic: 6 blocks with max KL values decreasing; group max > 0.1 for the first 6. + scores = {} + for i in range(6): + for j in range(3): + scores[f"blk{i}_n{j}"] = (6 - i) * (1.0 if j == 0 else 0.1) + # Add 4 low blocks below threshold + for i in range(6, 10): + for j in range(3): + scores[f"blk{i}_n{j}"] = 0.01 * (10 - i) + blocks = {f"blk{i}": [rf"^blk{i}_"] for i in range(10)} + + via_max = suggest_exclusion(scores, threshold=0.1, blocks=blocks, block_agg="max") + via_sum = suggest_exclusion( + scores, coverage=1.0, max_nodes=6, blocks=blocks, block_agg="sum" + ) + assert set(via_max) == set(via_sum) + + +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): + # b and c are near-tied (3.05/3.06 = 99.7%); coverage=0.75 cuts between them. + scores = {"a": 6.0, "b": 3.06, "c": 3.05, "d": 0.1} + with caplog.at_level(logging.WARNING, logger="modelopt.onnx"): + suggest_exclusion(scores, coverage=0.75) + assert "near-tie at the exclusion cut-off" in caplog.text + + def test_no_warning_when_cut_is_not_a_near_tie(self, caplog): + scores = {"a": 6.0, "b": 3.0, "c": 0.1, "d": 0.05} + with caplog.at_level(logging.WARNING, logger="modelopt.onnx"): + suggest_exclusion(scores, coverage=0.75) + assert "near-tie" not in caplog.text + + def test_warning_disabled_by_none(self, caplog): + scores = {"a": 5.0, "b": 4.99, "c": 0.1} + with caplog.at_level(logging.WARNING, logger="modelopt.onnx"): + suggest_exclusion(scores, coverage=0.5, near_tie_ratio=None) + assert "near-tie" not in caplog.text + + def test_threshold_mode_also_warns_on_near_tie(self, caplog): + scores = {"a": 6.0, "b": 3.06, "c": 3.05, "d": 0.1} + with caplog.at_level(logging.WARNING, logger="modelopt.onnx"): + suggest_exclusion(scores, threshold=3.056) + assert "near-tie" in caplog.text and "mode=threshold" in caplog.text + + +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_are_filtered_out_of_counts(self): + # Filter out unknown or duplicated entries + 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) + # Only "a" counts as excluded; "b" remains quantized + assert summary["num_excluded"] == 1 + assert summary["num_remaining_quantized"] == 1 + + def test_duplicate_excluded_names_counted_once(self): + scores = {"a": 5.0, "b": 5.0} + summary = summarize_exclusion(scores, ["a", "a"]) + assert summary["num_excluded"] == 1 + assert summary["num_remaining_quantized"] == 1 + assert summary["excluded_mass"] == pytest.approx(5.0) diff --git a/tests/unit/onnx/quantization/test_quantize_api.py b/tests/unit/onnx/quantization/test_quantize_api.py index f350d5d89f4..82180bdba01 100644 --- a/tests/unit/onnx/quantization/test_quantize_api.py +++ b/tests/unit/onnx/quantization/test_quantize_api.py @@ -22,7 +22,8 @@ import onnxruntime import pytest import torch -from _test_utils.onnx.lib_test_models import SimpleMLP, export_as_onnx +from _test_utils.onnx.lib_test_models import SimpleMLP, build_conv_concat_model, export_as_onnx +from _test_utils.onnx.quantization.utils import assert_nodes_are_quantized from packaging import version import modelopt.onnx.quantization as moq @@ -196,3 +197,45 @@ def test_quantize_opset_handling( assert output_opset == expected_opset, ( f"[{scenario_name}] Expected opset {expected_opset} for {quant_mode}, got {output_opset}" ) + + +def test_quantize_honors_nodes_to_quantize_allowlist(tmp_path): + """``nodes_to_quantize=[]`` inserts QDQ around the matched Conv only. + + Guards the primitive the ONNX sensitivity scanner relies on to isolate a single node for a + per-target probe; also documents the API contract of the flag itself. + """ + import onnx_graphsurgeon as gs + + from modelopt.onnx.utils import save_onnx + + onnx_model = build_conv_concat_model() + onnx_path = os.path.join(tmp_path, "conv_concat.onnx") + save_onnx(onnx_model, onnx_path) + + # Restrict quantization to the second Conv only (interior node with a real producer input). + keep = "conv2_conv/Conv2D" + moq.quantize( + onnx_path, + quantize_mode="int8", + nodes_to_quantize=[f"^{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)) + conv_nodes = {n.name: n for n in graph.nodes if n.op == "Conv"} + assert keep in conv_nodes, f"{keep} missing after quantization: {list(conv_nodes)}" + + # The selected Conv must have QDQ on its variable inputs; the other three must not. + assert assert_nodes_are_quantized([conv_nodes[keep]]) + for name, node in conv_nodes.items(): + if name == keep: + continue + for inp_idx, inp in enumerate(node.inputs): + if isinstance(inp, gs.Variable) and inp.inputs: + producer = node.i(inp_idx) + assert producer.op != "DequantizeLinear", ( + f"Unselected Conv '{name}' was quantized: input {inp_idx} traces to {producer.op}" + )