# Quantization - tensorplay.ao.quantization

Source: https://www.tensorplay.cn/docs/quantization.html

tensorplay.ao is the model optimization toolkit, and quantization is its
largest part: it converts a float model to run arithmetic on 8-bit integers
(roughly a 4x shrink in memory for weights), trading a controlled amount of
accuracy for it. The quantization stack lives in
tensorplay.ao.quantization, with the modules it produces under
tensorplay.ao.nn and the weight-pruning and numeric-comparison utilities
in sibling packages (see [pruning](/docs/pruning.html)).

Three workflows are supported, in increasing order of effort:

- Dynamic quantization — weights are quantized up front; activations are quantized on the fly at each use. One call, no calibration data: [quantize_dynamic()](/docs/generated/tensorplay.ao.quantization.quantize_dynamic.html#tensorplay.ao.quantization.quantize_dynamic).

- Static (post-training) quantization — weights and activations are both quantized ahead of time: insert observers with [prepare()](/docs/generated/tensorplay.ao.quantization.prepare.html#tensorplay.ao.quantization.prepare), calibrate on representative data, then bake the scales in with [convert()](/docs/generated/tensorplay.ao.quantization.convert.html#tensorplay.ao.quantization.convert).

- Quantization-aware training (QAT) — the same stubs, but with fake quantize in the training loop, so the weights adapt to the quantization noise before [convert()](/docs/generated/tensorplay.ao.quantization.convert.html#tensorplay.ao.quantization.convert) runs.

```
import tensorplay as tp
from tensorplay.ao.quantization import quantize_dynamic

model = tp.nn.Sequential(
    tp.nn.Linear(128, 64),
    tp.nn.ReLU(),
    tp.nn.Linear(64, 10),
)

# dynamic quantization: Linear weights go to int8, biases stay in float
qmodel = quantize_dynamic(model)
```

A quantized tensor carries its own scale and zero point; converting between
the quantized and float views is the job of the eager operators below, and
float_model - qmodel round trips are what the numeric-suite helpers
measure (see the last section).

## Eager operators

| tensorplay.ao.quantization.quantize_per_tensor |quantize_per_tensor(Tensor self, float scale, int zero_point, ScalarType dtype) -> Tensor quantize_per_tensor.tensor_qparams(Tensor self, Tensor scale, Tensor zero_point, ScalarType dtype) -> Tensor quantize_per_tensor.tensors(Tensor[] tensors, Tensor scales, Tensor zero_points, ScalarType dtype) -> Tensor[] |
| --- | --- |
| tensorplay.ao.quantization.quantize_per_channel | quantize_per_channel(Tensor self, Tensor scales, Tensor zero_points, int axis, ScalarType dtype) -> Tensor |
| tensorplay.ao.quantization.quantize_per_tensor_dynamic | quantize_per_tensor_dynamic(Tensor self, ScalarType dtype, bool reduce_range) -> Tensor |
| tensorplay.ao.quantization.dequantize | dequantize.self(Tensor self) -> Tensor dequantize.tensors(Tensor[] tensors) -> Tensor[] |
| tensorplay.ao.quantization.int_repr | int_repr(Tensor self) -> Tensor |
| tensorplay.ao.quantization.is_quantized | is_quantized(Tensor self) -> bool |
| tensorplay.ao.quantization.q_scale | q_scale(Tensor self) -> float |
| tensorplay.ao.quantization.q_zero_point | q_zero_point(Tensor self) -> int |
| tensorplay.ao.quantization.q_per_channel_scales | q_per_channel_scales(Tensor self) -> Tensor |
| tensorplay.ao.quantization.q_per_channel_zero_points | q_per_channel_zero_points(Tensor self) -> Tensor |
| tensorplay.ao.quantization.q_per_channel_axis | q_per_channel_axis(Tensor self) -> int |
| tensorplay.ao.quantization.qscheme | qscheme(Tensor self) -> int |
| tensorplay.ao.quantization.fake_quantize_per_tensor |  |
| tensorplay.ao.quantization.fake_quantize_per_channel |  |
| tensorplay.ao.quantization.quantized_linear | quantized_linear(Tensor input, Tensor weight, float input_scale, int input_zero_point, Tensor weight_scales, Tensor weight_zero_points, Tensor? bias=None, float out_scale=1.0, int out_zero_point=0) -> Tensor |
| tensorplay.ao.quantization.quantized_linear_dynamic | quantized_linear_dynamic(Tensor input, Tensor weight, Tensor weight_scales, Tensor weight_zero_points, Tensor? bias=None, bool reduce_range=False) -> Tensor |

- [quantize_per_tensor()](/docs/generated/tensorplay.ao.quantization.quantize_per_tensor.html#tensorplay.ao.quantization.quantize_per_tensor) / [quantize_per_channel()](/docs/generated/tensorplay.ao.quantization.quantize_per_channel.html#tensorplay.ao.quantization.quantize_per_channel) map a float tensor to the integer range with one (scale, zero_point) pair, or one pair per channel along axis; [quantize_per_tensor_dynamic()](/docs/generated/tensorplay.ao.quantization.quantize_per_tensor_dynamic.html#tensorplay.ao.quantization.quantize_per_tensor_dynamic) derives the scale from the tensor’s own range instead of a calibrated value.

- [dequantize()](/docs/generated/tensorplay.ao.quantization.dequantize.html#tensorplay.ao.quantization.dequantize) converts back to float; [int_repr()](/docs/generated/tensorplay.ao.quantization.int_repr.html#tensorplay.ao.quantization.int_repr) shows the stored integer values.

- The q_* accessors read the quantization parameters back off a quantized tensor, and [qscheme()](/docs/generated/tensorplay.ao.quantization.qscheme.html#tensorplay.ao.quantization.qscheme) reports which scheme (per-tensor or per-channel) it uses.

- fake_quantize_* simulate the rounding without leaving float — they are what training-time (QAT) quantization is built on.

- quantized_linear / quantized_linear_dynamic are the integer GEMM kernels behind quantized linear layers (static with calibrated input scales, dynamic with per-call input quantization).

## Observers and fake quantize

| tensorplay.ao.quantization.MinMaxObserver | |
| --- | --- |
| tensorplay.ao.quantization.PerChannelMinMaxObserver |  |
| tensorplay.ao.quantization.HistogramObserver |  |
| tensorplay.ao.quantization.MovingAverageMinMaxObserver |  |
| tensorplay.ao.quantization.MovingAveragePerChannelMinMaxObserver |  |
| tensorplay.ao.quantization.PlaceholderObserver |  |
| tensorplay.ao.quantization.FixedQParamsObserver |  |
| tensorplay.ao.quantization.FakeQuantize |  |
| tensorplay.ao.quantization.PerChannelFakeQuantize |  |
| tensorplay.ao.quantization.FixedQParamsFakeQuantize |  |
| tensorplay.ao.quantization.FusedMovingAvgObsFakeQuantize |  |
| tensorplay.ao.quantization.default_observer |  |
| tensorplay.ao.quantization.default_weight_observer |  |
| tensorplay.ao.quantization.default_dynamic_quant_observer |  |
| tensorplay.ao.quantization.get_observer_state_dict |  |
| tensorplay.ao.quantization.load_observer_state_dict |  |

Observers watch tensors pass through the model during calibration and
produce the (scale, zero_point) the conversion step needs.
MinMaxObserver tracks the running min/max (the default);
HistogramObserver keeps a histogram of values, giving a better scale when
outliers matter; the MovingAverage* variants smooth the min/max; and
PlaceholderObserver records nothing (for flows where parameters are
already known). FakeQuantize wraps an observer with the fake-quantize op,
which is the module QAT inserts into the graph.
[get_observer_state_dict()](/docs/generated/tensorplay.ao.quantization.get_observer_state_dict.html#tensorplay.ao.quantization.get_observer_state_dict) /
[load_observer_state_dict()](/docs/generated/tensorplay.ao.quantization.load_observer_state_dict.html#tensorplay.ao.quantization.load_observer_state_dict) save and
restore calibration results, so a calibration run need not be repeated.

Statistics, epsilon, histogram counts, simulation flags and quantization
parameters are registered buffers. They follow device moves and survive
state_dict checkpoints. Channel statistics resize when a checkpoint is
loaded into a fresh observer. Histogram calibration redistributes existing
counts when the observed range expands and searches the clipping range by
integrating reconstruction error over the source bins.

disable_observer() freezes statistics while retaining the simulated
rounding. disable_fake_quant() returns the floating input while allowing
statistics to continue updating. The backward mask uses the rounded grid
position before saturation, including ties and values outside the range.
_LearnableFakeQuantize, available from
tensorplay.ao.quantization._learnable_fake_quantize, additionally trains
scale and zero point with optional gradient scaling.

## Configuration and the eager workflow

| tensorplay.ao.quantization.QConfig |Describes how to quantize a layer: observer factories for activations and weights. |
| --- | --- |
| tensorplay.ao.quantization.QConfigMapping | Pattern-based assignment of qconfigs to modules. |
| tensorplay.ao.quantization.default_qconfig | observer factories for activations and weights. |
| tensorplay.ao.quantization.default_dynamic_qconfig | observer factories for activations and weights. |
| tensorplay.ao.quantization.default_per_channel_qconfig | observer factories for activations and weights. |
| tensorplay.ao.quantization.default_weight_only_qconfig | observer factories for activations and weights. |
| tensorplay.ao.quantization.QuantStub |  |
| tensorplay.ao.quantization.DeQuantStub |  |
| tensorplay.ao.quantization.prepare | Attach observers to float modules for post-training calibration. |
| tensorplay.ao.quantization.prepare_qat |  |
| tensorplay.ao.quantization.convert | Swap calibrated float modules for their quantized counterparts. |
| tensorplay.ao.quantization.quantize | Prepare, calibrate through run_fn(*run_args) and convert. |
| tensorplay.ao.quantization.quantize_dynamic | Swap modules for their dynamic-quantization counterparts. |
| tensorplay.ao.quantization.quantize_qat |  |
| tensorplay.ao.quantization.fuse_modules | Fuse a list of modules into a single module. |
| tensorplay.ao.quantization.fuse_modules_qat | Quantization-aware training variant of fuse_modules(). |
| tensorplay.ao.quantization.get_default_qconfig |  |
| tensorplay.ao.quantization.get_default_qat_qconfig |  |

A [QConfig](/docs/generated/tensorplay.ao.quantization.QConfig.html#tensorplay.ao.quantization.QConfig) pairs an activation observer
with a weight observer (the default_* functions provide the common
combinations); a [QConfigMapping](/docs/generated/tensorplay.ao.quantization.QConfigMapping.html#tensorplay.ao.quantization.QConfigMapping)
assigns QConfigs per submodule type or name. QuantStub / DeQuantStub
mark where activations enter and leave the quantized region.

The static workflow composes these:

```
import tensorplay as tp
from tensorplay.ao.quantization import (
    QuantStub, DeQuantStub, prepare, convert, default_per_channel_qconfig,
)

model = tp.nn.Sequential(
    QuantStub(),
    tp.nn.Linear(128, 64),
    tp.nn.ReLU(),
    tp.nn.Linear(64, 10),
    DeQuantStub(),
)
model.qconfig = default_per_channel_qconfig   # which observers to attach

m = prepare(model)               # attaches observers to the stubs and layers
for x, _ in calibration_data:    # run representative data through
    m(x)
model = convert(m)               # observers become scale constants,
                                 # float layers become quantized layers
```

Only submodules carrying a qconfig are quantized, and the model must
enter and leave its quantized region through QuantStub / DeQuantStub
— after conversion the stubs become the real Quantize / DeQuantize
modules and the interior runs on quantized tensors end to end.
[quantize()](/docs/generated/tensorplay.ao.quantization.quantize.html#tensorplay.ao.quantization.quantize) runs
prepare→calibrate (run_fn)→convert in one call for eager models;
[quantize_dynamic()](/docs/generated/tensorplay.ao.quantization.quantize_dynamic.html#tensorplay.ao.quantization.quantize_dynamic) needs no calibration
step at all. [fuse_modules()](/docs/generated/tensorplay.ao.quantization.fuse_modules.html#tensorplay.ao.quantization.fuse_modules) merges
adjacent layers (Conv+BN+ReLU, Linear+ReLU) into the single
fused modules below, which quantize better than the pieces separately.

For training, use prepare_qat on a model in training mode. Use
fuse_modules_qat before preparation to retain batch normalization during
training. Conv1d/2d/3d and Linear support weight simulation; fused Conv+BN,
Conv+BN+ReLU, Conv+ReLU, Linear+BN and Linear+ReLU modules retain their
training behavior. Freezing normalization statistics is independent of
freezing quantization statistics.

```
from tensorplay.ao.quantization import (
    prepare_qat, convert, default_per_channel_qat_qconfig,
    disable_observer,
)

model.train()
model.qconfig = default_per_channel_qat_qconfig
trained = prepare_qat(model)
optimizer = tp.optim.SGD(trained.parameters(), lr=0.01)
for x, target in training_data:
    optimizer.zero_grad()
    loss = loss_fn(trained(x), target)
    loss.backward()
    optimizer.step()
trained.apply(disable_observer)
inference = convert(trained.eval(), use_precomputed_fake_quant=True)
```

Construct the optimizer from trained.parameters(). Conversion folds
normalization parameters into the convolution or linear weights. Passing
use_precomputed_fake_quant=True preserves the trained weight grid during
conversion, including configured code bounds.

## Packed execution and checkpoints

Quantized Linear and Conv1d/2d/3d retain canonical quantized weights and a
reusable native execution cache. set_weight_bias rebuilds the cache when
weights change. Each inference call uses the activation tensor’s scale and
zero point. Weights support per-tensor and per-output-channel affine grids;
activations support signed and unsigned 8-bit storage. Grouped convolution
handles stride, padding and dilation in integer coordinates.

CPU execution uses tiled input gathering and AVX2 integer dot products when
available. CUDA execution uses shared operand tiles and packed four-byte
integer dot products. Both accumulate affine corrections in 64-bit integers.
LinearReLU and ConvReLU apply bias, output rounding and rectification in the
compute epilogue. BNReLU applies normalization and rectification in one
native kernel. Dynamic Linear and LinearReLU derive one activation grid
over the complete input on each call and return floating outputs.

state_dict saves canonical weights, bias and output quantization
parameters. Execution caches are nonpersistent and rebuilt when loading or
moving a module between devices. Saving a complete module also excludes
execution caches and rebuilds them during restoration. Create a module
with the same dimensions and convolution parameters before loading a state
dictionary; whole-module checkpoints retain those constructor parameters.
Whole-module loading requires registering the module classes with
tensorplay.serialization.safe_globals. State dictionaries load with the
default restricted reader.

Quantized tensor checkpoints preserve integer storage, strides, offsets and
the complete affine grid. Both ZIP and legacy stream formats support
per-tensor, per-channel and floating-zero-point channel grids. map_location
can restore their storage and metadata onto another device.

benchmark/bench_quantized_packing.py measures cached execution, rebuilding
the weights on every call and fused versus separate rectification. It checks
integer output equality before timing. The following measurements use a Core
Ultra 7 265K with four CPU threads and an RTX 3090. Each value is the median
of five batches of 100 calls, with device synchronization around CUDA timing.
All inputs, weights and outputs use signed 8-bit affine grids.

| Device |M × N × K |Direct (μs) |Cached (μs) |Rebuild (μs) |Direct / cached |
| --- | --- | --- | --- | --- | --- |
| CPU | 1 × 512 × 1024 | 244.6 | 21.1 | 213.2 | 11.6× |
| CPU | 32 × 256 × 512 | 520.6 | 48.1 | 113.0 | 10.8× |
| CPU | 256 × 512 × 1024 | 16361.0 | 1312.5 | 1495.4 | 12.5× |
| CUDA | 1 × 512 × 1024 | 108.8 | 16.0 | 99.8 | 6.8× |
| CUDA | 32 × 256 × 512 | 59.0 | 9.8 | 99.5 | 6.0× |
| CUDA | 256 × 512 × 1024 | 967.8 | 46.5 | 131.1 | 20.8× |

The direct baseline is the existing scalar integer linear operator. Fusion
removes a separate activation dispatch and output allocation; its total
latency depends on shape and device. For the 32-row case, fused and separate
rectification measured 58.3 and 55.2 μs on CPU, and 10.7 and 11.6 μs on CUDA.
The complete measurements are in benchmark/quantized_packing_results.json.

## Quantized modules

| tensorplay.ao.nn.quantized.Linear | |
| --- | --- |
| tensorplay.ao.nn.quantized.QuantizedLinear |  |
| tensorplay.ao.nn.quantized.Conv1d |  |
| tensorplay.ao.nn.quantized.Conv2d |  |
| tensorplay.ao.nn.quantized.Conv3d |  |
| tensorplay.ao.nn.quantized.BatchNorm2d | Quantized 2d batch norm over a quantized input. |
| tensorplay.ao.nn.quantized.BatchNorm3d | Quantized 3d batch norm over a quantized input. |
| tensorplay.ao.nn.quantized.LayerNorm | Normalize trailing dimensions and requantize with output qparams. |
| tensorplay.ao.nn.quantized.GroupNorm | Normalize channel groups and retain the input quantized dtype. |
| tensorplay.ao.nn.quantized.InstanceNorm1d | Normalize each channel's spatial elements of a batched 1D input. |
| tensorplay.ao.nn.quantized.InstanceNorm2d | Normalize each channel's spatial elements of a batched 2D input. |
| tensorplay.ao.nn.quantized.InstanceNorm3d | Normalize each channel's spatial elements of a batched 3D input. |
| tensorplay.ao.nn.quantized.ReLU | Rectified linear unit on a quantized tensor. |
| tensorplay.ao.nn.quantized.ReLU6 | Rectified linear unit clamped at 6 on a quantized tensor. |
| tensorplay.ao.nn.quantized.ELU | Exponential linear unit on a quantized tensor. |
| tensorplay.ao.nn.quantized.LeakyReLU | Leaky rectified linear unit on a quantized tensor. |
| tensorplay.ao.nn.quantized.Hardswish | HardSwish activation on a quantized tensor. |
| tensorplay.ao.nn.quantized.Hardsigmoid | HardSigmoid activation on a quantized tensor. |
| tensorplay.ao.nn.quantized.Sigmoid | Logistic sigmoid on a quantized tensor. |
| tensorplay.ao.nn.quantized.Tanh | Hyperbolic tangent on a quantized tensor. |
| tensorplay.ao.nn.quantized.MaxPool1d | Quantized 1d max pooling. |
| tensorplay.ao.nn.quantized.MaxPool2d | Quantized 2d max pooling. |
| tensorplay.ao.nn.quantized.MaxPool3d | Quantized 3d max pooling. |
| tensorplay.ao.nn.quantized.Quantize | Quantizes an incoming float tensor. |
| tensorplay.ao.nn.quantized.DeQuantize | Dequantizes an incoming tensor through its own affine parameters. |
| tensorplay.ao.nn.quantized.FloatFunctional | Float-domain operation wrappers with the quantized call signature. |
| tensorplay.ao.nn.quantized.QFunctional | Wrapper class for quantized operations with fixed output qparams. |
| tensorplay.ao.nn.quantized.FXFloatFunctional | Floating operation wrappers for graph-captured workflows. |
| tensorplay.ao.nn.quantized.dynamic.Linear |  |

These are the integer kernels convert maps float layers onto: the
Linear/Conv* compute layers, quantized BatchNorm, the
activation functions, and MaxPool. Quantize/DeQuantize are the module
forms of the eager ops at the boundary stubs. FloatFunctional /
QFunctional wrap arithmetic (add, cat, mul, …) so the same model code
runs its ops through quantized or float kernels;
FXFloatFunctional is the graph-capture-friendly variant.

LayerNorm, GroupNorm, and InstanceNorm1d/InstanceNorm2d/InstanceNorm3d
use native CPU and CUDA kernels. They accumulate statistics from integer codes
and requantize each normalized value without allocating a float copy of the
input. They accept per-tensor affine qint8, quint8, and qint32 inputs,
preserve the input dtype, and store output scale and zero_point as buffers
in the module’s state dictionary. Optional affine parameters use float32.
Group and instance normalization require batched inputs of at least three
dimensions; a four-dimensional channels-last input keeps that layout.
Instance normalization computes statistics from each input, including in
evaluation mode.

The default activation observer declares quint8 with a calibration range
of [0, 127]; the default weight observer declares qint8. The dtype refers
to the quantized tensor, while int_repr() exposes its plain integer codes.

## Fused modules

| tensorplay.ao.nn.intrinsic.LinearReLU |A linear transformation followed by a rectifier. |
| --- | --- |
| tensorplay.ao.nn.intrinsic.ConvReLU1d | A 1d convolution followed by a rectifier. |
| tensorplay.ao.nn.intrinsic.ConvReLU2d | A 2d convolution followed by a rectifier. |
| tensorplay.ao.nn.intrinsic.ConvReLU3d | A 3d convolution followed by a rectifier. |
| tensorplay.ao.nn.intrinsic.BNReLU2d | A 2d batch norm followed by a rectifier. |
| tensorplay.ao.nn.intrinsic.BNReLU3d | A 3d batch norm followed by a rectifier. |

fuse_modules produces these single-kernel combinations (Conv+ReLU,
Linear+ReLU, BN+ReLU), and convert then maps them onto their
fused quantized counterparts — one quantized kernel where the float model
had two or three layers.

## Quantization-aware training modules

| tensorplay.ao.nn.qat.Linear |A linear module with simulated per-output-channel weight quantization. |
| --- | --- |
| tensorplay.ao.nn.qat.Conv1d |  |
| tensorplay.ao.nn.qat.Conv2d |  |
| tensorplay.ao.nn.qat.Conv3d |  |

For QAT, tensorplay.ao.nn.qat provides the float modules with fake
quantize built in: ao.nn.qat.Linear and ao.nn.qat.Conv1d/Conv2d/
Conv3d quantize their weights on the fly through a
weight_fake_quant module (a per-output-channel PerChannelFakeQuantize,
or whatever qconfig.weight() produces), so training optimizes against
the quantized grid while the parameters stay float. The fused variants
(ao.nn.intrinsic.qat.LinearReLU, ConvReLU1d/2d/3d) add a rectifier on
top.

A QAT round trip: swap the float layer for its QAT twin, let prepare
attach the activation observers, train, then convert — the quantized
module that comes out bakes in exactly the grid the training loop
converged on:

```
import tensorplay as tp
from tensorplay.ao.nn.qat import Linear as QATLinear
from tensorplay.ao.quantization import (
    prepare, convert, default_per_channel_qconfig,
)

float_linear.qconfig = default_per_channel_qconfig
qat_linear = QATLinear.from_float(float_linear)
qat_linear.qconfig = default_per_channel_qconfig

qat_linear = prepare(qat_linear)     # observers now record activations
for x, _ in train_loader:             # ordinary training loop; the forward
    loss = loss_fn(qat_linear(x), y)  # pass runs through weight_fake_quant
    loss.backward()
q_linear = convert(qat_linear)        # QuantizedLinear with the trained grid
```

## Numeric comparison utilities

The tensorplay.ao.ns subpackage holds the helpers for measuring what
quantization actually cost: compare_weights (float vs quantized weight
tables side by side), prepare_model_with_stubs + Shadow +
compare_model_stub (a shadow copy of the float model runs next to the
quantized one), and prepare_model_outputs + compare_model_outputs
(per-activation comparisons). They are reached through the
tensorplay.ao.ns._numeric_suite module.

## Where to go next

- [pruning](/docs/pruning.html) — the other model-optimization axis: sparsity instead of precision.

- [nn](/docs/nn.html) — the float modules the quantized stack starts from.

- [the main namespace](/docs/tensorplay.html) — the tensor APIs the eager quantize operators are part of.