Quantization II: PTQ, QAT & LLM Quantization

Chapter 1: Calibration Methods

If you can't precompute activation ranges analytically, you compute them empirically. Calibration means running a small representative dataset through your trained FP32 model, recording the activation distributions at every layer, and deriving the best S and Z from those statistics before deployment. No retraining required.

Method 1: Min/Max (naive)

The simplest approach: collect min and max across all calibration samples, and use those as r_min and r_max. This guarantees no clipping — every observed value fits in the INT8 range. The problem is that a single outlier in your calibration set can force S to accommodate it, wasting resolution on everyone else.

Method 2: Percentile Clipping

Instead of using the absolute min and max, use the 99.9th or 99.99th percentile. You're deliberately clipping a tiny fraction of extreme values, accepting a small clipping error in exchange for much better resolution across the bulk of the distribution.

Worked example. Activation vector: [0.1, 0.3, −0.5, 0.8, 0.2, −0.9, 0.4, 0.6, −0.2, 52.0]

• Min/max: r_max = 52.0 → S = 52.0/127 ≈ 0.409. The value 0.1 quantizes to round(0.1/0.409) = 0. Every value below 0.2 maps to zero!
• 99th percentile clip: r_max = 0.9 (ignoring the 52.0 outlier) → S = 0.9/127 ≈ 0.0071. The value 0.1 maps to round(0.1/0.0071) = 14. Excellent resolution — 52.0 just clips to INT8_MAX.

Method 3: KL-Divergence (Entropy) Calibration

Developed by NVIDIA for TensorRT. Instead of minimizing MSE or avoiding clipping, you minimize the information loss from quantization. The idea: the FP32 activation distribution is your "truth." The INT8 representation is an approximation. KL divergence measures how much information the approximation loses.

You sweep a threshold T from a small value up to the absolute max. For each T, you saturate all values above T, quantize to INT8, dequantize, and measure KL(FP32 hist || INT8 hist). The T that minimizes KL divergence is your winner. This typically clips a few percent of the distribution but dramatically reduces entropy loss.

Method 4: MSE Minimization (OCTAV/Newton-Raphson)

Directly minimize E[(x − Q(x))²] over the clipping threshold. Under a Laplace(0, b) distribution, the optimal threshold for N-bit quantization has closed-form solutions: |r|_max = 2.83b (2-bit), 3.89b (3-bit), 5.03b (4-bit). The Laplace parameter b is estimated from your calibration set.

Chapter 2: Granularity Revisited

The outlier problem has a structural dimension: outliers are not random. In LLMs, they cluster in specific channels — the same input channel index has an outlier every single token, while its neighbors are well-behaved. This motivates quantization granularity: how much of the tensor shares a single (S, Z) pair?

Per-Tensor Quantization

One S and one Z for the entire weight matrix. Works well for large models with well-behaved distributions. Fails badly when different output channels have very different dynamic ranges — a common occurrence in early layers of MobileNet-style architectures.

Per-Channel Quantization

One S and Z per output channel. If the weight matrix is W ∈ ℝ^{C_out × C_in}, you have C_out scale factors. Each channel's scale is set by its own range, so a channel with large weights doesn't tax channels with small weights.

Why per-channel is standard for weights but not activations. Weight channels are fixed at export time — the C_out scales are constants stored alongside the model. Activation channels are different at every token. You'd need to compute per-channel scales at runtime, which costs extra ops and disrupts the INT8 GEMM pipeline. So weights get per-channel, activations get per-tensor (or per-token in some schemes).

Group Quantization

Group quantization is a middle ground. Instead of one scale per channel, you split each channel into groups of g elements and give each group its own scale. For a weight at INT4 with group size g=128, the effective bitwidth is 4 + 16/128 = 4.125 bits (the 16-bit scale overhead amortized over 128 values). This is the key innovation in AWQ and many modern LLM quantization schemes.

Chapter 3: The Straight-Through Estimator

Post-training quantization is fast but has limits. For aggressive quantization (4-bit and below), or for small models with limited representational capacity, PTQ often can't recover the accuracy lost by discretizing weights and activations. The solution: Quantization-Aware Training (QAT) — simulate quantization during training so the model learns to work within its discrete constraints.

But there's an immediate problem. Quantization uses the round() function: Q(w) = S · round(w/S). The round() function has zero gradient almost everywhere (it's piecewise constant), and undefined gradient at the integers. Backpropagation through a wall of zeros learns nothing.

The gradient is zero, so the weight update is zero, so training stops. The model freezes in place.

The STE Solution: Pretend Round is Identity

Bengio (2013) proposed the Straight-Through Estimator (STE): in the backward pass, replace ∂Q(w)/∂w with 1 (the identity). You pass the gradient straight through the quantizer as if it weren't there.

