`embedl_deploy.quantize` package#

Module contents:

Public quantization API.

Users import from here:

from embedl_deploy.quantize import quantize, QuantConfig, TensorQuantConfig

class embedl_deploy.quantize.CalibrationMethod(value)[source]#

Bases: Enum

Algorithm for collecting activation statistics during PTQ calibration.

Each member’s value is the torch.ao observer class instantiated and attached to a QuantStub on the first calibration forward pass.

HISTOGRAM = <class 'torch.ao.quantization.observer.HistogramObserver'>#: Build a histogram and search for the optimal quantization range.

MINMAX = <class 'torch.ao.quantization.observer.MinMaxObserver'>#: Track the global minimum and maximum (default, fastest).

MOVING_AVERAGE_MINMAX = <class 'torch.ao.quantization.observer.MovingAverageMinMaxObserver'>#: Exponential moving average of min/max — less sensitive to outliers.

Bases: object

Specifies which modules or module types to leave disabled during configure.

Each set accepts three kinds of entries:

A module type (e.g. nn.Conv2d) — matches any module of that type.
A module instance — matches that exact module.
A compiled regex (re.Pattern[str]) — matched via re.fullmatch against the original qualified name of each child module (set during recomposition from nn_module_stack).

smooth: set[type[Module] | Module | Pattern[str]]#: Modules to leave smooth quantization disabled for.

stub: set[type[Module] | Module | Pattern[str]]#: Modules to leave stub quantization disabled for.

weight: set[type[Module] | Module | Pattern[str]]#: Modules to leave weight quantization disabled for.

class embedl_deploy.quantize.Precision(value)[source]#

Bases: Enum

A quantized tensor format: signed bit width, negative for float.

The magnitude of value is the bit width and the sign encodes the format (negative = floating point), so INT8 and FP8 share n_bits == 8. Width and format are read only through the n_bits / is_float properties — a plain Enum defines no ordering, so comparisons like FP8 < INT4 raise TypeError instead of silently comparing the underlying ints.

Extend this enum to add a format (e.g. FP16 = -16) when demand appears; do not introduce a parallel bit-width scheme.

FP8 = -8#: 8-bit floating point.

INT16 = 16#: 16-bit signed integer.

INT4 = 4#: 4-bit signed integer.

INT8 = 8#: 8-bit signed integer.

UNQUANTIZED = 0#: No quantized domain; format-agnostic (not pinned to fp16/fp32).

property is_float: bool#: True for floating-point formats (encoded as a negative value).

property n_bits: int#: Bit width of the format (INT8 and FP8 both report 8).

class embedl_deploy.quantize.QuantConfig(activation: TensorQuantConfig = <factory>, weight: TensorQuantConfig = <factory>, smooth_quant: SmoothQuantConfig = <factory>, skip: ModulesToSkip = <factory>)[source]#

Bases: object

Top-level quantization configuration.

Bundles separate TensorQuantConfig instances for activations and weights so that each can be configured independently.

Parameters:

activation – Settings for activation (inter-layer) quantization.
weight – Settings for weight quantization (Conv kernels only; bias is always left in floating-point).
skip – Specifies which modules or module types to leave disabled during configure. See ModulesToSkip.

activation: TensorQuantConfig#: Settings for activation quantization.

skip: ModulesToSkip#: Modules or module types to leave disabled during configure.

smooth_quant: SmoothQuantConfig#: SmoothQuant settings applied during calibration.

weight: TensorQuantConfig#: Settings for weight quantization.

class embedl_deploy.quantize.QuantStub(consumers: set[Module], precision: Precision = Precision.INT8, symmetric: bool = True, calibration_method: CalibrationMethod = CalibrationMethod.MINMAX, *, fixed_calibration: tuple[float, int] | None = None)[source]#

Bases: Module

Quantize a floating-point tensor.

During calibration the module delegates statistics collection to a torch.ao observer selected by calibration_method. After calibration, scale and zero_point are derived from the observer and used by torch.fake_quantize_per_tensor_affine() in the forward pass.

Parameters:

consumers – Set of modules that consume this stub’s output.
precision – Quantized format (a Precision, default INT8).
symmetric – Symmetric or asymmetric quantization.
calibration_method – Algorithm used to collect activation statistics. Defaults to MINMAX.
fixed_calibration – Fixed (scale, zero_point) tuple. When provided, calibration will not override the values.

