# Copyright (C) 2026 Embedl AB
"""High-level quantization entry points.
Convenience wrappers that chain configuration, precision-state setup, Q/DQ stub
insertion, and — for PTQ — calibration into single calls (the ``configure`` ->
``prepare_state`` -> ``prepare_stubs`` -> calibration flow).
"""
import copy
import re
import warnings
from collections.abc import Callable
from typing import Any
from torch import fx, nn
from embedl_deploy._internal.core.modules import FusedModule
from embedl_deploy._internal.core.plan import propagate_shapes
from embedl_deploy._internal.core.quantize.calibrate import (
calibrate_qdq,
calibrate_smooth_quant,
)
from embedl_deploy._internal.core.quantize.config import (
Precision,
QuantConfig,
SkipEntry,
)
from embedl_deploy._internal.core.quantize.prepare import (
prepare_state,
prepare_stubs,
)
from embedl_deploy._internal.core.quantize.stubs import QuantStub
from embedl_deploy._internal.core.quantize.utils import (
get_model_quants,
resolve_consumer,
)
def _no_skip_match(
mods: set[nn.Module],
skip_set: set[SkipEntry],
) -> bool:
"""Return ``True`` if no child of any module in `mods` matches `skip_set`."""
types = tuple(s for s in skip_set if isinstance(s, type))
instances = {s for s in skip_set if isinstance(s, nn.Module)}
patterns = [s for s in skip_set if isinstance(s, re.Pattern)]
for mod in mods:
for child in mod.children():
if child in instances:
return False
if types and isinstance(child, types):
return False
if patterns:
qname = getattr(child, "qualified_name", None)
if qname is not None and any(
p.fullmatch(qname) for p in patterns
):
return False
return True
def _configure_stub(
stub: QuantStub,
consumer: FusedModule,
config: QuantConfig,
skip_set: set[SkipEntry],
) -> None:
"""Configure a single stub and settle its consumer's precision state."""
consumer.precision_deferred.upstream_skip = False
stub.enabled = _no_skip_match(stub.consumers, skip_set)
if stub.enabled:
if not stub.fixed_parameters:
stub.config = copy.copy(config.activation)
if consumer.precision_synced.output:
consumer.output_precision = stub.config.precision
elif consumer.precision_synced.output:
consumer.output_precision = Precision.UNQUANTIZED
consumer.precision_deferred.upstream_skip = True
def configure(
model: fx.GraphModule,
config: QuantConfig,
) -> None:
"""Configure quantization and per-module precision state, in place.
Walks every quantizer collected from `model` and:
* Enables each
:class:`~embedl_deploy._internal.core.quantize.stubs.QuantStub` not
excluded by ``config.skip`` and copies ``config.activation`` onto it
(stubs with ``fixed_parameters`` keep their own config).
* Enables each ``weight_fake_quant`` not excluded by ``config.skip`` and
copies ``config.weight`` onto it.
* Enables each ``smooth_quant_observer`` not excluded by ``config.skip``
and copies ``config.smooth_quant`` onto it.
A module's stubs and weight form a lock-stepped set: when its weight is
synced (``precision_synced.weight``), a skip on either ``config.skip.stub``
or ``config.skip.weight`` demotes the whole set together, keeping it
uniform. A module with an independent weight skips each side on its own.
A synced-output module's ``output_precision`` is set to track its input
stubs — the stubs' precision when enabled, ``Precision.UNQUANTIZED`` when a
skip disables them. A module whose output is not synced
(``precision_synced.output`` ``False``) is left untouched, so an output it
set for itself in its constructor survives a skip.
When a skip demotes a synced output, ``precision_deferred.upstream_skip``
is set so the upstream sweep treats that output as an intentional user
decision and does not cascade demotion through it to downstream neighbors.
A structural fp output is not marked because the skip did not cause it.
Only module state is touched: no Q/DQ nodes are inserted and no other graph
surgery is run.
:param model:
A ``GraphModule`` produced by the fusion step. Modified in-place.
:param config:
The :class:`~embedl_deploy._internal.core.quantize.config.QuantConfig`
controlling precision and which module types to skip.
"""
mq = get_model_quants(model, include_enabled=True, include_disabled=True)
for stub in mq.stubs:
consumer = resolve_consumer(stub)
skip_set = set(config.skip.stub)
if consumer.precision_synced.weight:
skip_set |= config.skip.weight
_configure_stub(stub, consumer, config, skip_set)
for wfq in mq.weight:
consumer = resolve_consumer(wfq)
skip_set = set(config.skip.weight)
if consumer.precision_synced.weight:
skip_set |= config.skip.stub
wfq.enabled = _no_skip_match(wfq.consumers, skip_set)
if wfq.enabled:
wfq.config = copy.copy(config.weight)
for obs in mq.smooth:
obs.enabled = _no_skip_match(obs.consumers, config.skip.smooth)
if obs.enabled:
obs.config = copy.copy(config.smooth_quant)
[docs]
def quantize(
model: fx.GraphModule,
args: tuple[Any, ...],
config: QuantConfig | None = None,
*,
forward_loop: Callable[[fx.GraphModule], None] | None = None,
freeze_weights: bool = False,
) -> fx.GraphModule:
"""Configure, insert Q/DQ stubs, optimize, and optionally calibrate.
Chains :func:`~embedl_deploy._internal.core.quantize.main.configure` →
:func:`~embedl_deploy._internal.core.quantize.prepare.prepare_state` →
:func:`~embedl_deploy._internal.core.quantize.prepare.prepare_stubs`, then
— when `forward_loop` is given —
:func:`~embedl_deploy._internal.core.quantize.calibrate.calibrate_smooth_quant`
→ :func:`~embedl_deploy._internal.core.quantize.calibrate.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 :func:`~embedl_deploy._internal.core.quantize.qat.prepare_qat`
next).
:param model:
A ``GraphModule`` produced by the fusion step.
:param args:
The arguments to use for shape propagation necessary to get tensor meta
data required for calibration.
:param config:
Optional :class:`~embedl_deploy._internal.core.quantize.config.QuantConfig`.
Defaults to 8-bit symmetric.
:param 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).
:param 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
:func:`~embedl_deploy._internal.core.quantize.main.freeze_weight_quantization`
before export once training is complete.
:returns:
The quantized ``GraphModule``, with calibrated stubs when a
`forward_loop` was supplied.
"""
propagate_shapes(model, args)
config = config or QuantConfig()
configure(model, config)
prepare_state(model)
prepare_stubs(model)
if forward_loop is not None:
calibrate_smooth_quant(model, forward_loop)
calibrate_qdq(model, forward_loop)
if freeze_weights:
freeze_weight_quantization(model)
return model
[docs]
def freeze_weight_quantization(model: fx.GraphModule) -> None:
"""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
:func:`~embedl_deploy._internal.core.quantize.main.quantize` unless
``freeze_weights=False`` is passed. QAT users should call this explicitly
before ONNX export once training is complete.
"""
mq = get_model_quants(model)
for wfq in mq.weight:
mod = next(iter(wfq.consumers))
if not isinstance(mod, FusedModule):
continue
weight = mod.quantized_weight
if weight is not None:
wfq.freeze(weight)
else:
warnings.warn(
f"{type(mod).__name__} has a WeightFakeQuantize but "
f"quantized_weight returns None — override the property",
stacklevel=1,
)
