Quantization

Quantization techniques perform computations and store tensors in lower precision data types like 8-bit integer instead of floating point precision. There are multiple approaches to quantize a deep learning model categorized as:

• Post-training quantization (PTQ)
• Quantization aware training (QAT)

In post-training quantization, the model is trained in floating point precision and later converted to the lower precision data type. There are two types of post-training quantization:

1. Static quantization: quantizes the weights and activations of the model. Quantizing the activations statically requires data to be calibrated (i.e., recording the activation values to compute the optimal quantization parameters with representative data).
2. Dynamic quantization: quantized the weights ahead of time (like static quantization) but the activations are dynamically at runtime.

Sometimes post-training quantization is not able to achieve acceptable task accuracy. In general, this is where quantization-aware training (QAT) can be used: during training, fake-quantization modules are inserted in the forward and backward passes to simulate quantization effects, allowing the model to learn representations that are more robust to reduced precision.

Burn does not currently support QAT. Only post-training quantization (PTQ) is implemented at this time.

Module Quantization

Quantizing the weights of your model after training is quite simple. We have access to the weight tensors and can collect their statistics, such as the min and max value when using MinMaxCalibration, to compute the quantization parameters.

use burn::module::Quantizer;
use burn::tensor::quantization::{Calibration, QuantScheme, QuantValue, ScaleDtype};

// Quantization config
let scheme = QuantScheme::default()
    .per_block([32], ScaleDtype::F16)
    .with_value(QuantValue::Q4F);
let mut quantizer = Quantizer::new(Calibration::MinMax, scheme);

// Quantize the weights
let model = model.quantize_weights(&mut quantizer);

Calibration

Calibration is the step during quantization where the range of all floating-point tensors is computed. This is pretty straightforward for weights since the actual range is known at quantization-time (weights are static), but activations require more attention.

To compute the quantization parameters, Burn supports the following Calibration methods.

Quantization Scheme

A quantization scheme defines how an input is quantized, including the representation of quantized values, storage format, granularity, and how the values are scaled.

#![allow(unused)]
fn main() {
let scheme = QuantScheme::default()
    .with_mode(QuantMode::Symmetric)          // Quantization mode
    .per_block([2, 16], ScaleDtype::F16)      // One scale per block of values, stored as f16
    .with_value(QuantValue::Q8S)              // Data type of quantized values, independent of how they're stored
    .with_store(QuantStore::Native);          // Storage format for quantized values
}

A scheme carries up to two scale levels, each set in any order and each owning the type its scales are stored in. per_tensor is one scale for the whole tensor, which is also what a scheme with no level set resolves to; per_block is one scale per block of values.

#![allow(unused)]
fn main() {
// Block scales in ue4m3, normalized by a single per-tensor f32 scale.
let scheme = QuantScheme::default()
    .per_block([16], ScaleDtype::UE4M3)
    .per_tensor(ScaleDtype::F32);
}

Two levels exist so the block scales can live in a narrow type: the per-tensor scale absorbs the tensor's dynamic range, and the block type only has to cover the spread between blocks.

Quantization Mode

Scale Levels

Block dimensions are relative to the tensor rank and describe per-axis rectangles, not flat chunks of contiguous values. For example, [16] divides the trailing dimension into groups of 16, while [2, 16] divides the final two dimensions into rectangular blocks. Dynamic calibration currently requires every affected tensor dimension to be evenly divisible by its block extent.

Setting both nests the blocks inside the tensor: the per-tensor scale is the factor the block scales are relative to, and it has to be F32.

Quantization Value

Quantization Store

Native storage is not supported for sub-byte quantization values.

Scale Data Type

A narrower scale type stores less per block, but it also has a much smaller range. The smallest positive UE4M3 value is 2^-9; Burn rounds a smaller positive scale up to 2^-9 rather than down to zero. The resulting scale can nevertheless be so coarse that the block's values quantize to zero. A two-level scheme lets the per-tensor F32 scale carry the absolute magnitude, but the differences between block scales can still exceed UE4M3's range and make smaller blocks quantize too coarsely.