aimet_onnx.experimental.adascale

Top level APIs

aimet_onnx.experimental.adascale.adascale_optimizer.apply_adascale(sim, inputs, adascale_model_config, num_iterations=1500, *, topology=None, loss_fn=None)
Parameters:

  • sim (QuantizationSimModel) – Quantization Sim model

  • inputs (Collection[Dict[str, ndarray]]) – (Collection[Dict[str, np.ndarray]]): The set of input samples to use during optimization.

  • adascale_model_config (AdaScaleModelConfig) – Adascale model config. There are pre-defined configs for Llama, Qwen2, Mistral, Qwen3, Phi3. For other models use AdaScaleModelConfig

  • num_iterations (int) – Number of iterations to optimize for during AdaScale

  • topology (Optional[LlmTopology]) – Decoder-stack topology of the model held by sim, from analyze_llm_topology(). AdaScale optimizes one decoder block at a time and takes the block boundaries from it — structure discovery belongs to llm_topology, not to AdaScale. Analyze the float model before building the sim: a topology describes the model rather than the sim, so it is derived once and reused. Optional today: when omitted the topology is discovered internally from the sim’s graph and a warning is raised. Passing it explicitly is recommended, and may become required in a future release.

  • loss_fn (Optional[Callable[[Tensor, Tensor, int], Tensor]]) – Loss function with signature loss_fn(fp_out, quant_out, data_idx) returning a scalar loss tensor. data_idx is the index of the current input in inputs order. Defaults to MSE loss if None.

Example usage:
>>> model = DummyModel()
>>> inputs = ...
>>> adascale_model_config = adascale_model_config['llama']
>>> topology = analyze_llm_topology(model)   # analyze the float model
>>> sim = QuantizationSimModel(model)
>>> apply_adascale(sim, inputs, adascale_model_config, num_iterations=num_iterations,
...                topology=topology)
>>> sim.compute_encodings(...)
>>> sim.export(...)

  1. apply_adascale modifies the weights in-place in the model

  2. compute encodings should not be called before the apply_adascale call

  3. Activation quantizers will remain uninitialized throughout the feature, and so compute encodings needs to be called by the user afterwards. This is so activation encodings will be computed with updated weights taken into account.

Warning: This feature is currently considered experimental pending API changes

class aimet_onnx.experimental.adascale.adascale_optimizer.AdaScaleModelConfig(model_type, beta_gamma_lr=0.001, scales_lr=0.0005)[source]

Model topology

AdaScale optimizes one decoder block at a time and takes those block boundaries from an LlmTopology, which describes the structure of the decoder stack. Analyze the float model before creating the sim, then pass the result as topology:

topology = analyze_llm_topology(onnx_model)
sim = QuantizationSimModel(onnx_model, ...)
apply_adascale(sim, inputs, adascale_model_config, topology=topology)
aimet_onnx.experimental.llm_topology.analyze_llm_topology(model, active_norms_per_block=2, expected_num_blocks=None, role_patterns=None, ir_model=None)[source]

Analyze model end-to-end and return a name-based LlmTopology.

Runs the whole pipeline on a private onnx_ir copy of model: strip quantizers, fuse RMSNorms, detect active norms and block boundaries, build the per-block topology, and infer hidden_size / head_dim.

Parameters:

  • model (ModelProto) – ONNX ModelProto to analyze. Not mutated.

  • active_norms_per_block (int) – Active norms per decoder block (see get_decoder_block_boundaries()). Defaults to 2.

  • expected_num_blocks (Optional[int]) – If given, validated against the detected count.

  • role_patterns (Optional[Dict[LinearRole, Pattern]]) – Optional module-name → role override (see classify_linear_role()).

  • ir_model (Optional[Model]) – Pre-built analysis IR for model — as returned by build_analysis_ir(), i.e. quantizer-stripped and RMSNorm-fused. Built here when None. Pass one only to avoid a second from_proto of a large model when the caller already holds it; a faithful (unfused) IR will not detect norms and must not be passed.

Return type:

LlmTopology

Returns:

LlmTopology with block/backbone roles, active_norms, hidden_size and head_dim populated. head_dim is None when the export exposes no past_value graph input to derive it from.

class aimet_onnx.experimental.llm_topology.LlmTopology(embed_tokens=<factory>, lm_head=<factory>, blocks=<factory>, past_key_input_names=<factory>, past_key_output_names=<factory>, past_value_input_names=<factory>, past_value_output_names=<factory>, active_norms=None, hidden_size=None, head_dim=None)[source]

Topology of an ONNX decoder-stack model: blocks + backbone-level roles + dims.

Parameters:

  • embed_tokens (List[str]) – Names of the token-embedding ``Gather``(s) that produce the initial residual-stream activations.

  • lm_head (List[str]) – Names of the vocabulary-projection linear(s) downstream of the final norm.

  • blocks (List[BlockTopology]) – Per-decoder-block topology in topological order.

  • past_key_input_names (List[str]) – Raw past_key_* graph inputs in declaration order, collected tolerantly (empty for prefill-only exports without a KV-cache). Pairing these to blocks and validating that their count matches the block count are the consumer’s responsibility (e.g. R3) — R1-only and prefill-only flows do not require KV-cache inputs.

  • past_key_output_names (List[str]) – Key-cache graph outputs in declaration order.

  • past_value_input_names (List[str]) – Value-cache graph inputs in declaration order.

  • past_value_output_names (List[str]) – Value-cache graph outputs in declaration order.

  • active_norms (Optional[List[ActiveNorm]]) – Active norms in topological order used to build the topology.

  • hidden_size (Optional[int]) – Residual-stream hidden dimension (None if not inferred; analyze_llm_topology() fills it).

  • head_dim (Optional[int]) – Per-head dimension (None when it could not be derived, e.g. an export without KV-cache inputs).