[ONNX] Add quantization sensitivity ranking + exclusion picker

[gcunhase]

What does this PR do?

Type of change: new feature (ONNX PTQ tooling)

Adds modelopt.onnx.quantization.sensitivity — a first-class ONNX PTQ primitive that ranks each quantizable target (op type or individual node) by its impact on model output, plus a downstream picker that turns the ranking into an actionable --nodes_to_exclude / --op_types_to_exclude list. Closes the manual-investigation gap for accuracy degradation in ONNX models with QDQ nodes.

• sensitivity.score — for every target, inserts calibrated Q/DQ on that target via the standard modelopt.onnx.quantization.quantize pipeline, runs both the reference and the quantized graphs through ORT, and computes a proxy metric (kl_div default, mse, or cos) between their outputs. Higher = more accuracy loss if quantized. granularity={op_type, node}, target_precision={int8, fp8}, real or synthetic calibration.
• sensitivity.suggest_exclusion — coverage mode (cumulative-mass, architecture-portable) or threshold mode (absolute cutoff), plus an optional blocks= / block_agg= argument for block-level aggregation where per-node picking would fragment precision within a block.
• sensitivity.summarize_exclusion — one-call summary of what an exclusion set covers.
• CLI: python -m modelopt.onnx.quantization.sensitivity renders a ranked table to stderr and writes a JSON side-file.

Small supporting extension to modelopt/onnx/quantization/quantize.py — symmetric nodes_to_quantize argument alongside the existing nodes_to_exclude, enabling per-node granularity for sensitivity scoring and any future single-node workflow. Also adds get_op_types_in_graph to modelopt/onnx/utils.py and TopK to is_fusible_reduction_op in modelopt/onnx/op_types.py.

Usage

Python code:

from modelopt.onnx.quantization import quantize
from modelopt.onnx.quantization.sensitivity import score, suggest_exclusion, summarize_exclusion

# 1. Rank the quantizable targets in the graph.
result = score(
    onnx_path="coatnet-0.onnx",
    calibration_data="imagenet_calib_500.npz",
    granularity="op_type",     # or "node"
    metric="kl_div",           # or "mse", "cos"
    target_precision="int8",
)
# result["scores"] is a dict {op_type_or_node_name: metric_value}, higher = more sensitive.

# 2. Turn the ranking into an exclusion list.
excluded = suggest_exclusion(result["scores"], coverage=0.90)
print(summarize_exclusion(result["scores"], excluded))

# 3. Quantize with the exclusion applied.
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",
)

CLI equivalent:

python -m modelopt.onnx.quantization.sensitivity \
    --onnx_path coatnet-0.onnx \
    --calibration_data_path imagenet_calib_500.npz \
    --granularity op_type \
    --metric kl_div

Block-level aggregation on transformer architectures (ViT-tiny example, picks whole transformer blocks instead of individual nodes):

blocks = {f"blocks.{n}": [rf"^/blocks/blocks\.{n}/"] for n in range(12)}
excluded = suggest_exclusion(
    result["scores"], threshold=0.1, blocks=blocks, block_agg="max",
)

Testing

Four-tier test layout (ordered from lightest to heaviest, tests/gpu/onnx/quantization/test_sensitivity.py):

• Tier 1 (unit, seconds) — synthetic-random-calibration regression guard: LN > Conv directionally holds even under calibration_data=None.
• Tier 2 (unit, seconds, parametrized over kl_div / mse / cos) — synthetic 2-Conv + 1-MatMul + 1-LayerNorm graph with deterministic real inputs: LayerNormalization scores highest of all ops.
• Tier 3 (@pytest.mark.slow, ~14 min) — CoAtNet-0 op-type ranking on 500-sample ImageNet calibration: top-4 = Add / Mul / LayerNormalization / ReduceMean (all > 1.5 KL), Conv sits ~10× below. Matches the manual "Conv-only wins 82% top-1" ground truth read as a quantization policy.
• Tier 4 (@pytest.mark.slow_gpu, ~30–60 min) — CoAtNet-0 per-node ranking: LN / MHA nodes in top-10, individual Conv nodes in bottom-10.

Slow fixtures (Tiers 3–4) resolve via MODELOPT_ONNX_ACCURACY_MODELS_DIR; missing fixtures pytest.skip cleanly.

