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)[source]¶

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)[source]¶

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)[source]¶

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)[source]¶

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,
)
```
Pass group_by_op_name() to key the results by owning ONNX node name instead of by tensor name.

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

Naming¶

Quantizers are keyed by tensor name, which for exported models is often opaque (a torch Linear weight becomes an initializer named onnx::MatMul_9772). These helpers resolve the owning ONNX node name – e.g. /model/layers.0/self_attn/q_proj/MatMul – so results read in terms of the model’s structure and are keyed the way aimet_onnx.lite_mp.flip_layers_to_higher_precision() expects.

aimet_onnx.analysis.get_quantizer_op_names(sim)[source]¶

Map each quantizer name to the name of the ONNX node that owns it.

Quantizers are keyed by tensor name, which for exported models is frequently unreadable (onnx::MatMul_9772), while the node the tensor belongs to usually carries the module path it came from (/model/layers.0/self_attn/q_proj/MatMul). This resolves the latter for every quantizer in sim, mirroring the ownership rules of QuantizationSimModel.get_op_quantizers():

a parameter tensor belongs to the node it parameterizes;

an activation tensor belongs to the node that produces it;

a tensor with no producer (e.g. a graph input) belongs to the node that consumes it.

The returned names are connected-graph op names – the same keys accepted by aimet_onnx.lite_mp.flip_layers_to_higher_precision().

Parameters:: sim (QuantizationSimModel) – QuantizationSimModel to inspect.
Return type:: Dict[str, str]
Returns:: {quantizer_name: op_name}. Quantizers that belong to no op (rare) are absent from the mapping.

aimet_onnx.analysis.group_by_op_name(sim)[source]¶

Build a group_fn that keys sensitivity scores by owning ONNX node name.

Pass the result as the group_fn of analyze_per_quantizer_sensitivity() to get readable, node-named results instead of raw tensor names:

scores = analyze_per_quantizer_sensitivity(
    sim, metric, group_fn=group_by_op_name(sim)
)
flip_layers_to_higher_precision(sim, scores, percent_to_flip=10)

Quantizers that share an owning node are swept together as one group, so the result is per-node rather than per-quantizer wherever a node has more than one enabled quantizer (e.g. both a weight and an output quantizer). Quantizers with no owning node are skipped.

Parameters:: sim (QuantizationSimModel) – QuantizationSimModel to be analyzed.
Return type:: Callable[[str], Optional[str]]
Returns:: Callable mapping a quantizer name to its owning op name, or None if it has none.

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, details=None, details_label='Details')[source]¶

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.
details (Optional[Dict[str, str]]) – Optional {name: detail} dict. When given, the detail string is shown as an extra hover row and table column, letting the plot carry a secondary identifier – e.g. plot by ONNX node name (see aimet_onnx.analysis.get_quantizer_op_names()) while still showing the underlying quantizer tensor name. Names absent from the dict get an empty cell.
details_label (str) – Column and tooltip label used for details.

Return type:

None

aimet_onnx.analysis.save_sensitivity_results(scores, save_path='./sensitivity_results.json', details=None)[source]¶

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.
details (Optional[Dict[str, str]]) – Optional {name: detail} dict; a matching entry is written alongside each ranking entry as "detail". Use it to record what a name stands for, e.g. the quantizer tensor names behind a node name (see aimet_onnx.analysis.get_quantizer_op_names()).

Return type:

None

aimet_onnx.analysis.load_sensitivity_results(input_path)[source]¶

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.