Quantization - tensorplay.ao.quantization

Eager operators

• quantize_per_tensor() / 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() derives the scale from the tensor’s own range instead of a calibrated value.

• dequantize() converts back to float; int_repr() shows the stored integer values.

• The q_* accessors read the quantization parameters back off a quantized tensor, and 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

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() / 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

A QConfig pairs an activation observer with a weight observer (the default_* functions provide the common combinations); a 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() runs prepare→calibrate (run_fn)→convert in one call for eager models; quantize_dynamic() needs no calibration step at all. 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.

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

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

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

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 — the other model-optimization axis: sparsity instead of precision.

• nn — the float modules the quantized stack starts from.

• the main namespace — the tensor APIs the eager quantize operators are part of.