algo_cfg#

Compile a quantize config into an ordered list of scoped calibration stages.

This is the compile half of the calibration plan (the calibration_plan mode in mode is the execute half). It is deliberately side-effect free: it reads the already-quantized model’s structure — quantizer and linear names — to resolve globs and validate, but it mutates nothing, runs no forward and touches no data. Consequences the design leans on:

  • bad configs fail fast, before any expensive calibration runs;

  • it is testable without running a model;

  • the resulting plan is a pure function of (config, model structure), so it is identical on every rank — which is what keeps predicate scoping from desynchronizing collectives in distributed calibration.

Both surfaces lower here. algorithm="max" becomes the single all-"*" stage, so the legacy whole-model path is a special case of the scoped one rather than a second engine.

Classes

AlgoCapabilities

What one calibration algorithm consumes, produces and needs to run.

AlgoStage

One algorithm applied to one scope — the unit of work in a calibration plan.

Functions

compile_algo_cfg

Lower a quantize config into an ordered, validated list of scoped stages.

describe_plan

Human-readable plan dump, used by the demos and for debugging.

plan_hash

A stable hash of the plan.

stage_predicate

Build the should_process write-mask for a stage.
class AlgoCapabilities#

Bases: object

What one calibration algorithm consumes, produces and needs to run.

Only the fields the compiler actually uses today are declared. produces / requires are a small open vocabulary of state tokens:

weight, weight_amax, input_amax, pre_quant_scale, acts.

__init__(granularity, role, requires=frozenset({}), produces=frozenset({}), needs_forward=True, self_forwards=False, requires_absent=frozenset({}))#
Parameters:

  • granularity (str)

  • role (str)

  • requires (frozenset[str])

  • produces (frozenset[str])

  • needs_forward (bool)

  • self_forwards (bool)

  • requires_absent (frozenset[str])

Return type:

None

granularity: str#
needs_forward: bool = True#
produces: frozenset[str] = frozenset({})#
requires: frozenset[str] = frozenset({})#
requires_absent: frozenset[str] = frozenset({})#
role: str#
self_forwards: bool = False#
property shareable_forward: bool#

Whether this stage could ride a forward pass shared with other stages.

exception AlgoCfgValidationError#

Bases: ValueError

Raised when an algo_cfg cannot be lowered into a valid plan.

class AlgoStage#

Bases: object

One algorithm applied to one scope — the unit of work in a calibration plan.

__init__(algo, cfg, scope, selector, order, entry, exclude=())#
Parameters:

  • algo (str | None)

  • cfg (dict)

  • scope (str)

  • selector (str)

  • order (int)

  • entry (int)

  • exclude (tuple[tuple[str, str], ...])

Return type:

None

algo: str | None#
property capabilities: AlgoCapabilities | None#

Declared capabilities of this stage’s algorithm, or None if undeclared.

Looked up rather than copied onto the stage: capabilities describe the algorithm, so a stage that carried its own copy could drift from the registry.

cfg: dict#
entry: int#
exclude: tuple[tuple[str, str], ...] = ()#
key()#

Execution-relevant identity, used for the plan hash.

entry is deliberately excluded: it records where in the config a stage came from, which is provenance, not behaviour. Leaving it out is what makes the legacy algorithm="max" plan and the explicit [{"quantizer_name": "*", "cfg": ["max"]}] plan hash identical – the same execution, written two ways.

Return type:

tuple

order: int#
scope: str#
selector: str#
compile_algo_cfg(config, model=None, strict=True)#

Lower a quantize config into an ordered, validated list of scoped stages.

Pure: reads model structure only, mutates nothing, runs no forward.

Parameters:

  • config (QuantizeConfig | dict) – a QuantizeConfig or a mapping with algo_cfg / algorithm.

  • model (Module | None) – the already-quantized model. Required for the model-aware validation (scope resolution, roles, fused siblings, dependencies); when None only the config-only rules run.

  • strict (bool) – raise AlgoCfgValidationError on violations. False downgrades them to warnings, which is what lets a knowingly-broken pipeline be run for demonstration purposes.

Returns:

The ordered plan. Stages run in list order.

Return type:

list[AlgoStage]

describe_plan(plan, model=None)#

Human-readable plan dump, used by the demos and for debugging.

Parameters:

  • plan (list[AlgoStage])

  • model (Module | None)

Return type:

str

plan_hash(plan)#

A stable hash of the plan.

Distributed calibration is safe only if every rank runs the same stages over the same scopes — otherwise a predicate skips a quantizer on one rank and its amax all-reduce never matches, which hangs. Comparing this hash across ranks turns that silent deadlock into a clear error.

Parameters:

plan (list[AlgoStage])

Return type:

str

stage_predicate(model, stage)#

Build the should_process write-mask for a stage.

The predicate is AND-ed into each algorithm’s existing is_enabled filter, so a stage writes only its targets and never toggles enable-state — reads (and hence the activations seen by search-based algorithms like AWQ and GPTQ) are unchanged.

Parameters:
Return type:

Callable[[str], bool]