compute_parameters() → None[source]#

Derive scale and zero_point from the observer.

Raises:: RuntimeError – If no data was observed during calibration.

forward(x: Tensor) → Tensor[source]#: Fake-quantize x, updating observer stats if calibrating.

scale: Tensor#

zero_point: Tensor#

class embedl_deploy.quantize.SmoothQuantConfig(alpha: float = 0.5)[source]#

Bases: object

SmoothQuant migration settings.

Controls the per-channel weight/activation redistribution applied by calibrate_smooth_quant().

Parameters:: alpha – Migration strength in [0, 1]. 0 keeps all difficulty on activations; 1 pushes it entirely to weights.

alpha: float = 0.5#: Migration strength in [0, 1].

class embedl_deploy.quantize.TensorQuantConfig(precision: Precision = Precision.INT8, symmetric: bool = True, per_channel: bool = True, calibration_method: CalibrationMethod = CalibrationMethod.MINMAX)[source]#

Bases: object

Quantization settings for a single tensor class (activation or weight).

Parameters:

precision – The quantized format as a Precision member. Must be a quantized integer format (INT4/INT8); UNQUANTIZED and floating-point formats are rejected at construction.
symmetric – When True the quantized range is centered on zero (zero_point = 0). When False an asymmetric range is used, allowing better coverage of distributions that are not centered on zero.
per_channel – When True and when used for weight quantization, a separate scale/zero-point is computed along the output-channel axis (axis 0). Defaults to True for production quantization of Conv/Linear weights.

calibration_method: CalibrationMethod = <class 'torch.ao.quantization.observer.MinMaxObserver'>#: Calibration algorithm used by QuantStub to collect activation statistics. Only relevant for activation configs; weight quantization computes scale/zero-point on-the-fly.

per_channel: bool = True#: Per-channel quantization (weight only).

precision: Precision = 8#: Quantized format (integer width; floating-point not yet supported).

property quant_max: int#: Maximum representable integer value.

property quant_min: int#: Minimum representable integer value.

symmetric: bool = True#: Symmetric (zero_point = 0) or asymmetric range.

class embedl_deploy.quantize.WeightFakeQuantize(consumers: set[Module], precision: Precision = Precision.INT8, symmetric: bool = True, per_channel: bool = True, *, channel_axis: int = 0)[source]#

Bases: Module

Fake-quantize a weight tensor during the forward pass.

Unlike QuantStub (which requires a calibration pass), this module computes scale and zero_point on-the-fly from the weight tensor. Correct for QAT where weights change each step.

Parameters:

consumers – Set of modules that consume this module’s output.
precision – Quantized format (a Precision).
symmetric – Symmetric (zero_point = 0) or asymmetric.
per_channel – Use per-channel quantization along channel_axis.
channel_axis – The axis along which per-channel scales are computed (default 0, i.e. output channels).

consumers: set[Module]#

forward(weight: Tensor) → Tensor[source]#: Fake-quantize weight using stored (frozen) or on-the-fly scale.

freeze(weight: Tensor) → None[source]#

Compute scale/zero_point from weight and store as constant buffers.

After calling this, forward() uses the stored constants instead of recomputing them each call. ONNX export will therefore emit a Constant node for the scale rather than the full Abs → ReduceMax → Div arithmetic, which TensorRT requires for explicit-quantization fusion.

Parameters:: weight – The weight tensor to calibrate against (typically conv.weight).

scale: Tensor | None#

zero_point: Tensor | None#

embedl_deploy.quantize.calibrate_qdq(model: GraphModule, forward_loop: Callable[[GraphModule], None]) → None[source]#

Calibrate Q/DQ stubs by running the user’s forward loop.

Temporarily disables all enabled QuantStub and WeightFakeQuantize modules, switches non-fixed stubs into calibration mode, invokes forward_loop once, then finalizes scale / zero_point from the observed min/max ranges and re-enables all modules.

The model is modified in-place.

Parameters:

model – A configured GraphModule whose fused modules have been set up by configure().
forward_loop –
(model) -> None callable that runs representative data through the model. The caller controls batch size, device placement, and iteration count. Example:
```
def forward_loop(model):
    for batch in calib_loader:
        model(batch)
```