Is this mathematically correct? No. The true gradient is zero; the STE gradient is 1. It's a biased estimator — it introduces approximation error. But it works in practice. The intuition: the STE says "the gradient I would get if the network were continuous is a good enough proxy for updating my real weights."

Full Forward/Backward Picture

In QAT, you maintain a full-precision shadow copy of the weights W_fp32. In the forward pass, you compute Q(W) and run the layer with quantized weights. Gradients flow back through Q using the STE, updating W_fp32. The small gradient steps accumulate in full precision without rounding error. At inference, you export only the quantized weights.

Chapter 4: Quantization-Aware Training

With the STE in hand, you can build a QAT training loop. The key concept is the fake-quantization node (also called simulated quantization): a differentiable approximation that behaves like a quantizer in the forward pass but allows gradient flow in the backward pass.

python
import torch
import torch.nn as nn

class FakeQuantize(torch.autograd.Function):
    """
    Forward: round to N-bit quantization grid (fake-quant).
    Backward: Straight-Through Estimator — pass gradient unchanged.
    """
    @staticmethod
    def forward(ctx, x, scale, zero_point, bits):
        q_min = -2 ** (bits - 1)
        q_max = 2 ** (bits - 1) - 1
        # Quantize: x → integer grid
        q = torch.clamp(torch.round(x / scale + zero_point), q_min, q_max)
        # Dequantize: back to float (still in forward, but now on grid)
        x_fake_quant = scale * (q - zero_point)
        # Save a mask: gradient passes through if |x| <= range
        ctx.save_for_backward(
            (x >= scale * q_min).float() * (x <= scale * q_max).float()
        )
        return x_fake_quant

    @staticmethod
    def backward(ctx, grad_output):
        mask, = ctx.saved_tensors
        # STE: pass gradient through where not clipped, else 0
        grad_input = grad_output * mask
        return grad_input, None, None, None


class QATLinear(nn.Module):
    """Linear layer with fake-quant on weights and activations."""
    def __init__(self, in_f, out_f, bits=8):
        super().__init__()
        self.linear = nn.Linear(in_f, out_f)
        self.bits = bits
        # Learnable scale (initialize from weight range)
        self.w_scale = nn.Parameter(
            (self.linear.weight.abs().max() / (2 ** (bits - 1) - 1)).reshape(1)
        )

    def forward(self, x):
        # Fake-quantize weights
        w_fq = FakeQuantize.apply(self.linear.weight, self.w_scale, 0, self.bits)
        # Fake-quantize activations (scale from per-batch absmax)
        a_scale = x.abs().max() / (2 ** (self.bits - 1) - 1)
        x_fq = FakeQuantize.apply(x, a_scale, 0, self.bits)
        return nn.functional.linear(x_fq, w_fq, self.linear.bias)

BatchNorm Folding for QAT

At inference, BatchNorm layers are almost always folded into the preceding convolution/linear layer (the mean/variance subtraction and scale/shift become constant modifications to W and b). For QAT to be accurate, you must simulate this folding during training — otherwise the QAT sees weights that will be different at inference.

Chapter 5: PTQ vs QAT: When to Use What

PTQ and QAT are not competing — they're tools for different situations. The decision tree comes down to three factors: bitwidth, model size, and whether you have access to training data and compute.

When PTQ is sufficient

For INT8 quantization on large models (ResNets, BERT, large Transformers), PTQ with good calibration (percentile or KL) typically achieves less than 0.5% accuracy drop. The model has enough representational capacity to absorb the quantization error. PTQ takes minutes; QAT takes days.

When QAT is necessary

For INT4 and below, or for small efficient models (MobileNet, EfficientNet-lite), PTQ degrades unacceptably. MobileNetV1 at INT8 per-tensor loses 11.8% accuracy with PTQ — that's the difference between a useful model and a useless one. QAT recovers it to 0.2% loss. The training cost is justified.

The Bitwidth Decision

A practical heuristic: use PTQ for INT8, consider QAT for INT4 and below. The transition point moves depending on model size — larger models can absorb INT4 PTQ (see GPTQ, AWQ for LLMs). Smaller models need QAT or at minimum per-channel calibration and careful range tuning.

Chapter 6: LLM Outliers & SmoothQuant

Large language models beyond 6.7B parameters develop a peculiar quantization pathology. Certain hidden-state channels consistently activate with values 10–100× larger than their neighbors. The phenomenon is systematic: the same channel indices spike across every token, every sequence, every context. It's a structural feature of how large Transformers represent information — not a distribution artifact you can calibrate away.

