aimet_onnx.analysis.sensitivity

Warning

This feature is under heavy development and API changes may occur without notice in future versions.

Sensitivity analysis measures how sensitive a calibrated QuantizationSimModel is to quantization at per-quantizer granularity, and drives mixed-precision decisions.

The analysis takes a SensitivityMetric – a named evaluation function (onnxruntime.InferenceSession -> float) plus ranking semantics – and returns a {name: score} dict ordered most-sensitive-first. That dict feeds directly into aimet_onnx.lite_mp.flip_layers_to_higher_precision() to raise the most sensitive units to a higher precision. To restrict the sweep to a subset (e.g. KV-cache tensors only), pass a group_fn that returns None for the quantizers to exclude.

Metric

class aimet_onnx.analysis.SensitivityMetric(name, eval_fn, higher_is_worse=True)

A named evaluation function with ranking semantics.

Parameters:

• name (str) – Human-readable metric name (used in logs and plots).

• eval_fn (Callable[[InferenceSession], float]) – Callable taking an onnxruntime.InferenceSession and returning a scalar score.

• higher_is_worse (bool) – 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.

sensitivity_score(score)

Map a raw score to a value where larger always means more sensitive.

Return type:

float

aimet_onnx.analysis.make_topk_logit_psnr_metric(fp_session, inputs, k=10, logit_output_index=0)

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.

Parameters:

• fp_session (InferenceSession) – ORT inference session for the floating-point model.

• inputs (Iterable[Dict[str, ndarray]]) – Iterable of input feed dicts ({input_name: np.ndarray}).

• k (int) – Number of top FP vocabulary indices to compare per position.

• logit_output_index (int) – Index of the logits tensor in the model outputs.

Return type:

SensitivityMetric

Returns:

A SensitivityMetric wrapping the top-k logit PSNR eval fn.

Analysis

aimet_onnx.analysis.analyze_per_quantizer_sensitivity(sim, metric, group_fn=None)

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 aimet_onnx.analyze_per_layer_sensitivity().

Parameters:

• sim (QuantizationSimModel) – Calibrated QuantizationSimModel to analyze. Its quantizers’ enabled state is restored on return.

• metric (SensitivityMetric) – SensitivityMetric used to score each group.

• group_fn (Optional[Callable[[str], Optional[str]]]) –

  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 type:

Dict[str, float]

Returns:

Dict mapping group key to its metric score, ordered most-sensitive-first per metric.

For op-level sensitivity, see aimet_onnx.analyze_per_layer_sensitivity().

Visualization and persistence

The {name: score} dict returned by the analysis functions can be rendered as an interactive HTML chart or serialized to / from JSON.

aimet_onnx.analysis.save_sensitivity_plot(scores, metric, save_path='./sensitivity_plot.html', highlight_patterns=None)

Render an interactive sensitivity chart and save it as standalone HTML.

Points are plotted along the x-axis in the order of scores (so passing a dict in topological/layer order plots them that way); the y-axis is the metric score. The sensitivity rank (most-sensitive-first) is still shown in the hover tooltip and table. Interactive features:

• A numeric threshold input; entries scoring below it are listed in a table and a guide line is drawn.

• Toggle buttons (one per highlight_patterns entry) that highlight points whose name contains the pattern and add them to the table.

Parameters:

• scores (Dict[str, float]) – {name: score} dict from an analysis function. Iteration order of this dict sets the x-axis order of the plotted points.

• metric (SensitivityMetric) – The SensitivityMetric used to produce scores.

• save_path (str) – Output HTML path (must end with .html).

• highlight_patterns (Optional[List[str]]) – Name substrings for the highlight toggles. If omitted, defaults to common transformer projection weight names.

Return type:

None

aimet_onnx.analysis.save_sensitivity_results(scores, save_path='./sensitivity_results.json')

Save sensitivity scores to JSON, preserving their order.

scores from the analysis functions is already ranked most-sensitive-first; that order is preserved in the ranking list. Inverse of load_sensitivity_results().

Parameters:

• scores (Dict[str, float]) – {name: score} dict from an analysis function.

• save_path (str) – Output JSON path.

Return type:

None

aimet_onnx.analysis.load_sensitivity_results(input_path)

Load a sensitivity JSON written by save_sensitivity_results().

Parameters:

input_path (str) – Path to the JSON file.

Return type:

Dict[str, float]

Returns:

{name: score} dict, preserving the file’s ranked order.