Raises:

ValueError – If the model contains no enabled QuantStub or WeightFakeQuantize modules.
RuntimeError – If any stub did not observe finite values during the loop.
Exception – Propagates any exception from forward_loop after restoring model state.

embedl_deploy.quantize.calibrate_smooth_quant(model: GraphModule, forward_loop: Callable[[GraphModule], None]) → None[source]#

Calibrate and apply SmoothQuant to a fused model in-place.

Migrates quantization difficulty from activations to weights for every LayerNorm → Linear pair that has an enabled SmoothQuantObserver with populated downstream_linears. Instruments each observer’s layer_norm with forward hooks, temporarily disables all enabled stubs and weight fake-quantize modules, and runs forward_loop on the model.

Must be called after transform() (fusion) and configure(), and before calibrate_qdq().

Parameters:

model – A fused GraphModule whose observers have been enabled by configure.
forward_loop – (model) -> None callable that runs representative data through the model. The caller controls batch size, device placement, and iteration count.

Raises:

Exception – Propagates any exception from forward_loop after restoring model state.

embedl_deploy.quantize.disable_fake_quant(model: Module) → Module[source]#

Disable fake quantization throughout the model.

Parameters:: model – The model to modify in-place.
Returns:: The same model, for method chaining.

embedl_deploy.quantize.enable_fake_quant(model: Module) → Module[source]#

Enable fake quantization in all stubs and weight quantizers.

Parameters:: model – The model to modify in-place.
Returns:: The same model, for method chaining.

embedl_deploy.quantize.freeze_bn_stats(model: Module) → Module[source]#

Freeze BatchNorm running statistics.

Puts all BatchNorm*d layers into eval mode so running_mean and running_var are no longer updated. Affine parameters remain trainable.

Parameters:: model – The model to modify in-place.
Returns:: The same model, for method chaining.

embedl_deploy.quantize.freeze_weight_quantization(model: GraphModule) → None[source]#

Freeze all WeightFakeQuantize scale/zero_point buffers.

After calibration (or QAT training) the weights are fixed for export, so we compute the scale once and store it as a constant buffer. This ensures ONNX export emits a Constant node for the scale rather than the dynamic Abs → ReduceMax → Div arithmetic, which TensorRT requires for explicit-quantization Q/DQ fusion.

This is called automatically by quantize() unless freeze_weights=False is passed. QAT users should call this explicitly before ONNX export once training is complete.

embedl_deploy.quantize.prepare_qat(model: Module) → Module[source]#

Prepare a quantized model for quantization-aware training.

Sets the model to training mode so that QuantStub nodes propagate gradients through STE and WeightFakeQuantize nodes apply fake-quant to weights.

Parameters:: model – A quantized nn.Module.
Returns:: The same model, in-place, for method chaining.

embedl_deploy.quantize.quantize(model: GraphModule, args: tuple[Any, ...], config: QuantConfig | None = None, *, forward_loop: Callable[[GraphModule], None] | None = None, freeze_weights: bool = False) → GraphModule[source]#

Configure, insert Q/DQ stubs, optimize, and optionally calibrate.

Chains configure() → prepare_state() → prepare_stubs(), then — when forward_loop is given — calibrate_smooth_quant() → calibrate_qdq().

With a forward_loop this is the PTQ entry point: the stubs are calibrated from representative data. Without one it is the QAT entry point: the stubs are inserted and configured but left uncalibrated for a subsequent training loop (call prepare_qat() next).

Parameters:

model – A GraphModule produced by the fusion step.
args – The arguments to use for shape propagation necessary to get tensor meta data required for calibration.
config – Optional QuantConfig. Defaults to 8-bit symmetric.
forward_loop – Optional (model) -> None callable that runs representative data through the model. The caller controls batch size, device placement, and iteration count. When omitted, calibration is skipped (QAT).
freeze_weights – When True, weight scales are computed from the current weights and stored as constant buffers after calibration. This is required for ONNX/TensorRT export (PTQ workflow). Defaults to False to preserve the original behavior (dynamic on-the-fly scale computation, suitable for QAT). Call freeze_weight_quantization() before export once training is complete.

Returns:

The quantized GraphModule, with calibrated stubs when a forward_loop was supplied.

embedl_deploy.quantize package#

`embedl_deploy.quantize` package#