Picker-side unit tests (tests/unit/onnx/quantization/test_sensitivity_picker.py) cover coverage / threshold / near-tie warning / max_nodes / min_score_floor / blocks semantics. tests/unit/onnx/quantization/test_nodes_to_quantize.py covers the new nodes_to_quantize filter on quantize.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: ✅ Pure additions: new modelopt.onnx.quantization.sensitivity package, a new nodes_to_quantize argument on quantize (default None preserves existing behavior), and two small helpers in modelopt/onnx/op_types.py + modelopt/onnx/utils.py. No existing signatures changed.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: ✅ No new PIP dependencies; primitive uses existing numpy / onnx / onnxruntime / modelopt.onnx.quantization.quantize machinery only.
• Did you write any new necessary tests?: ✅ Four-tier layout above; unit tests for picker + nodes_to_quantize.
• Did you update Changelog?: ✅ New Features → Quantization: one-line entry covering the primitive, picker, and optional block-level aggregation.
• Did you get Claude approval on this PR?: N/A — will run /claude review after opening.

Additional Information

Motivation — post-training quantization on ONNX hybrid architectures currently has no automated way to identify sensitivity-driving ops. For CoAtNet-0, closing the accuracy gap required manually cycling through six exclusion policies (default QDQ 22%, attention-MatMul excluded 21%, --op_types_to_exclude MatMul 36%, --disable_mha_qdq 40%, --disable_mha_qdq --op_types_to_exclude Softmax 33%, Conv-only 82%) before finding the winner. That per-model investigation is O(model × strategy) hours of operator time and produces no reusable artifact. This primitive replaces the manual sweep with a single ranking call.

Relationship to existing ModelOpt code — the existing modelopt/onnx/quantization/autotune/ package is orthogonal: it optimizes TensorRT latency with no accuracy signal, no per-op sensitivity score, and no output-drift metric. modelopt.torch.quantization.model_quant.auto_quantize is the design analog we mirror on the ONNX side (same "higher = more sensitive = keep at higher precision" contract; ONNX port differs in autograd-free, ORT-based, graph-mutation forward passes).

Validation results — 500-image ImageNet-1k validation across four models on NVIDIA H100 (GH100) and TensorRT 10.16.2.11. Metrics: GPU Compute Time median (ms) for latency and Top-1 for accuracy. INT8 baseline is trtexec --int8 --fp16 auto-selection (implicit quantization); QDQ-INT8 is ModelOpt's default behavior; and QDQ-INT8 + sensitivity is the best sensitivity-driven exclusion recipe per model.

Model	INT8 baseline (--int8 --fp16)	QDQ-INT8 (default)	QDQ-INT8 + sensitivity exclusions	Best recipe
CoAtNet-0	80.6% / 1.047 ms	22.4% / 1.220 ms	81.4% / 1.218 ms	per-node coverage=0.90 (26/345 excluded)
MobileNetV3-L	68.9% / 0.331 ms	48.3% / 0.501 ms	69.0% / 0.606 ms	per-node KL > 0.002 (19/139 excluded)
ResNet-50	78.9% / 0.278 ms	57.9% / 0.276 ms	79.6% / 0.294 ms	per-node KL > 0.002 (9/121 excluded)
ViT-tiny	76.0% / 0.525 ms	6.7% / 0.357 ms	75.2% / 0.408 ms	block-level blocks.7-11 + /norm/LN (101/244 excluded)

Synthetic-calibration caveat — calibration_data=None is a supported fallback but produces directional rankings only. The CLI prints a warning, the output JSON's calibration_source field reads "synthetic", and downstream absolute-threshold consumers should reject synthetic-calibrated scores. Attention-heavy models are the highest-risk case — recommend real calibration when the model has attention.

(Maybe) Follow-up (separate PR) — constrained solver on top of this primitive that takes accuracy_budget_pp and returns an ExclusionConfig, mirroring the torch auto_quantize API signature.

🤖 Generated with Claude Code

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

_{Comment @coderabbitai help to get the list of available commands.}

[cjluo-nv]

Bot review (claude-opus-5) — DM the bot to share feedback.

Summary

New ONNX PTQ subsystem (modelopt.onnx.quantization.sensitivity: score / picker / metrics / CLI) plus a 317-line docs section, ~2000 added lines. The core idea is sound and the picker unit tests are good, but there are a few correctness/coverage issues I'd like resolved before this lands.

Design gate (architectural-change protocol)

Problem: ONNX PTQ has no automated way to find which ops/nodes drive accuracy loss, so operators hand-sweep exclusion policies per model.

Alternatives checked in-repo:

• modelopt/onnx/quantization/autotune/ — the PR body explicitly addresses this (latency-driven, no accuracy signal). Reasonable.
• modelopt.torch.quantization.model_quant.auto_quantize — body names it as the design analog. Reasonable.

So the top-level "why a new subsystem" question is addressed; I'm not blocking on it. Two smaller reuse questions the body does not address:

1. score._resolve_calibration_data / _load_calibration_from_path and __main__._load_calibration re-implement calibration ingestion that already exists in modelopt/onnx/quantization/calib_utils.py (CalibrationDataProvider normalizes ndarray/dict/batch-splitting; RandomDataProvider+gen_random_inputs covers the synthetic fallback) and in modelopt/onnx/quantization/__main__.py (npz/npy loading with the --trust_calibration_data/validate_file_size security path, which the new CLI does not inherit). This is now a third loader with slightly different semantics.
2. A second python -m ... entrypoint vs. a flag on the existing ONNX PTQ CLI — worth a sentence in the body.

Blocking-ish findings

1. op_type and node granularity are not equivalent probes (see inline on score.py): node targets pass only nodes_to_quantize and leave op_types_to_quantize=None, so ORT's effective allow-list is the post-configure_ort registry, which has Relu/Sigmoid/Softmax/Concat/Transpose/... deleted. Those nodes silently get no Q/DQ → score 0.0 ("safe to quantize"), while the same op scores non-zero in op-type mode. Since the headline results are per-node, this matters.
2. op_types.py change is not a pure addition. Adding "Gelu" to get_activation_ops() also changes QDQAutotunerBase.get_ort_quantization_config() (op_types_needing_output_quant), i.e. autotune Q/DQ placement for any Gelu model. No test, and the PR body describes a different op_types change (TopK → is_fusible_reduction_op, which is already in main).
3. blocks / block_agg has zero tests despite being the documented best recipe for ViT-tiny and the body claiming coverage for it. metrics.py and the whole CLI module (_render_ranked_table, _load_calibration, _default_output_json) are also untested.
4. PR description is out of sync with the diff: it claims a nodes_to_quantize extension to quantize.py, but quantize.py is not in the diff — that argument already exists on main, so tests/unit/onnx/quantization/test_nodes_to_quantize.py is a test of pre-existing behavior (fine to add, but please re-word the body so reviewers know what actually changed).
5. Size: 2003 lines. The picker + its tests, the scorer + CLI, and the docs section are three fairly independent units and would review much better split.

Answer to your docs question

I'd put it in a new dedicated guide page, mirroring docs/source/guides/9_autotune.rst, with only a short pointer paragraph at the end of _onnx_quantization.rst. Reasons: (a) 317 lines is already longer than several complete sections of the ONNX PTQ guide and pushes the "how do I quantize an ONNX model" narrative off the page; (b) the content is a workflow with its own CLI, its own concepts (coverage vs threshold, block aggregation, near-tie), and its own troubleshooting — same shape as autotune; (c) a separate page lets you use .. argparse:: for the new CLI the way 9_autotune.rst does. The parts that are reproduction recipes rather than API docs — the timm export snippet, the ImageNet-500 NPZ preparation, and the per-model validation table — belong in examples/onnx_ptq/README.md, since they need real datasets/checkpoints and will drift with model versions.

No prompt-injection content observed in the PR text.

[cjluo-nv]

Bot comment.

Node-granularity probes pass only nodes_to_quantize and leave op_types_to_quantize unset, while op-type probes pass op_types_to_quantize=[op]. Those two paths are not equivalent: quantize() → configure_ort() deletes Relu, Sigmoid, Softmax, Concat, Gather, Transpose, Reshape, Pad, ... from QLinearOpsRegistry/QDQRegistry for every op type not in op_types_to_quantize, and ORT's should_quantize_node gates on the resulting allow-list. So a per-node probe on e.g. a Softmax, Relu or Sigmoid node inserts no Q/DQ at all and is recorded as 0.0 — indistinguishable from "quantizing this node is harmless", which is exactly the signal suggest_exclusion consumes.

Since the per-node ranking is what the headline CoAtNet/MobileNet/ResNet results are built on, this looks like a real correctness gap. Suggest also passing op_types_to_quantize=[node.op_type] alongside the node regex so both granularities probe the same registry, and adding a unit test that a per-node probe on an activation op (e.g. the Sigmoid/Relu in a small synthetic graph) produces a non-zero score.

[cjluo-nv]

Bot comment.

Two related robustness points around the probe loop:

