# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries.
# SPDX-License-Identifier: BSD-3-Clause
"""Per-quantizer quantization sensitivity analysis for ONNX models.
:func:`analyze_per_quantizer_sensitivity` measures how sensitive a calibrated
:class:`QuantizationSimModel` is to quantization at a finer granularity than the
op-level :func:`aimet_onnx.analyze_per_layer_sensitivity`, by enabling one
quantizer (or group of quantizers) at a time. An optional ``group_fn`` selects
and groups the quantizers to sweep -- e.g. pass a ``group_fn`` that returns
``None`` for everything but the KV-cache tensors to run a KV-cache-only sweep.
It accepts a :class:`SensitivityMetric` (a named ``eval_fn`` plus ranking
semantics) and returns a ``{name: score}`` dict ranked most-sensitive-first.
Feed that dict into :func:`aimet_onnx.lite_mp.flip_layers_to_higher_precision`
(keyed by op name) to raise the most sensitive units to a higher precision.
The metric ``eval_fn`` takes an ``onnxruntime.InferenceSession`` and returns a
float, matching the existing :func:`aimet_onnx.analyze_per_layer_sensitivity`
and :func:`aimet_onnx.utils.make_psnr_eval_fn` convention -- there is no
dependency on any particular model harness.
"""
from typing import Callable, Dict, Iterable, List, Optional
import numpy as np
import onnxruntime as ort
from tqdm import tqdm
from aimet_onnx.quantsim import QuantizationSimModel
from aimet_onnx.utils import disable_quantizers
from aimet_onnx.common.utils import AimetLogger, compute_psnr
logger = AimetLogger.get_area_logger(AimetLogger.LogAreas.Quant)
[docs]
class SensitivityMetric:
"""A named evaluation function with ranking semantics.
:param name: Human-readable metric name (used in logs and plots).
:param eval_fn: Callable taking an ``onnxruntime.InferenceSession`` and
returning a scalar score.
:param higher_is_worse: If ``True`` (e.g. perplexity, KL divergence), a
larger score means greater sensitivity. If ``False`` (e.g. PSNR,
accuracy), a smaller score means greater sensitivity. Controls the
ordering used when ranking units most-sensitive-first.
"""
def __init__(
self,
name: str,
eval_fn: Callable[[ort.InferenceSession], float],
higher_is_worse: bool = True,
):
if not callable(eval_fn):
raise ValueError(f"eval_fn is expected to be callable; got {type(eval_fn)}")
self.name = name
self.eval_fn = eval_fn
self.higher_is_worse = higher_is_worse
def __call__(self, session: ort.InferenceSession) -> float:
return self.eval_fn(session)
[docs]
def sensitivity_score(self, score: float) -> float:
"""Map a raw score to a value where larger always means more sensitive."""
return score if self.higher_is_worse else -score
[docs]
def make_topk_logit_psnr_metric(
fp_session: ort.InferenceSession,
inputs: Iterable[Dict[str, np.ndarray]],
k: int = 10,
logit_output_index: int = 0,
) -> SensitivityMetric:
"""Build a top-k logit PSNR metric for LLM-style models.
For each input sample, the reference (floating-point) logits are compared
against the sim's logits, restricted to the top-``k`` vocabulary indices of
the FP logits at each position. PSNR is computed over those slices,
aggregated across all samples. Higher PSNR is better, so
``higher_is_worse=False`` -- a drop in PSNR indicates greater sensitivity.
The FP logits are captured once, up front, from ``fp_session``.
:param fp_session: ORT inference session for the floating-point model.
:param inputs: Iterable of input feed dicts (``{input_name: np.ndarray}``).
:param k: Number of top FP vocabulary indices to compare per position.
:param logit_output_index: Index of the logits tensor in the model outputs.
:return: A :class:`SensitivityMetric` wrapping the top-k logit PSNR eval fn.
"""
inputs = list(inputs)
def _topk_slices(logits: np.ndarray, topk_idx: np.ndarray) -> np.ndarray:
# logits: [..., vocab]; topk_idx: [..., k] along the last axis.
return np.take_along_axis(logits, topk_idx, axis=-1)
# The reference logits, their top-k indices, and the gathered FP slice depend
# only on the (fixed) FP model -- not on the sim under test. Precompute them
# once here so the per-eval hot path (called once per quantizer in a sweep)
# only runs the quantized session, gathers with the cached indices, and
# computes PSNR. This hoists the expensive top-k off the vocab axis
# (~150k wide) out of the inner loop entirely.
#
# argpartition (O(vocab)) rather than argsort (O(vocab log vocab)): we only
# need *which* k indices are largest, not their order, and PSNR over the
# gathered slice is order-independent.
topk_indices: List[np.ndarray] = []
fp_slices: List[np.ndarray] = []
for feed in inputs:
fp_arr = np.asarray(
fp_session.run(None, feed)[logit_output_index], dtype=np.float32
)
effective_k = min(k, fp_arr.shape[-1])
idx = np.argpartition(fp_arr, -effective_k, axis=-1)[..., -effective_k:]
topk_indices.append(idx)
fp_slices.append(_topk_slices(fp_arr, idx).reshape(-1))
expected = np.concatenate(fp_slices)
def _eval(session: ort.InferenceSession) -> float:
qt_slices: List[np.ndarray] = []
for feed, idx in zip(inputs, topk_indices):
qt_arr = np.asarray(
session.run(None, feed)[logit_output_index], dtype=np.float32
)
qt_slices.append(_topk_slices(qt_arr, idx).reshape(-1))
actual = np.concatenate(qt_slices)
return compute_psnr(expected, actual)
return SensitivityMetric(
name=f"Top{k}LogitPSNR", eval_fn=_eval, higher_is_worse=False
)
[docs]
def analyze_per_quantizer_sensitivity(
sim: QuantizationSimModel,
metric: SensitivityMetric,
group_fn: Optional[Callable[[str], Optional[str]]] = None,
) -> Dict[str, float]:
"""Analyze sensitivity by enabling one quantizer (or quantizer group) at a time.
All quantizers are first disabled. Then, for each group, that group's
quantizers are enabled, the metric is evaluated on ``sim.session``, and the
group is disabled again. This isolates each unit's contribution to
quantization error at a finer granularity than the op-level
:func:`aimet_onnx.analyze_per_layer_sensitivity`.
:param sim: Calibrated QuantizationSimModel to analyze. Its quantizers'
enabled state is restored on return.
:param metric: :class:`SensitivityMetric` used to score each group.
:param group_fn: Optional callable mapping a quantizer name to a group key.
Quantizers sharing a key are enabled together; a key of ``None`` skips
that quantizer. If omitted, each currently-enabled quantizer is its own
group (keyed by its own name). To restrict the sweep to a subset (e.g.
KV-cache tensors), return ``None`` for the quantizers to exclude::
analyze_per_quantizer_sensitivity(
sim, metric,
group_fn=lambda name: name if name in kv_names else None,
)
:return: Dict mapping group key to its metric score, ordered
most-sensitive-first per ``metric``.
"""
if group_fn is None:
group_fn = lambda name: name # noqa: E731 - one quantizer per group
groups: Dict[str, List[str]] = {}
for name, quantizer in sim.qc_quantize_op_dict.items():
if not quantizer.enabled:
continue
key = group_fn(name)
if key is None:
continue
groups.setdefault(key, []).append(name)
if not groups:
raise RuntimeError(
"No enabled quantizers found to analyze. Ensure the sim is "
"calibrated and has enabled quantizers."
)
logger.info("Analyzing per-quantizer sensitivity over %d groups", len(groups))
scores: Dict[str, float] = {}
# Disable every quantizer for the duration of the sweep (restored on exit),
# then enable one group at a time to isolate its contribution.
with disable_quantizers(sim, sim.qc_quantize_op_dict.keys()):
for key, names in tqdm(
groups.items(),
total=len(groups),
desc=f"Per-quantizer sensitivity ({metric.name})",
):
for name in names:
sim.qc_quantize_op_dict[name].enabled = True
try:
scores[key] = metric(sim.session)
finally:
for name in names:
sim.qc_quantize_op_dict[name].enabled = False
return _sorted_by_sensitivity(scores, metric)
def _sorted_by_sensitivity(
scores: Dict[str, float], metric: SensitivityMetric
) -> Dict[str, float]:
"""Return ``scores`` reordered most-sensitive-first per ``metric``."""
return dict(
sorted(
scores.items(),
key=lambda kv: kv[1],
reverse=metric.higher_is_worse,
)
)