The LLM Transformer hidden state is Y = XW. In standard INT8, you'd quantize both X (activations) and W (weights) independently, then multiply the integers. But X has channels with max|X_j| ≈ 70 while others have max|X_k| ≈ 0.5. A single per-tensor scale of 70/127 ≈ 0.55 means channel k's values all round to zero. The matrix multiply is destroyed.

SmoothQuant: Migrate the Difficulty

The key insight from Lin et al. (2022): weights are easy to quantize, activations are hard to quantize — but we can balance this difficulty by migrating it from activations to weights. Weights can absorb more difficulty because they're quantized per-channel. Activations must use a single scale per-token.

The mathematical trick: introduce a per-channel diagonal scale matrix S. Because Y = XW, you can insert S and S⁻¹ without changing the result:

Now X̂ = X / s_j (divide each activation channel by s_j) and Ŵ = s_j × W_j,:  (multiply each corresponding weight row by s_j). The scale migrates the magnitude from activations to weights.

Deriving the Optimal Scale

How big should s_j be? Too large — weights become hard to quantize. Too small — activations stay hard. Song Han's group derived a formula with a migration strength parameter α ∈ [0, 1]:

Worked numbers. Say channel j has max|X_j| = 60 and max|W_j| = 2. With α = 0.5:

After smoothing: max|X̂_j| = 60 / 5.48 ≈ 10.9, max|Ŵ_j| = 2 × 5.48 ≈ 10.96. The outlier activation dropped from 60 to 10.9 (5.5× improvement). The weight only grew from 2 to 10.96 — still manageable with per-channel quantization. Both are now in a similar range, making per-token activation quantization and per-channel weight quantization both effective.

The α tradeoff. α = 0: all difficulty stays in activations (original problem). α = 1: all difficulty moves to weights (weights become huge, activations perfect). Sweet spot α ≈ 0.5 balances both, making the joint quantization error minimal. The optimal α can be tuned per-layer on calibration data.

python
import torch

def compute_smoothquant_scale(X_calibration, W, alpha=0.5):
    """
    X_calibration: collected activation tensor [n_tokens, d_in]
    W: weight matrix [d_out, d_in]
    Returns: per-channel scale s of shape [d_in]
    """
    # Per-channel max of activations (over all tokens)
    act_max = X_calibration.abs().max(dim=0).values  # [d_in]
    # Per-input-channel max of weight rows
    w_max = W.abs().max(dim=0).values               # [d_in]
    # Balancing scale
    s = act_max.pow(alpha) / w_max.pow(1 - alpha)
    s = s.clamp(min=1e-5)  # avoid division by zero
    return s

def apply_smooth_quant(X, W, s):
    """Fold the scale: smooth activations, absorb into weights."""
    X_smooth = X / s.unsqueeze(0)   # divide activation channels
    W_smooth = W * s.unsqueeze(0)   # multiply weight rows
    # Now quantize X_smooth and W_smooth independently with INT8
    return X_smooth, W_smooth

Chapter 7: GPTQ & AWQ

SmoothQuant enables W8A8 (both weights and activations quantized to INT8) — excellent for batch serving on data-center GPUs. But single-query LLM inference on a local machine has a different bottleneck: memory bandwidth. At batch size 1, the GPU compute is idle while it waits for 65B parameters to transfer from HBM to registers. W8A8 halves the memory but still needs a fast GPU. For edge devices, you need W4A16: 4-bit weights, 16-bit activations — only weights are compressed, and they're decompressed on-the-fly to FP16 for compute.

GPTQ: Second-Order Weight Quantization

GPTQ (Frantar et al., 2022) is a layer-wise PTQ method that finds the best 4-bit weights by minimizing the output error: for each layer, find the quantized Ŵ that minimizes ||WX − ŴX||_F². It uses the second-order Hessian information — the curvature of the loss surface — to make smarter rounding decisions than round-to-nearest.

Key insight: not all weight rounding errors are equal. A weight with high Hessian curvature (high second derivative) contributes more to output error when rounded badly than a weight with low curvature. GPTQ quantizes weights in order of decreasing importance (Hessian magnitude) and compensates remaining weights after each rounding to absorb the error.

AWQ: Activation-Aware Weight Quantization

AWQ (Lin et al., 2023) starts from a striking observation: 1% of weight channels are disproportionately important for model quality. Keeping just those 1% in FP16 while quantizing the remaining 99% to INT4 dramatically reduces perplexity. But which 1% to keep?

The answer: look at activations, not weights. Weight channels paired with activation channels that have large magnitudes are the most important — a large activation multiplied by a weight error produces a large output error. These salient channels are identified by running calibration data and measuring the average activation magnitude per input channel.

Instead of keeping salient channels in FP16 (which breaks hardware uniformity), AWQ uses the same scale trick as SmoothQuant, but just for the salient channels:

Scaling up W_salient by s before quantization concentrates the quantization grid around the important values, reducing their error. The inverse scale is applied to the corresponding activation, which is cheap since activations are large in those channels anyway.

Why does scaling up a weight channel reduce its quantization error? Quantization error is proportional to the step size Δ = max(|w|)/(2^N−1). When you scale up a channel by s before quantization, max(|w·s|) = s·max(|w|), and Δ' = s·Δ. But the output error is Δ'·(1/s)·x = Δ·x — the same as before in the unscaled domain. However, within the channel, other weights also got scaled by s, so their quantization is finer (more of the range is covered). For salient channels where x is large, the improvement is dramatic.

Cross-Layer Equalization (CLE) for PTQ

Before SmoothQuant, Song Han's group (and Markus et al., ICCV 2019) proposed Cross-Layer Equalization (CLE) for CNN PTQ. The idea: consecutive layers separated by a positive-homogeneous activation function (like ReLU) can be rescaled without changing the network output. ReLU is positive-homogeneous: ReLU(sx) = s·ReLU(x) for s > 0.

For two consecutive linear layers W⁽¹⁾ and W⁽²⁾ with a ReLU between them, introduce a per-channel scale s_c:

The optimal per-channel scale that equalizes quantization ranges is: s_c = √(max|W⁽¹⁾[c]| / max|W⁽²⁾[:,c]|). This geometric mean balances the per-channel range between the two layers. CLE is data-free (no calibration data needed) and can precede any other PTQ method.

PTQ Calibration Loop in Practice

python
import torch
from collections import defaultdict

def run_ptq_calibration(model, calibration_loader, bits=8):
    """Collect activation statistics and compute per-layer quant params."""
    model.eval()
    act_stats = defaultdict(lambda: {'min': float('inf'), 'max': float('-inf'), 'all': []})
    hooks = []

    for name, module in model.named_modules():
        if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d)):
            def make_hook(n):
                def hook(m, inp, out):
                    x = out.detach()
                    act_stats[n]['min'] = min(act_stats[n]['min'], x.min().item())
                    act_stats[n]['max'] = max(act_stats[n]['max'], x.max().item())
                    act_stats[n]['all'].extend(x.abs().flatten().tolist()[:512])  # sample
                return hook
            hooks.append(module.register_forward_hook(make_hook(name)))

    with torch.no_grad():
        for i, (x, _) in enumerate(calibration_loader):
            model(x)
            if i >= 32: break  # 32 batches is typically sufficient

    for h in hooks: h.remove()

    quant_params = {}
    q_max = (1 << (bits - 1)) - 1
    for name, stats in act_stats.items():
        # Percentile calibration (99.9th percentile)
        vals = sorted(stats['all'])
        thresh = vals[int(len(vals) * 0.999)]
        scale = thresh / q_max
        quant_params[name] = {'scale': scale, 'zero_point': 0}

    return quant_params

Chapter 8: Showcase: SmoothQuant & AWQ Explorer

This is the payoff chapter. Three linked canvases let you see the full quantization difficulty story, the SmoothQuant migration in action, and the AWQ salient-channel protection — all with live controls.

Chapter 9: Connections & Cheat Sheet

Complete Decision Guide: PTQ vs QAT vs LLM Methods

Calibration Methods Cheat Sheet

SmoothQuant One-Liner

Multiply each activation channel by s_j⁻¹ and each weight row by s_j, where s_j = max(|X_j|)^α/max(|W_j|)^1−α. The scaling is mathematically equivalent (Y = XW = (X/s)·(sW)) and is absorbed offline into weights and the preceding LayerNorm.

GPTQ One-Liner

Layer-wise second-order PTQ: quantize weights one-by-one in order of Hessian curvature, then compensate the remaining unquantized weights to absorb the just-introduced error. Enables INT4 on 175B GPT in ~4 GPU-hours.

AWQ One-Liner

Identify the 1% of input channels with the largest activation magnitudes. Scale those weight channels up by s before INT4 quantization (and scale the activations down correspondingly). The error on salient channels drops ∝ 1/s.

W8A8 vs W4A16 Guide

Related Lessons

• TinyML L5: Quantization I — affine quantization, K-means, scale/zero-point, quantized matmul. The foundation for this lesson.
• CS229S: Sparsity & Quantization — broader quantization ecosystem including mixed-precision and hardware perspectives.
• TinyML L7: Neural Architecture Search I — building architectures that need less compression by design.
• CS224N L9: PEFT & LoRA — QLoRA combines 4-bit quantization with LoRA fine-tuning for efficient LLM adaptation.

"What I cannot create, I do not understand." — Feynman. You've now built every piece of the quantization stack from scratch — the calibration loop, the fake-quant node, the STE backward pass, the SmoothQuant scale formula. You understand quantization.