1. except Exception: ... continue drops the target from scores entirely, so a target that fails to quantize is silently absent from both the ranking and summarize_exclusion's num_previously_quantized. Consider recording it explicitly (e.g. a failed list in the returned dict) so downstream consumers can tell "failed" from "not probed".
2. There's no verification that the probe actually inserted Q/DQ. As the docstring notes, ops that slip past the filter but can't be quantized "are reported with score 0.0" — but 0.0 is semantically "quantizing this is free", which is the opposite of "unknown". Cheap fix: after quantize(), count QuantizeLinear nodes in probe_path and, if zero, log a warning and record the target as unprobed rather than 0.0.

[cjluo-nv]

Bot comment.

This is not a pure addition. get_activation_ops() is consumed by QDQAutotunerBase.get_ort_quantization_config() (modelopt/onnx/quantization/autotune/autotuner_base.py) to build op_types_needing_output_quant, so adding Gelu changes autotune's Q/DQ placement for any model containing Gelu. That's a behavior change to a shipped feature with no test and no mention in the PR body (the body instead describes a TopK addition to is_fusible_reduction_op, which is already on main).

Please either (a) call this out in the PR body/CHANGELOG and add an autotune-side test asserting the new op_types_needing_output_quant behavior, or (b) keep the sensitivity scope local — e.g. add Gelu to _default_op_types_scope in score.py instead of mutating the shared helper.

[cjluo-nv]

Bot comment.

The PR body says these tests cover "blocks semantics", but there is no test here exercising blocks= / block_agg=. That path is the documented best recipe for ViT-tiny (101/244 nodes excluded) and it has non-trivial semantics worth pinning: first-match-wins across blocks iteration order, unmatched nodes becoming singleton groups, sum/max/mean aggregation, the union-of-members return value (and its ordering), and the ValueError on an invalid block_agg. The docs also assert that threshold=0.1, block_agg="max" and coverage=1.0, max_nodes=6, block_agg="sum" yield the same 101-node set — that equivalence claim is a good test case.

Also missing: direct unit tests for metrics.py (identical inputs → kl_div/mse/cos == 0; orthogonal vectors → cos_dist == 1; 0-D tensor handling in _flatten_per_sample), which is pure numpy and costs milliseconds.

[cjluo-nv]

Bot comment.

_load_calibration partially duplicates score._load_calibration_from_path (which already handles .npz, .npy and directories) — the CLI eagerly loads .npz but delegates the other two forms, so there are now two code paths with different behavior for the same flag. Suggest just passing the path straight through to score().

Separately, the existing ONNX PTQ CLI (modelopt/onnx/quantization/__main__.py) guards calibration/model loading with validate_file_size() and the --trust_calibration_data opt-in for pickled data; this new entrypoint inherits neither. allow_pickle=False covers the RCE case, but the size guard and the consistent flag surface are worth mirroring.

[cjluo-nv]

Bot comment.

Two things about placement/comments here:

1. Tiers 1 and 2 run with calibration_eps=("cpu",) on a tiny synthetic graph — they don't need a GPU, so putting them under tests/gpu/ means the only fast regression guards for score() won't run in the unit CI job. Suggest moving tiers 1–2 to tests/unit/onnx/quantization/ and keeping only the slow/slow_gpu CoAtNet tiers here.
2. This comment says the score() default is get_autotuner_quantizable_ops() and that it excludes LayerNormalization. That's stale: score() uses _default_op_types_scope(), which includes normalization ops via is_normalization_op. Either drop _SYNTHETIC_OP_SCOPE and let the default scope apply (which would also make the test exercise the default path), or fix the rationale.

[codecov]

Codecov Report

❌ Patch coverage is 34.37500% with 210 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.78%. Comparing base (73d7784) to head (956bd85).

Files with missing lines	Patch %	Lines
modelopt/onnx/quantization/sensitivity/score.py	29.41%	96 Missing ⚠️
modelopt/onnx/quantization/sensitivity/__main__.py	0.00%	72 Missing ⚠️
modelopt/onnx/quantization/sensitivity/picker.py	72.50%	22 Missing ⚠️
modelopt/onnx/quantization/sensitivity/metrics.py	29.62%	19 Missing ⚠️
modelopt/onnx/utils.py	50.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2240      +/-   ##
==========================================
- Coverage   79.01%   78.78%   -0.24%     
==========================================
  Files         523      528       +5     
  Lines       60695    61015     +320     
==========================================
+ Hits        47960    48072     +112     
- Misses      12735    12943     +208     

Flag	Coverage Δ
unit	55.56% <34.37%> (-0.11%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.