diff --git a/CHANGELOG.rst b/CHANGELOG.rst old mode 100755 new mode 100644 index 6687ebd31ea..2690967f23e --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -10,6 +10,7 @@ Changelog - Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference and ``mtq.preserve_quantizer_attributes_context`` for restoring temporary quantizer property and type changes. Temporary folding snapshots affected fake-quant weights on a configurable device and restores them with their quantizer state; retained pre-quant scales are inactive, while shared weights, shared quantizers, and ``SequentialQuantizer`` weights are unsupported. - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Instead of setting the global scale from the largest per-block amax seen during calibration (plain ``max``, which leaves no room above it so any larger activation saturates), it anchors the scale to a low percentile of the per-block amax distribution, leaving the rest of the FP8 block-scale range as headroom: ``amax = max(rho * anchor, upper)``, where ``anchor`` and ``upper`` are the per-block amaxes at ``anchor_percentile`` (default 1) and ``upper_percentile`` (default 99.99; set to 100 to never clip calibration data), and ``rho`` (default 16384) is the headroom factor. Applies only to NVFP4 dynamic-block input quantizers; ``SequentialQuantizer`` activation quantizers raise. Weight scales are an orthogonal axis selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian``), so one recipe can combine a weight calibration with this activation policy in a single pass. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` with only the calibration algorithm swapped and exports a standard NVFP4 checkpoint. +- Add ``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..e4738081a29 100644 --- a/docs/source/guides/_onnx_quantization.rst +++ b/docs/source/guides/_onnx_quantization.rst @@ -121,3 +121,320 @@ 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 + + import numpy as np, onnx, timm, torch + from datasets import load_dataset + from timm.data import resolve_model_data_config, create_transform + + # 1. Export the timm checkpoint to ONNX. + model = timm.create_model("coatnet_0_rw_224.sw_in1k", pretrained=True).eval() + cfg = resolve_model_data_config(model) + dummy = torch.randn(1, *cfg["input_size"]) # (1, 3, 224, 224) + torch.onnx.export( + model, dummy, "coatnet-0.onnx", + input_names=["input"], output_names=["output"], + opset_version=17, + ) + + # 2. Prepare the calibration NPZ with matching preprocessing. + m = onnx.load("coatnet-0.onnx") + input_name = m.graph.input[0].name + tfm = create_transform(**cfg, is_training=False) + ds = load_dataset("ILSVRC/imagenet-1k", split="validation", streaming=True) + samples = [tfm(ex["image"].convert("RGB")).numpy() + for i, ex in enumerate(ds) if i < 500] + np.savez("imagenet_calib_500.npz", + **{input_name: np.stack(samples).astype(np.float32)}) + +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 + + # 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/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..117ad9e5416 --- /dev/null +++ b/modelopt/onnx/quantization/sensitivity/__init__.py @@ -0,0 +1,33 @@ +# 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.""" + +from modelopt.onnx.quantization.sensitivity.picker import suggest_exclusion, summarize_exclusion +from modelopt.onnx.quantization.sensitivity.score import ( + CalibrationSource, + Granularity, + Metric, + score, +) + +__all__ = [ + "CalibrationSource", + "Granularity", + "Metric", + "score", + "suggest_exclusion", + "summarize_exclusion", +] diff --git a/modelopt/onnx/quantization/sensitivity/__main__.py b/modelopt/onnx/quantization/sensitivity/__main__.py new file mode 100644 index 00000000000..913aa1cd4c3 --- /dev/null +++ b/modelopt/onnx/quantization/sensitivity/__main__.py @@ -0,0 +1,242 @@ +# 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 + +import numpy as np + +from modelopt.onnx.logging_config import logger +from modelopt.onnx.quantization.sensitivity.score import Granularity, Metric, score + + +def _default_output_json(onnx_path: str) -> str: + """Derive the default ``--output_json`` path next to the input ONNX file.""" + stem, _ = os.path.splitext(os.path.basename(onnx_path)) + return os.path.join(os.path.dirname(os.path.abspath(onnx_path)), f"{stem}.sensitivity.json") + + +def _load_calibration(path: str | None) -> str | dict | None: + """Return calibration input for :func:`score`. + + If ``path`` is a ``.npz`` file, load it eagerly so the caller sees a proper ``dict``. + Directories and ``.npy`` files are passed through as strings so :func:`score` uses its + path-loader. + + Args: + path: Filesystem location or ``None`` for the synthetic-random fallback. + + Returns: + The value to hand to :func:`score` as ``calibration_data``. + """ + if path is None: + return None + if os.path.isdir(path) or path.endswith(".npy"): + return path + if path.endswith(".npz"): + payload = np.load(path, allow_pickle=False) + return {key: payload[key] for key in payload.files} + raise ValueError(f"Unsupported calibration_data_path: {path}") + + +def _render_ranked_table(result: dict, show_zero_scores: bool = False) -> str: + """Format a sensitivity result as a two-column, high-to-low ranked table. + + Args: + result: The return value of :func:`score`. + show_zero_scores: If False (default), hide targets whose drift score is exactly ``0.0``. + + Returns: + A newline-joined string with a header, one row per non-hidden target, and highest / lowest + markers. A trailing footer notes the count of hidden zero-score rows when applicable. + """ + scores = result["scores"] + header = ( + f"Sensitivity scan ({result['target_precision']} / " + f"{result['metric']} / {result['granularity']}):" + ) + if not scores: + return header + "\n (no quantizable targets found)" + + ranked = sorted(scores.items(), key=lambda kv: kv[1], reverse=True) + hidden = 0 + if not show_zero_scores: + visible = [(n, v) for n, v in ranked if v != 0.0] + hidden = len(ranked) - len(visible) + ranked = visible + + if not ranked: + return ( + header + + f"\n (all {hidden} target(s) scored 0.0 -- pass --show_zero_scores to see them)" + ) + + name_width = max(len(name) for name, _ in ranked) + lines = [header] + for i, (name, value) in enumerate(ranked): + marker = "" + if i == 0: + marker = " <-- highest impact" + elif i == len(ranked) - 1: + marker = " <-- lowest impact" + lines.append(f" {name:<{name_width}} {value:.3f}{marker}") + if hidden: + lines.append( + f" ({hidden} target(s) with score 0.0 hidden; pass --show_zero_scores or read the JSON)" + ) + return "\n".join(lines) + + +def get_parser() -> argparse.ArgumentParser: + """Build the argparse parser for the sensitivity CLI.""" + parser = argparse.ArgumentParser( + prog="modelopt.onnx.quantization.sensitivity", + description=( + "Rank ONNX quantization targets (op types or individual nodes) by their impact on " + "model output. Emits a ranked table to stderr and a JSON file for downstream tooling." + ), + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "--onnx_path", required=True, type=str, help="Path to the input ONNX model." + ) + parser.add_argument( + "--calibration_data_path", + type=str, + default=None, + help=( + "Real calibration data (.npy, .npz, or a directory of .npz files). If omitted, " + "falls back to synthetic random tensors and produces directional-only rankings." + ), + ) + parser.add_argument( + "--num_calib_samples", + type=int, + default=100, + help="Number of synthetic samples generated when --calibration_data_path is omitted.", + ) + parser.add_argument( + "--granularity", + type=str, + default=Granularity.OP_TYPE.value, + choices=[g.value for g in Granularity], + help="Scan granularity: 'op_type' (fast, one probe per type) or 'node' (per-instance).", + ) + parser.add_argument( + "--metric", + type=str, + default=Metric.KL_DIV.value, + choices=[m.value for m in Metric], + help="Proxy metric between FP-reference and quantized activations.", + ) + parser.add_argument( + "--target_precision", + type=str, + default="int8", + choices=["int8", "fp8"], + help="Precision to probe per target.", + ) + parser.add_argument( + "--calibration_method", + type=str, + default="entropy", + choices=["entropy", "max"], + help="Calibration method threaded through to quantize().", + ) + parser.add_argument( + "--calibration_eps", + type=str, + nargs="+", + default=["cuda:0", "cpu"], + help="ORT execution providers, in priority order.", + ) + parser.add_argument( + "--op_types_scope", + type=str, + nargs="+", + default=None, + help=( + "Optional whitelist of op types to probe. Defaults to every unique op type actually " + "present in the ONNX graph." + ), + ) + parser.add_argument( + "--output_json", + type=str, + default=None, + help="Where to write the sensitivity JSON. Defaults to .sensitivity.json.", + ) + parser.add_argument( + "--show_zero_scores", + action="store_true", + help="Include zero-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) + + if args.calibration_data_path is None: + logger.warning( + "Synthetic random calibration -- scores are directional-only; do not pair with " + "absolute thresholds. See calibration_source in the output JSON." + ) + + calibration_data = _load_calibration(args.calibration_data_path) + + result = score( + onnx_path=args.onnx_path, + calibration_data=calibration_data, + num_synthetic_samples=args.num_calib_samples, + target_precision=args.target_precision, + granularity=args.granularity, + metric=args.metric, + calibration_method=args.calibration_method, + calibration_eps=args.calibration_eps, + op_types_scope=args.op_types_scope, + ) + 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..0955fb87db6 --- /dev/null +++ b/modelopt/onnx/quantization/sensitivity/metrics.py @@ -0,0 +1,98 @@ +# 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 and quantized activations. 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. + + Args: + fp16_act: FP16 reference activations. + quant_act: Activations from the quantized model with the same shape as ``fp16_act``. + + Returns: + Mean cosine distance across the sample axis, as a Python float in ``[0, 2]``. + """ + p = _flatten_per_sample(fp16_act).astype(np.float64) + q = _flatten_per_sample(quant_act).astype(np.float64) + dot = np.sum(p * q, axis=-1) + norm = np.linalg.norm(p, axis=-1) * np.linalg.norm(q, axis=-1) + cos = dot / (norm + _EPS) + return float(np.mean(1.0 - cos)) diff --git a/modelopt/onnx/quantization/sensitivity/picker.py b/modelopt/onnx/quantization/sensitivity/picker.py new file mode 100644 index 00000000000..c1e61a49f11 --- /dev/null +++ b/modelopt/onnx/quantization/sensitivity/picker.py @@ -0,0 +1,253 @@ +# 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 + + +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) picks the largest target set whose cumulative score stays at or + below ``coverage * total_mass``. Threshold mode (when ``threshold`` is set; ``coverage`` + is then ignored) picks every target whose individual score exceeds ``threshold``. + + Args: + scores: Per-target 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). + """ + total_mass = sum(scores.values()) + excluded_mass = sum(scores.get(name, 0.0) for name in excluded) + coverage_pct = 100.0 * excluded_mass / total_mass if total_mass > 0.0 else 0.0 + return { + "coverage_pct": coverage_pct, + "num_excluded": len(excluded), + "num_previously_quantized": len(scores), + "num_remaining_quantized": len(scores) - len(excluded), + "excluded_mass": excluded_mass, + "total_mass": total_mass, + } diff --git a/modelopt/onnx/quantization/sensitivity/score.py b/modelopt/onnx/quantization/sensitivity/score.py new file mode 100644 index 00000000000..0946754ba11 --- /dev/null +++ b/modelopt/onnx/quantization/sensitivity/score.py @@ -0,0 +1,434 @@ +# 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 FP16 and quantized activations.""" + + KL_DIV = "kl_div" + MSE = "mse" + COS = "cos" + + +class Granularity(str, Enum): + """Enumeration granularity for sensitivity targets.""" + + OP_TYPE = "op_type" + NODE = "node" + + +class CalibrationSource(str, Enum): + """Origin of the calibration data used for scoring.""" + + REAL = "real" + SYNTHETIC = "synthetic" + + +_METRIC_FUNCS: dict[str, Callable[[np.ndarray, np.ndarray], float]] = { + Metric.KL_DIV.value: kl_div, + Metric.MSE.value: mse, + Metric.COS.value: cos_dist, +} + +# Fixed seed for the synthetic-random calibration fallback so that repeated invocations produce +# identical inputs and, therefore, comparable rankings within one machine. +_SYNTHETIC_SEED = 0 + + +def _default_op_types_scope(onnx_model: onnx.ModelProto) -> set[str]: + """Return op types worth probing by default. + + 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: Sequence[str] = ("cuda:0", "cpu"), + op_types_scope: Sequence[str] | None = None, + work_dir: str | None = None, +) -> dict: + """Rank quantization targets by their impact on model output. + + Runs one reference forward pass over calibration data on the unquantized ``onnx_path``, then for + each target (op type or node) invokes :func:`modelopt.onnx.quantization.quantize` to insert + calibrated Q/DQ nodes on just that target, re-runs the model, and computes ``metric`` between + the reference and quantized graph outputs. Scores are summed across output tensors and averaged + across the calibration samples inside each metric function; higher score means more accuracy + loss if the target is quantized. + + Args: + onnx_path: Path to the ONNX model to score. The model is treated as the FP-precision + reference and is quantized once per target below. + calibration_data: Calibration inputs. Accepts a ``dict[str, np.ndarray]`` (batch-first), + a ``Sequence[dict[str, np.ndarray]]`` of single-sample dicts, a raw ``np.ndarray`` + (single-input models only), or a path to real data on disk (``.npy`` file, ``.npz`` + file, or directory of ``.npz`` files). Passing ``None`` falls back to synthetic random + tensors of the ONNX's declared input shapes. Synthetic random calibration produces + directional rankings only; see :class:`CalibrationSource` in the returned dict. + num_synthetic_samples: Number of synthetic samples generated when + ``calibration_data is None``. Ignored otherwise. + target_precision: Quantization mode passed through to + :func:`modelopt.onnx.quantization.quantize` for each per-target probe. Supported values + are ``"int8"`` and ``"fp8"``. + granularity: ``"op_type"`` scores each quantizable op type once (one probe per type); + ``"node"`` scores each individual quantizable node (one probe per node), which is + substantially more expensive but pinpoints single-node offenders. + metric: One of :class:`Metric` values -- ``"kl_div"`` (default), ``"mse"``, or ``"cos"``. + calibration_method: Passed through to :func:`modelopt.onnx.quantization.quantize` (defaults + to ``"entropy"`` for int8/fp8). + calibration_eps: ONNXRuntime execution providers to use for both the reference and the + per-target forward passes, and for calibration inside :func:`quantize`. Same schema as + the ``--calibration_eps`` CLI flag. + op_types_scope: Optional whitelist of op types to probe. If omitted, defaults to the + intersection of ops present in ``onnx_path`` and the union of ORT's default + quantizable set, activation ops, normalization ops, and fusible reduction ops + (see :func:`_default_op_types_scope`). Graph plumbing (``Cast`` / ``Constant`` / + ``Shape`` / ...) is skipped by default because it produces zero-drift probes. + Ops that slip past the filter but that the underlying + :func:`modelopt.onnx.quantization.quantize` still cannot quantize are reported + with score ``0.0``. + work_dir: Directory to place intermediate per-target quantized ONNX files. Defaults to a + fresh temporary directory that is removed after the call returns. + + Returns: + A dict with keys: + + * ``scores``: mapping of ``op_type`` (op-type granularity) or ``node_name`` (node + granularity) to the summed metric across graph outputs. + * ``calibration_source``: ``"real"`` if the caller supplied calibration data, ``"synthetic"`` + when the primitive fell back to random tensors. + * ``num_calibration_samples``: number of samples used for the scoring pass. + * ``metric``: the metric name as passed in. + * ``granularity``: ``"op_type"`` or ``"node"``. + * ``target_precision``: the requested quantization precision. + """ + if metric not in _METRIC_FUNCS: + raise ValueError( + f"Unknown metric '{metric}'. Expected one of {list(_METRIC_FUNCS.keys())}." + ) + if granularity not in (Granularity.OP_TYPE.value, Granularity.NODE.value): + raise ValueError(f"Unknown granularity '{granularity}'. Expected 'op_type' or 'node'.") + if target_precision not in ("int8", "fp8"): + raise ValueError( + f"Unsupported target_precision '{target_precision}'. Expected 'int8' or 'fp8'." + ) + + onnx_model = onnx.load(onnx_path) + calib_dict, calibration_source = _resolve_calibration_data( + onnx_model, calibration_data, num_synthetic_samples + ) + num_samples = _num_samples(calib_dict) + logger.info( + f"Sensitivity scan on {onnx_path}: {calibration_source.value} calibration, " + f"{num_samples} samples, granularity={granularity}, metric={metric}, " + f"target_precision={target_precision}" + ) + + quantizable_ops = set(op_types_scope) if op_types_scope else _default_op_types_scope(onnx_model) + if granularity == Granularity.OP_TYPE.value: + targets = _enumerate_op_type_targets(onnx_model, quantizable_ops) + else: + targets = _enumerate_node_targets(onnx_model, quantizable_ops) + if not targets: + logger.warning("No quantizable targets found under the requested scope.") + + metric_fn = _METRIC_FUNCS[metric] + calibration_eps_list = list(calibration_eps) + ref_outputs = _run_inference(onnx_path, calib_dict, calibration_eps_list) + + scores: dict[str, float] = {} + use_tempdir = work_dir is None + tmp_ctx = tempfile.TemporaryDirectory() if use_tempdir else None + target_dir = tmp_ctx.name if tmp_ctx is not None else work_dir + assert target_dir is not None + try: + os.makedirs(target_dir, exist_ok=True) + wall_start = time.monotonic() + for idx, (target_name, quantize_kwargs) in enumerate(targets, start=1): + probe_path = os.path.join( + target_dir, f"probe_{_sanitize_filename(target_name)}.quant.onnx" + ) + step_start = time.monotonic() + try: + quantize( + onnx_path=onnx_path, + quantize_mode=target_precision, + calibration_data=calib_dict, + calibration_method=calibration_method, + calibration_eps=calibration_eps_list, + output_path=probe_path, + # Keep non-quantized ops at fp32 to avoid I/O dtype drift between the reference + # and quantized graphs -- the metric then reflects pure Q/DQ distortion. + high_precision_dtype="fp32", + keep_intermediate_files=False, + **quantize_kwargs, + ) + except Exception as e: + logger.warning( + f"[{idx}/{len(targets)}] quantize() failed for target '{target_name}': {e}" + ) + continue + quant_outputs = _run_inference(probe_path, calib_dict, calibration_eps_list) + scores[target_name] = _pair_metric(ref_outputs, quant_outputs, metric_fn) + logger.info( + f"[{idx}/{len(targets)}] scored '{target_name}' = {scores[target_name]:.6g} " + f"(step {time.monotonic() - step_start:.1f}s, total {time.monotonic() - wall_start:.1f}s)" + ) + finally: + if tmp_ctx is not None: + tmp_ctx.cleanup() + + return { + "scores": scores, + "calibration_source": calibration_source.value, + "num_calibration_samples": num_samples, + "metric": metric, + "granularity": granularity, + "target_precision": target_precision, + } + + +def _resolve_calibration_data( + onnx_model: onnx.ModelProto, + calibration_data: ( + Sequence[dict[str, np.ndarray]] | dict[str, np.ndarray] | np.ndarray | str | None + ), + num_synthetic_samples: int, +) -> tuple[dict[str, np.ndarray], CalibrationSource]: + """Normalize any accepted calibration input into a batch-first ``dict[str, ndarray]``. + + Args: + onnx_model: Loaded ONNX model, used to resolve input names and shapes when the caller + passes an ``ndarray`` (single-input models) or ``None`` (synthetic fallback). + calibration_data: One of the forms documented on :func:`score`. + num_synthetic_samples: Number of synthetic samples to generate when ``calibration_data`` is + ``None``. + + Returns: + A tuple ``(calib_dict, source)`` where ``calib_dict`` has each input as a batch-first + numpy array and ``source`` is either ``CalibrationSource.REAL`` or + ``CalibrationSource.SYNTHETIC``. + """ + input_names = get_input_names(onnx_model) + if calibration_data is None: + # np.random is used inside gen_random_inputs; reseed here so the fallback is deterministic + # across invocations on the same model. + np.random.seed(_SYNTHETIC_SEED) + samples = [gen_random_inputs(onnx_model) for _ in range(num_synthetic_samples)] + return _stack_sample_list(samples), CalibrationSource.SYNTHETIC + if isinstance(calibration_data, str): + return _load_calibration_from_path(calibration_data, input_names), CalibrationSource.REAL + if isinstance(calibration_data, np.ndarray): + assert len(input_names) == 1, ( + "ndarray calibration_data is only valid for single-input models." + ) + return {input_names[0]: calibration_data}, CalibrationSource.REAL + if isinstance(calibration_data, dict): + return {k: np.asarray(v) for k, v in calibration_data.items()}, CalibrationSource.REAL + # Sequence[dict] + return _stack_sample_list(list(calibration_data)), CalibrationSource.REAL + + +def _load_calibration_from_path(path: str, input_names: list[str]) -> dict[str, np.ndarray]: + """Load real calibration data from ``.npy``, ``.npz``, or a directory of ``.npz`` files. + + Args: + path: Filesystem location. + input_names: ONNX input names, used to attach ``.npy`` arrays to the sole input. + + Returns: + Batch-first ``dict[str, ndarray]``. + """ + if os.path.isdir(path): + files = sorted(glob.glob(os.path.join(path, "*.npz"))) + assert files, f"No .npz files found under directory {path}" + parts: dict[str, list[np.ndarray]] = {} + for f in files: + payload = np.load(f, allow_pickle=False) + for key in payload.files: + parts.setdefault(key, []).append(payload[key]) + return {k: np.concatenate(v, axis=0) for k, v in parts.items()} + if path.endswith(".npz"): + payload = np.load(path, allow_pickle=False) + return {key: payload[key] for key in payload.files} + if path.endswith(".npy"): + arr = np.load(path, allow_pickle=False) + assert len(input_names) == 1, ( + f"{path} is a single-tensor .npy but the model has {len(input_names)} inputs." + ) + return {input_names[0]: arr} + raise ValueError(f"Unsupported calibration_data path: {path}") + + +def _stack_sample_list(samples: Sequence[dict[str, np.ndarray]]) -> dict[str, np.ndarray]: + """Concatenate a sequence of single-sample dicts into one batch-first dict.""" + assert samples, "Empty calibration sample sequence." + keys = list(samples[0].keys()) + return {k: np.concatenate([np.asarray(s[k]) for s in samples], axis=0) for k in keys} + + +def _num_samples(calib_dict: dict[str, np.ndarray]) -> int: + """Return the batch-axis length of the first array in ``calib_dict``.""" + first = next(iter(calib_dict.values())) + return int(first.shape[0]) + + +def _enumerate_op_type_targets( + onnx_model: onnx.ModelProto, quantizable_ops: set[str] +) -> list[tuple[str, dict]]: + """Return one probe per op type present in the model and in ``quantizable_ops``. + + Args: + onnx_model: Loaded model to enumerate. + quantizable_ops: Whitelist of op types considered quantizable. + + Returns: + List of ``(op_type, quantize_kwargs)`` pairs where ``quantize_kwargs`` restricts + :func:`quantize` to that op type only. + """ + present = {node.op_type for node in onnx_model.graph.node} + scoped = sorted(present & quantizable_ops) + return [(op, {"op_types_to_quantize": [op]}) for op in scoped] + + +def _enumerate_node_targets( + onnx_model: onnx.ModelProto, quantizable_ops: set[str] +) -> list[tuple[str, dict]]: + """Return one probe per named quantizable node. + + Args: + onnx_model: Loaded model to enumerate. + quantizable_ops: Whitelist of op types considered quantizable. + + Returns: + List of ``(node_name, quantize_kwargs)`` pairs where ``quantize_kwargs`` restricts + :func:`quantize` to a regex matching that node only. + """ + targets: list[tuple[str, dict]] = [] + for node in onnx_model.graph.node: + if node.op_type not in quantizable_ops or not node.name: + continue + regex = f"^{re.escape(node.name)}$" + targets.append((node.name, {"nodes_to_quantize": [regex]})) + return targets + + +def _run_inference( + onnx_path: str, calib_dict: dict[str, np.ndarray], calibration_eps: list[str] +) -> list[np.ndarray]: + """Run every sample through ORT and stack outputs along the batch axis. + + Args: + onnx_path: ONNX file to load into an ORT ``InferenceSession``. + calib_dict: Batch-first input dict. + calibration_eps: ORT execution providers, same schema as + :func:`quantize`'s ``calibration_eps``. + + Returns: + List of numpy arrays, one per graph output, each shaped ``(num_samples, ...)``. + """ + session = create_inference_session(onnx_path, calibration_eps) + num_output = len(session.get_outputs()) + num_samples = _num_samples(calib_dict) + per_output: list[list[np.ndarray]] = [[] for _ in range(num_output)] + for i in range(num_samples): + feed = {name: arr[i : i + 1] for name, arr in calib_dict.items()} + outputs = session.run(None, feed) + for j, out in enumerate(outputs): + per_output[j].append(np.asarray(out)) + return [np.concatenate(chunks, axis=0) for chunks in per_output] + + +def _pair_metric( + ref_outputs: list[np.ndarray], + quant_outputs: list[np.ndarray], + metric_fn: Callable[[np.ndarray, np.ndarray], float], +) -> float: + """Sum the metric across matched graph outputs of the reference and quantized models.""" + return float(sum(metric_fn(ref, quant) for ref, quant in zip(ref_outputs, quant_outputs))) + + +def _sanitize_filename(name: str) -> str: + """Turn an arbitrary op/node name into a filesystem-safe token.""" + return re.sub(r"[^A-Za-z0-9._-]", "_", name)[:80] or "unnamed" diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py index f8b5a41a41a..70c1d8b001c 100644 --- a/modelopt/onnx/utils.py +++ b/modelopt/onnx/utils.py @@ -316,6 +316,18 @@ def get_tensor_by_name( return tensor_val or tensor_init or tensor_inp or tensor_out +def get_op_types_in_graph(onnx_model: onnx.ModelProto) -> set[str]: + """Return the set of unique op types that appear as nodes in the graph. + + Args: + onnx_model: Loaded ONNX model. + + Returns: + Set of unique op-type strings appearing in ``onnx_model.graph.node``. + """ + return {node.op_type for node in onnx_model.graph.node if node.op_type} + + def gen_random_inputs( model: onnx.ModelProto, shapes_spec: str | None = None ) -> dict[str, np.ndarray]: diff --git a/tests/gpu/onnx/quantization/test_sensitivity.py b/tests/gpu/onnx/quantization/test_sensitivity.py new file mode 100644 index 00000000000..d32ff5fafe5 --- /dev/null +++ b/tests/gpu/onnx/quantization/test_sensitivity.py @@ -0,0 +1,267 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the ONNX quantization sensitivity primitive. + +Tiers, from lightest to heaviest: + +1. Synthetic-random calibration smoke test -- LayerNorm > Conv directionally. +2. Synthetic graph + deterministic real inputs -- LayerNorm scores highest. +3. CoAtNet-0 op-type integration (``@pytest.mark.slow`` + real ImageNet calibration). +4. CoAtNet-0 per-node integration (``@pytest.mark.slow_gpu`` + real ImageNet calibration). + +Tiers 3 and 4 read a pre-staged CoAtNet-0 ONNX + calibration ``.npz`` from a fixtures directory +resolved via ``MODELOPT_ONNX_ACCURACY_MODELS_DIR`` (default ``/tmp``). Missing fixtures ``pytest.skip`` +cleanly. +""" + +from __future__ import annotations + +import os + +import numpy as np +import onnx +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.quantization.sensitivity import score + +_INPUT_NAME = "input" +_OUTPUT_NAME = "output" +_C_IN = 8 +_C_MID = 16 +_H = W = 16 +_MATMUL_DIM = _C_MID * _H * W +_LOGITS = 32 +_FIXTURE_DIR = os.environ.get("MODELOPT_ONNX_ACCURACY_MODELS_DIR", "/tmp") +# Ops covered by the synthetic Conv+MatMul+LN graph. Passed explicitly because the score() +# default -- get_autotuner_quantizable_ops() -- excludes LayerNormalization even though the ModelOpt +# quantize() path registers it via configure_ort. +_SYNTHETIC_OP_SCOPE = ["Conv", "MatMul", "LayerNormalization"] + + +def _build_conv_mm_ln_onnx(path: str, opset: int = 17) -> None: + """Build a small 2-Conv + 1-MatMul + 1-LayerNorm ONNX for deterministic sensitivity tests.""" + rng = np.random.default_rng(0) + w1 = rng.standard_normal((_C_MID, _C_IN, 3, 3)).astype(np.float32) * 0.1 + b1 = np.zeros((_C_MID,), dtype=np.float32) + w2 = rng.standard_normal((_C_MID, _C_MID, 3, 3)).astype(np.float32) * 0.1 + b2 = np.zeros((_C_MID,), dtype=np.float32) + mm = rng.standard_normal((_MATMUL_DIM, _LOGITS)).astype(np.float32) * 0.05 + ln_scale = np.ones((_LOGITS,), dtype=np.float32) + ln_bias = np.zeros((_LOGITS,), dtype=np.float32) + + initializers = [ + numpy_helper.from_array(w1, "w1"), + numpy_helper.from_array(b1, "b1"), + numpy_helper.from_array(w2, "w2"), + numpy_helper.from_array(b2, "b2"), + numpy_helper.from_array(mm, "mm_w"), + numpy_helper.from_array(ln_scale, "ln_scale"), + numpy_helper.from_array(ln_bias, "ln_bias"), + ] + + nodes = [ + helper.make_node( + "Conv", + ["input", "w1", "b1"], + ["conv1_out"], + name="conv_1", + pads=[1, 1, 1, 1], + strides=[1, 1], + ), + helper.make_node( + "Conv", + ["conv1_out", "w2", "b2"], + ["conv2_out"], + name="conv_2", + pads=[1, 1, 1, 1], + strides=[1, 1], + ), + helper.make_node("Flatten", ["conv2_out"], ["flat_out"], name="flatten_1", axis=1), + helper.make_node("MatMul", ["flat_out", "mm_w"], ["mm_out"], name="matmul_1"), + helper.make_node( + "LayerNormalization", + ["mm_out", "ln_scale", "ln_bias"], + [_OUTPUT_NAME], + name="layernorm_1", + axis=-1, + epsilon=1e-5, + ), + ] + + graph = helper.make_graph( + nodes=nodes, + name="sens_test_graph", + inputs=[helper.make_tensor_value_info(_INPUT_NAME, TensorProto.FLOAT, [1, _C_IN, _H, W])], + outputs=[helper.make_tensor_value_info(_OUTPUT_NAME, TensorProto.FLOAT, [1, _LOGITS])], + initializer=initializers, + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)], ir_version=8) + onnx.save(model, path) + + +def _deterministic_calibration(num_samples: int = 8) -> dict[str, np.ndarray]: + """Fixed-seed calibration data for the synthetic sensitivity graph.""" + rng = np.random.default_rng(42) + return {_INPUT_NAME: rng.standard_normal((num_samples, _C_IN, _H, W)).astype(np.float32)} + + +def _assert_ln_over_conv(scores: dict[str, float]) -> None: + """Directional invariant: LayerNormalization must rank strictly above Conv.""" + assert "LayerNormalization" in scores, f"LayerNorm missing from scores: {scores}" + assert "Conv" in scores, f"Conv missing from scores: {scores}" + assert scores["LayerNormalization"] > scores["Conv"], ( + f"Expected LayerNormalization > Conv, got {scores}" + ) + + +def test_synthetic_random_calibration_directional(tmp_path): + """Tier 1: with ``calibration_data=None`` -- LN > Conv holds.""" + onnx_path = str(tmp_path / "sens_synth.onnx") + _build_conv_mm_ln_onnx(onnx_path) + + result = score( + onnx_path, + calibration_data=None, + num_synthetic_samples=8, + metric="kl_div", + target_precision="int8", + granularity="op_type", + calibration_eps=("cpu",), + op_types_scope=_SYNTHETIC_OP_SCOPE, + ) + assert result["calibration_source"] == "synthetic" + _assert_ln_over_conv(result["scores"]) + + +@pytest.mark.parametrize("metric", ["kl_div", "mse", "cos"]) +def test_synthetic_deterministic_ln_highest(tmp_path, metric): + """Tier 2: synthetic graph + deterministic real inputs -- LN scores highest of all ops.""" + onnx_path = str(tmp_path / "sens_synth.onnx") + _build_conv_mm_ln_onnx(onnx_path) + calib = _deterministic_calibration() + + result = score( + onnx_path, + calibration_data=calib, + metric=metric, + target_precision="int8", + granularity="op_type", + calibration_eps=("cpu",), + op_types_scope=_SYNTHETIC_OP_SCOPE, + ) + assert result["calibration_source"] == "real" + assert result["num_calibration_samples"] == 8 + scores = result["scores"] + assert scores, "No scores produced for synthetic graph." + # Highest-scoring op should be LayerNormalization. + top_op = max(scores.items(), key=lambda kv: kv[1])[0] + assert top_op == "LayerNormalization", ( + f"Expected LayerNormalization to be the top-ranked op, got '{top_op}' from {scores}" + ) + _assert_ln_over_conv(scores) + + +def _require_fixture(name: str) -> str: + """Return a fixture path or ``pytest.skip`` if it isn't staged on this host.""" + path = os.path.join(_FIXTURE_DIR, name) + if not os.path.exists(path): + pytest.skip(f"Sensitivity fixture missing: {path}") + return path + + +@pytest.fixture(scope="module") +def coatnet_fixtures() -> tuple[str, str]: + """CoAtNet-0 baseline ONNX + 500-sample ImageNet calibration for tier 3 / 4 tests.""" + return ( + _require_fixture("coatnet-0_rw_inpsize_1x3x224x224_opsetv_17_simplified.onnx"), + _require_fixture("imagenet_calib_500.npz"), + ) + + +@pytest.mark.slow +def test_coatnet_op_type_matches_manual_groundtruth(coatnet_fixtures): + """Tier 3: CoAtNet-0 op-type ranking must surface the ops that ``--op_types_to_quantize + Conv`` implicitly avoids. + + On CoAtNet-0 with 500-sample ImageNet calibration and ``kl_div``, the top-4 are + ``Add`` / ``Mul`` / ``LayerNormalization`` / ``ReduceMean`` (all > 1.5 KL) and + ``Conv`` sits ~10x below, matching the manual "Conv-only wins 82% top-1" ground truth. + Full ranking is documented in :doc:`_onnx_quantization`. + + Wall-clock ~14 min on H100. Fixtures (override root via ``MODELOPT_ONNX_ACCURACY_MODELS_DIR``): + ``coatnet-0_rw_inpsize_1x3x224x224_opsetv_17_simplified.onnx`` and ``imagenet_calib_500.npz``. + """ + onnx_path, calib_path = coatnet_fixtures + + result = score( + onnx_path, + calibration_data=calib_path, + metric="kl_div", + target_precision="int8", + granularity="op_type", + calibration_eps=("cuda:0", "cpu"), + ) + assert result["calibration_source"] == "real" + scores = result["scores"] + ranked = sorted(scores.items(), key=lambda kv: kv[1], reverse=True) + + top4 = {name for name, _ in ranked[:4]} + assert {"Add", "Mul", "LayerNormalization", "ReduceMean"}.issubset(top4), ( + f"Top-4 sensitive ops should include Add / Mul / LayerNormalization / " + f"ReduceMean (all > 1.5 KL), got {ranked}" + ) + # Conv sits ~10x below the top-4 -- justifies the Conv-only quantization policy. + assert scores["Conv"] < 0.5, ( + f"Conv score {scores['Conv']:.3f} unexpectedly high (top-4 are all > 1.5)" + ) + # These cluster at ~0 -- primitive won't recommend excluding them because there's + # nothing to exclude. + for op in ("Softmax", "Gemm", "GlobalAveragePool"): + assert scores.get(op, 0.0) < 0.001, f"{op} score {scores.get(op, 0.0):.3g} should be ~0" + + +@pytest.mark.slow_gpu +def test_coatnet_per_node_matches_manual_groundtruth(coatnet_fixtures): + """Tier 4: CoAtNet-0 per-node ranking (LN/MHA nodes top, Conv nodes bottom). + + Wall clock ~30-60 min; gated behind ``@pytest.mark.slow_gpu`` so default CI stays fast. + """ + onnx_path, calib_path = coatnet_fixtures + + result = score( + onnx_path, + calibration_data=calib_path, + metric="kl_div", + target_precision="int8", + granularity="node", + calibration_eps=("cuda:0", "cpu"), + ) + assert result["calibration_source"] == "real" + ranked = sorted(result["scores"].items(), key=lambda kv: kv[1], reverse=True) + assert len(ranked) >= 20, "Per-node ranking is unexpectedly short." + top_k = 10 + bottom_k = 10 + top_names = [name for name, _ in ranked[:top_k]] + bottom_names = [name for name, _ in ranked[-bottom_k:]] + # LayerNorm / MHA subgraph nodes dominate the top of the ranking. + assert any("layernorm" in n.lower() or "attn" in n.lower() for n in top_names), ( + f"Expected LN or MHA nodes in top-{top_k}, got {top_names}" + ) + # Individual Conv nodes cluster at the bottom (Conv-only ground truth). + assert any("conv" in n.lower() for n in bottom_names), ( + f"Expected Conv nodes in bottom-{bottom_k}, got {bottom_names}" + ) diff --git a/tests/unit/onnx/quantization/test_nodes_to_quantize.py b/tests/unit/onnx/quantization/test_nodes_to_quantize.py new file mode 100644 index 00000000000..3cb19f75208 --- /dev/null +++ b/tests/unit/onnx/quantization/test_nodes_to_quantize.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the ``nodes_to_quantize`` allow-list filter (symmetric with ``nodes_to_exclude``). + +Builds a small two-Conv ONNX graph and asserts that ``nodes_to_quantize=["conv_keep"]`` produces +Q/DQ around ``conv_keep`` only, leaving ``conv_skip`` in its original precision. This is the +primitive the ONNX sensitivity scanner relies on to isolate a single node for a per-target probe. +""" + +from __future__ import annotations + +import os + +import numpy as np +import onnx +import onnx_graphsurgeon as gs +from onnx import TensorProto, helper, numpy_helper + +import modelopt.onnx.quantization as moq + + +def _build_two_conv_onnx(path: str, opset: int = 17) -> None: + """Emit a 2-Conv ONNX with the node names the test filters on.""" + rng = np.random.default_rng(0) + w1 = rng.standard_normal((16, 16, 3, 3)).astype(np.float32) * 0.1 + b1 = np.zeros((16,), dtype=np.float32) + w2 = rng.standard_normal((16, 16, 3, 3)).astype(np.float32) * 0.1 + b2 = np.zeros((16,), dtype=np.float32) + + nodes = [ + helper.make_node( + "Conv", + ["input", "w1", "b1"], + ["conv_keep_out"], + name="conv_keep", + pads=[1, 1, 1, 1], + strides=[1, 1], + ), + helper.make_node( + "Conv", + ["conv_keep_out", "w2", "b2"], + ["output"], + name="conv_skip", + pads=[1, 1, 1, 1], + strides=[1, 1], + ), + ] + initializers = [ + numpy_helper.from_array(w1, "w1"), + numpy_helper.from_array(b1, "b1"), + numpy_helper.from_array(w2, "w2"), + numpy_helper.from_array(b2, "b2"), + ] + graph = helper.make_graph( + nodes=nodes, + name="nodes_to_quantize_test", + inputs=[helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 16, 8, 8])], + outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 16, 8, 8])], + initializer=initializers, + ) + onnx.save( + helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)], ir_version=8), + path, + ) + + +def _has_dq_predecessor(node: gs.Node, input_idx: int) -> bool: + """Return True when the input at ``input_idx`` of ``node`` is produced by DequantizeLinear. + + Returns False (rather than raising) when the input has no producer (graph input) or when + the producer chain is shorter than expected. + """ + inp = node.inputs[input_idx] + if not isinstance(inp, gs.Variable) or not inp.inputs: + return False + producer = inp.inputs[0] + if producer.op == "Cast": + cast_inp = producer.inputs[0] if producer.inputs else None + if not isinstance(cast_inp, gs.Variable) or not cast_inp.inputs: + return False + producer = cast_inp.inputs[0] + return producer.op == "DequantizeLinear" + + +def test_nodes_to_quantize_restricts_qdq_to_single_conv(tmp_path): + """`nodes_to_quantize=["conv_keep"]` inserts Q/DQ around conv_keep only.""" + onnx_path = str(tmp_path / "two_conv.onnx") + _build_two_conv_onnx(onnx_path) + calibration_data = { + "input": np.random.default_rng(0).standard_normal((2, 16, 8, 8)).astype(np.float32) + } + + moq.quantize( + onnx_path, + quantize_mode="int8", + calibration_data=calibration_data, + calibration_eps=["cpu"], + nodes_to_quantize=["^conv_keep$"], + high_precision_dtype="fp32", + ) + + quantized_path = onnx_path.replace(".onnx", ".quant.onnx") + assert os.path.isfile(quantized_path) + + graph = gs.import_onnx(onnx.load(quantized_path)) + keep_nodes = [n for n in graph.nodes if n.name == "conv_keep"] + skip_nodes = [n for n in graph.nodes if n.name == "conv_skip"] + assert len(keep_nodes) == 1, ( + f"conv_keep not found in quantized graph: {[n.name for n in graph.nodes]}" + ) + assert len(skip_nodes) == 1, ( + f"conv_skip not found in quantized graph: {[n.name for n in graph.nodes]}" + ) + + # conv_keep must have DQ on its activation input; conv_skip must not. + assert _has_dq_predecessor(keep_nodes[0], 0), ( + "conv_keep is not quantized despite nodes_to_quantize=['conv_keep']" + ) + assert not _has_dq_predecessor(skip_nodes[0], 0), ( + "conv_skip was quantized but nodes_to_quantize only listed conv_keep" + ) diff --git a/tests/unit/onnx/quantization/test_sensitivity_picker.py b/tests/unit/onnx/quantization/test_sensitivity_picker.py new file mode 100644 index 00000000000..085b46683c0 --- /dev/null +++ b/tests/unit/onnx/quantization/test_sensitivity_picker.py @@ -0,0 +1,218 @@ +# 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): + # Total = 10. coverage=0.5 -> target 5. top-1 is 4 (fits), top-2 would + # be 7 (crosses 5) -> stop at 1 node. + scores = {"a": 4.0, "b": 3.0, "c": 2.0, "d": 1.0} + assert suggest_exclusion(scores, coverage=0.5) == ["a"] + + def test_includes_second_when_it_fits(self): + # Total = 10. coverage=0.8 -> target 8. top-1 (4) + top-2 (7) both fit, + # top-3 would be 9 (crosses 8) -> stop at 2 nodes. + scores = {"a": 4.0, "b": 3.0, "c": 2.0, "d": 1.0} + assert suggest_exclusion(scores, coverage=0.8) == ["a", "b"] + + def test_full_coverage_returns_all_nodes(self): + # coverage=1.0 -> target = total, everything fits exactly. + scores = {"a": 1.0, "b": 2.0, "c": 3.0} + result = suggest_exclusion(scores, coverage=1.0) + assert set(result) == {"a", "b", "c"} + + def test_returns_sorted_by_kl_desc(self): + scores = {"low": 0.1, "high": 0.9, "mid": 0.5} + # coverage=1.0 -> everything fits, and result is sorted by KL desc. + assert suggest_exclusion(scores, coverage=1.0) == ["high", "mid", "low"] + + def test_top_node_alone_exceeds_target(self): + # Total = 10, coverage=0.2 -> target 2. Top node (5) alone exceeds + # target, so nothing is included. + scores = {"a": 5.0, "b": 3.0, "c": 2.0} + assert suggest_exclusion(scores, coverage=0.2) == [] + + def test_zero_target_returns_empty(self): + scores = {"a": 5.0, "b": 3.0} + assert suggest_exclusion(scores, coverage=0.0) == [] + + def test_zero_total_returns_empty(self): + assert suggest_exclusion({"a": 0.0, "b": 0.0}, coverage=0.9) == [] + + def test_empty_scores_returns_empty(self): + assert suggest_exclusion({}, coverage=0.9) == [] + + def test_max_nodes_caps_exclusion_set(self): + # 10 nodes at KL 10..1. Total = 55. coverage=1.0 would include all, + # but max_nodes=3 caps at 3. + scores = {chr(ord("a") + i): 10.0 - i for i in range(10)} + assert suggest_exclusion(scores, coverage=1.0, max_nodes=3) == ["a", "b", "c"] + + def test_min_score_floor_stops_before_low_nodes(self): + # Even at coverage=1.0, nodes below the floor are excluded from the set. + scores = {"hi_1": 5.0, "hi_2": 4.0, "trivial_1": 0.001, "trivial_2": 0.0001} + result = suggest_exclusion(scores, coverage=1.0, min_score_floor=0.01) + assert result == ["hi_1", "hi_2"] + + def test_vit_like_distribution_undershoots_cleanly(self): + # Mimics ViT-tiny's distribution: 15 nodes at KL ~3.8-6.7, then a + # sharp drop to ~3.05 for ranks 16-17, then a long tail. + big = {f"top_{i}": 6.7 - i * 0.2 for i in range(15)} # ranks 1-15, KL ~6.7 down to ~3.9 + borderline = {"rank_16": 3.057, "rank_17": 3.055} + tail = {f"tail_{i}": 0.5 - i * 0.02 for i in range(30)} + scores = {**big, **borderline, **tail} + total = sum(scores.values()) + result = suggest_exclusion(scores, coverage=0.90) + excluded_mass = sum(scores[n] for n in result) + # Actual coverage never exceeds requested. + assert excluded_mass <= 0.90 * total + # But should still capture most of the mass with fewer than the total. + assert len(result) < len(scores) + + +class TestThresholdMode: + """Tests the absolute-KL cutoff semantic: exclude all nodes above threshold.""" + + def test_picks_all_above_absolute_threshold(self): + scores = {"a": 5.0, "b": 3.0, "c": 1.0, "d": 0.5, "e": 0.05} + assert suggest_exclusion(scores, threshold=1.0) == ["a", "b"] + + def test_returns_sorted_by_kl_desc(self): + scores = {"low_hit": 0.6, "high_hit": 0.9, "mid_hit": 0.75, "miss": 0.1} + assert suggest_exclusion(scores, threshold=0.5) == ["high_hit", "mid_hit", "low_hit"] + + def test_boundary_score_is_excluded_from_set(self): + # A score exactly at the threshold does NOT get excluded (strict >). + scores = {"above": 0.11, "at": 0.10, "below": 0.09} + assert suggest_exclusion(scores, threshold=0.10) == ["above"] + + def test_no_nodes_above_threshold_returns_empty(self): + assert suggest_exclusion({"a": 0.01, "b": 0.005}, threshold=1.0) == [] + + def test_threshold_overrides_coverage(self): + scores = {"a": 5.0, "b": 3.0, "c": 2.98, "d": 2.0} + # coverage=0.99 would try to include most; threshold overrides. + assert suggest_exclusion(scores, coverage=0.99, threshold=2.5) == ["a", "b", "c"] + + def test_max_nodes_still_caps_threshold_mode(self): + # All 10 nodes have score > 5.0 but max_nodes=3 caps at 3. + scores = {chr(ord("a") + i): 10.0 - i * 0.1 for i in range(10)} + assert suggest_exclusion(scores, threshold=5.0, max_nodes=3) == ["a", "b", "c"] + + def test_min_score_floor_composes_with_threshold(self): + # threshold=0.1 would normally include all three, but min_score_floor=1.0 + # short-circuits after "a" (b=0.5 is below the floor). + scores = {"a": 5.0, "b": 0.5, "c": 0.3} + assert suggest_exclusion(scores, threshold=0.1, min_score_floor=1.0) == ["a"] + + +class TestNearTieWarning: + """Warning fires when the cut-off between included and excluded is a near-tie.""" + + def test_warning_fires_on_near_tied_cutoff(self, caplog): + # Ranks 16 and 17 are near-tied at KL 3.06 vs 3.05 (99.7% ratio); coverage=0.94 + # cuts between them. + scores = { + f"node_{i:02d}": kl + for i, kl in enumerate( + [ + 6.7, + 5.7, + 4.6, + 4.3, + 4.1, + 4.1, + 4.0, + 4.0, + 3.8, + 3.8, + 3.8, + 3.8, + 3.7, + 3.7, + 3.7, + 3.06, + 3.05, + 0.8, + 0.5, + 0.1, + ], + 1, + ) + } + with caplog.at_level(logging.WARNING, logger="modelopt.onnx"): + suggest_exclusion(scores, coverage=0.94) + messages = [r.message for r in caplog.records] + assert any("near-tie at the exclusion cut-off" in m for m in messages) + + def test_no_warning_when_cut_is_not_a_near_tie(self, caplog): + # ViT-like distribution where coverage=0.75 cuts between very different KL values. + scores = { + f"node_{i:02d}": kl for i, kl in enumerate([6.7, 5.7, 4.6, 4.3, 4.1, 0.5, 0.2, 0.1], 1) + } + with caplog.at_level(logging.WARNING, logger="modelopt.onnx"): + suggest_exclusion(scores, coverage=0.75) + messages = [r.message for r in caplog.records] + assert not any("near-tie" in m for m in messages) + + def test_warning_disabled_by_none(self, caplog): + # Setting near_tie_ratio=None disables the warning entirely. + scores = {"a": 5.0, "b": 4.99, "c": 0.1} + with caplog.at_level(logging.WARNING, logger="modelopt.onnx"): + suggest_exclusion(scores, coverage=0.5, near_tie_ratio=None) + messages = [r.message for r in caplog.records] + assert not any("near-tie" in m for m in messages) + + def test_threshold_mode_also_warns_on_near_tie(self, caplog): + # threshold=3.056 cuts between KL 3.06 (above threshold) and 3.05 (below) -- near-tie. + scores = {"a": 6.7, "b": 5.7, "c": 3.06, "d": 3.05, "e": 0.1} + with caplog.at_level(logging.WARNING, logger="modelopt.onnx"): + suggest_exclusion(scores, threshold=3.056) + messages = [r.message for r in caplog.records] + assert any("near-tie" in m and "mode=threshold" in m for m in messages) + + +class TestSummarizeExclusion: + def test_reports_coverage_pct_and_counts(self): + scores = {"a": 4.0, "b": 3.0, "c": 2.0, "d": 1.0} + summary = summarize_exclusion(scores, ["a", "b"]) + assert summary["num_excluded"] == 2 + assert summary["num_previously_quantized"] == 4 + assert summary["num_remaining_quantized"] == 2 + assert summary["coverage_pct"] == pytest.approx(70.0) + assert summary["excluded_mass"] == pytest.approx(7.0) + assert summary["total_mass"] == pytest.approx(10.0) + + def test_empty_scores_zero_coverage(self): + summary = summarize_exclusion({}, []) + assert summary["coverage_pct"] == 0.0 + assert summary["num_excluded"] == 0 + + def test_missing_node_names_default_zero(self): + scores = {"a": 5.0, "b": 5.0} + summary = summarize_exclusion(scores, ["a", "unknown"]) + assert summary["excluded_mass"] == pytest.approx(5.0) + assert summary["coverage_pct"] == pytest.approx(50.0) + assert summary["num_excluded"] == 2