Quantization, Pruning, and Compression

Memory. Bytes per weight: fp32 4, bf16 2, int8 1, int4 0.5 plus scales (§29.3). A 7B model is 13 GiB in bf16, about 3.4 GiB in int4 with group-128 fp16 scales.

Bandwidth. Decoding at small batch reads every weight once per token and does about 2 FLOPs per weight read, far below a GPU's ridge point. Latency ≈ weight bytes / bandwidth, so halving bytes nearly halves latency even if the kernel dequantizes to bf16 (when quantization pays).

Compute. int8 and fp8 tensor cores run about twice the bf16 rate, but only when both operands are low precision, so compute-bound prefill gains only from activation quantization or 2:4 sparsity.

An affine quantizer maps reals onto integer codes \([q_{\min}, q_{\max}]\) with a scale \(s > 0\) and integer zero point \(z\):

Asymmetric (uint8). Widen the observed \([\alpha, \beta]\) to include 0 and map it onto \([0, 255]\): \(s = (\beta - \alpha)/255\), \(z = \operatorname{round}(-\alpha/s)\). Real 0.0 is then exactly code \(z\), so zero padding and ReLU zeros quantize without bias (Jacob et al., 2018).

Symmetric (int8). Fix \(z = 0\) and \(s = \max|x|/127\), codes in \([-127, 127]\). Dropping \(-128\) keeps \(x \mapsto -x\) exact. Symmetric suits roughly centred weights. Asymmetric suits one-sided post-ReLU activations, where a symmetric grid would waste half its codes.

Error. In range, rounding moves \(x/s\) by at most \(\tfrac12\), so \(|x - \hat x| \le s/2\). If the error is uniform on \([-s/2, s/2]\), its MSE is \(\frac{1}{s}\int_{-s/2}^{s/2} e^2\,de = s^2/12\). Out of range, clipping error is unbounded. Calibration is that trade-off: a max-based range never clips but inflates \(s\); a 99.99th-percentile or MSE-optimal range clips a few outliers to shrink \(s\) for everything else (design space).

By hand. \(x = [-1.0, 0.0, 0.5, 2.0]\), asymmetric uint8: \(s = 3/255 \approx 0.011765\), \(z = \operatorname{round}(1/s) = 85\).

\(0.5/s = 42.5\) is a tie, and torch.round rounds half to even (42), so the error is exactly \(s/2\). Fence 1 checks this table against torch.quantize_per_tensor.

One scale per tensor lets the largest value set \(s\) for everyone. In fence 2, one weight row scaled ×50 leaves the other seven rows using codes no larger than ±5 of ±127, and their MSE is about 1,000× worse than with per-row scales. LLM activations have exactly such outlier channels (§29.9).

Per-channel. Linear.weight is (out_features, in_features) and output \(y_o = \sum_i W_{oi} x_i\) uses one row. A (out,) scale vector, broadcast as (out, 1), factors out of each dot product, so it costs nothing in the matmul (§29.4). "Per-channel" in PyTorch's eager API and here means per output channel, axis 0.

Per-group. int4 has 16 codes, so one scale per row is too coarse. Groups of \(g\) consecutive in-features (32, 64, 128) each get a scale, which now varies inside the dot product, so kernels rescale partial sums per group. For a (4096, 4096) weight with fp16 scales:

Asymmetric groups add a zero point each. torchao's Int4WeightOnlyConfig is uint4 with a floating-point zero point per group.

Packing. torch.int4 exists in 2.8 but reports itemsize 1, so two codes are packed into one uint8: the two's-complement low nibble (q & 0xF) of the even element goes in bits 0–3, the odd one in bits 4–7, and unpacking masks, shifts and sign-extends (u - 16 when u >= 8). Production kernels interleave nibbles for tensor-core fragments (torchao's tensor_core_tiled layout for tinygemm; README); only the byte order differs. Fence 3 round-trips the simple layout.

With symmetric weights (scale \(s_w\) per row) and activations (\(s_x\)):

The sum is integer arithmetic: int8 products accumulated in int32, since 128 products of up to \(127^2\) already need 22 bits. One float multiply per output rescales. With an activation zero point \(z_x\):

one int32 per row, not a second matmul. Weight zero points would add a per-token \(z_w \sum_i q^x_i\) term, one reason weights stay symmetric. The lab's StaticInt8Linear is exactly these two lines. The int8 × int8 → int32 GEMM is exposed as the private torch._int_mm, which ran on CPU in 2.8.

Weight-only kernels instead load codes, dequantize to bf16 in registers, and run a bf16 matmul: less bandwidth, same FLOPs. Right for decode, useless for compute-bound prefill.

Native quantized tensors. torch.quantize_per_tensor(x, s, z, torch.quint8) and torch.quantize_per_channel(w, scales, zero_points, axis=0, dtype=torch.qint8) return tensors with .int_repr(), .q_scale(), .q_per_channel_scales() and .dequantize() (API reference). Quantized modules run on the global quantized engine, torch.backends.quantized.engine: fbgemm or x86 on x86, qnnpack on ARM. On the reference arm64 Mac, supported_engines was ['qnnpack', 'none'] with default 'none', so portable code picks from supported_engines and restores the old value in finally.

Dynamic quantization stores int8 weights and quantizes each activation batch from its own range at run time, with no calibration. Engine-guarded, since arm64 defaults to 'none':

from torch.ao.quantization import quantize_dynamic     # DeprecationWarning on call in 2.8
old = torch.backends.quantized.engine
try:
    torch.backends.quantized.engine = next(e for e in ("fbgemm", "x86", "qnnpack")
                                           if e in torch.backends.quantized.supported_engines)
    qmodel = quantize_dynamic(model, {nn.Linear}, dtype=torch.qint8)
finally:
    torch.backends.quantized.engine = old

Each Linear becomes a DynamicQuantizedLinear (source) with a per-tensor qint8 weight. The call emitted DeprecationWarning: torch.ao.quantization is deprecated and will be removed in 2.10, whose text in torch/ao/quantization/utils.py is also the migration map:

The 2.8 quantization page still documents eager and FX; read it for concepts only.

Static PTQ fixes activation scales ahead of time from observers run on calibration batches: MinMaxObserver (extremes), MovingAverageMinMaxObserver (EMA of them), HistogramObserver (error-minimizing range, so deliberate clipping), PerChannelMinMaxObserver (weights). The eager flow marks the int region with QuantStub/DeQuantStub, calls prepare, calibrates, and calls convert. First fuse Conv-BN-ReLU: BatchNorm folds into conv weights (the folding in module 24), and fusing ReLU removes a requantization between ops.

The lab implements the algorithm instead: a forward pre-hook on each Linear feeds a small MinMaxObserver over 8 batches. The maintained path is PT2E: export, let a backend Quantizer annotate, calibrate, convert (torchao v0.13.0 paths):

from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e, convert_pt2e
from torchao.quantization.pt2e.quantizer.x86_inductor_quantizer import (
    X86InductorQuantizer, get_default_x86_inductor_quantization_config)

ep = torch.export.export(model.eval(), (example_x,))
quantizer = X86InductorQuantizer().set_global(get_default_x86_inductor_quantization_config())
prepared = prepare_pt2e(ep.module(), quantizer)   # inserts observers
for xb in calibration_batches:
    prepared(xb)
quantized = convert_pt2e(prepared)                 # explicit quantize/dequantize nodes

A backend (Inductor on x86, XNNPACK in ExecuTorch, module 28) pattern-matches dequantize → op → quantize into an integer kernel.

QAT inserts fake quantization, \(\mathrm{fq}(x) = s\,(\operatorname{clamp}(\operatorname{round}(x/s)+z) - z)\) computed in float, on weights and activations in the forward pass, so the loss sees rounding and weights settle where it hurts less. Conversion is as in PTQ.

But round's gradient is 0 almost everywhere, and autograd returns exactly that (fence 4). The straight-through estimator (Bengio et al., 2013) treats round as identity in backward but keeps the clamp:

Outside the range it is 0: nudging a clipped value does not change the output. A custom autograd.Function (module 08) returns grad * inside. The idiom x + (fq(x) - x).detach() has forward fq(x) but gradient 1 everywhere, so it needs a mask, as the lab's fake_quant_ste adds with torch.where. torch.fake_quantize_per_tensor_affine and ..._per_channel_affine have the masked STE built in. LSQ (Esser et al., 2020) also learns \(s\) through the STE, with its gradient scaled by \(1/\sqrt{N q_{\max}}\); it matters most at 2–4 bits.

torchao's QAT (QAT README; IntXQuantizationAwareTrainingConfig is the older name), not executed here:

from torchao.quantization import quantize_, Int8DynamicActivationInt4WeightConfig
from torchao.quantization.qat import QATConfig

base = Int8DynamicActivationInt4WeightConfig(group_size=32)
quantize_(model, QATConfig(base, step="prepare"))   # Linear -> FakeQuantizedLinear
train_loop(model)
quantize_(model, QATConfig(base, step="convert"))   # real quantization with base

torchao is not installed in this course's CI, so this section is fragments against v0.13.0:

from torchao.quantization import (quantize_, Int8WeightOnlyConfig, Int4WeightOnlyConfig,
    Int8DynamicActivationInt8WeightConfig, Float8DynamicActivationFloat8WeightConfig, PerRow)

quantize_(model, Int8WeightOnlyConfig())                    # A16W8 per-channel
quantize_(model, Int4WeightOnlyConfig(group_size=128))      # A16W4, tinygemm layout
quantize_(model, Int8DynamicActivationInt8WeightConfig())   # A8W8 dynamic
quantize_(model, Float8DynamicActivationFloat8WeightConfig(granularity=PerRow()))
model = torch.compile(model)

Apply one config per model. Each makes linear.weight a tensor subclass (AffineQuantizedTensor and successors) of codes, scales and layout that dispatches F.linear to a fused kernel (module 20); torch.compile fuses the dequantization. Int4WeightOnlyConfig(use_hqq=True) picks scales with HQQ's optimizer instead of min/max. torchao.float8.convert_to_float8_training swaps Linears for fp8 training (float8 README; formats in module 17 and Micikevicius et al., 2022). Block-scaled MX formats appear with Triton in module 25.

torchao's README reports for Llama-3.1-8B decode on an H100 (vendor-reported): bf16 126.9 tokens/s, int8 weight-only 198.9, int4 group-64 241.4, with wikitext perplexity 7.54, 7.56 and 8.44. Its A100 Llama-2-7B rows show int8dq at 9.6 tokens/s against bf16's 107.4: activation quantization buys compute that a bandwidth-bound batch-1 loop does not need. The README gives no cause for an 11× slowdown, which is too large to be the quantization arithmetic alone; read it as "this kernel path is not tuned for decode". Match the regime to the bottleneck and measure the kernel you will ship.

LLMs break naive int8 through activation outliers: a few hidden dimensions are consistently tens of times larger than the rest.

• LLM.int8() (Dettmers et al., 2022) splits the matmul by hidden dimension: every feature column of \(X\) holding at least one value with magnitude above a threshold \(\alpha\) (the paper finds \(\alpha = 6.0\) sufficient), together with the matching rows of \(W\), runs in fp16; the rest, about 99.9 % of dimensions, runs in int8 with vector-wise scales (per row of \(X\), per column of \(W\)).
• SmoothQuant (Xiao et al., 2023) rewrites \(Y = (X\,\mathrm{diag}(s)^{-1})(\mathrm{diag}(s)\,W)\) with \(s_j = \max|X_j|^\alpha / \max|W_j|^{1-\alpha}\). The migration strength \(\alpha = 0.5\) is the paper's sweet spot for OPT and BLOOM; GLM-130B, with stronger outliers, needs 0.75. The product is unchanged, both factors become int8-friendly, and \(s^{-1}\) folds into the preceding LayerNorm.
• GPTQ (Frantar et al., 2023) quantizes columns one at a time and updates the remaining columns to cancel each column's error, using the inverse Hessian of \(H = 2XX^\top\) from calibration activations.
• AWQ (Lin et al., 2024) starts from the observation that keeping only 0.1–1 % of weight channels in fp16, chosen by activation magnitude rather than weight magnitude, recovers most of int3/int4's loss. Mixed precision is awkward for kernels, so AWQ instead scales those channels up before quantization (and the activations down), with a per-channel scale grid-searched on calibration data and no backpropagation.
• KV-cache quantization applies the same arithmetic to cached keys and values (KV quantization). QLoRA's NF4 returns in module 36.

Serving detail: GPTQ, AWQ, SmoothQuant and formats and kernels.

torch.nn.utils.prune is a reparametrization (module 10). prune.l1_unstructured(module, "weight", amount=0.5) renames the parameter weight_orig, adds a weight_mask buffer, and registers a forward pre-hook computing weight = weight_orig * weight_mask (source). Gradients reach weight_orig through the mask, so pruned entries stay dead. prune.remove bakes in the product as a plain Parameter (fence 5). prune.global_unstructured ranks across modules; prune.ln_structured(module, "weight", amount, n=2, dim=0) zeros whole output rows.

Zeros are not speed. Dense kernels multiply zeros at full price, and unstructured sparse formats rarely beat dense GPU GEMMs below very high sparsity. ln_structured also only zeros rows: FLOPs fall only after slicing the weight, and the next layer's input columns, to a smaller dense shape.

2:4 semi-structured sparsity (Mishra et al., 2021) zeros exactly 2 of every 4 contiguous values along the reduction dimension. Kernels store the 2 values plus 2-bit indices, and Ampere+ sparse tensor cores offer 2× theoretical math throughput (vendor-reported). The semi-structured tutorial reports 1.3× on a 2:4 BERT on an A100. On CPU, torch.sparse.to_sparse_semi_structured (source) raised RuntimeError ... Only CUDA tensors are currently supported, so the lab simulates the mask. References: pruning tutorial, torchao sparsity README.

A student matches a teacher's softened outputs and the labels (Hinton et al., 2015):

\(T > 1\) exposes how the teacher ranks wrong classes. The KL gradient per student logit is \(\frac{1}{T}(p_s - p_t)\), and for large \(T\) the probability gap itself shrinks as \(1/T\), so soft gradients scale as \(1/T^2\); the \(T^2\) factor keeps their weight against CE stable as \(T\) is tuned. F.kl_div takes the student's log-probabilities as input (module 13):

kd = F.kl_div(F.log_softmax(z_s / T, -1), F.log_softmax(z_t / T, -1),
              log_target=True, reduction="batchmean") * T * T
loss = alpha * kd + (1 - alpha) * F.cross_entropy(z_s, y)

"batchmean" is the true KL; "mean" also divides by \(V\). The distillation tutorial uses the same \(T^2\). The standard edge recipe: distill, then quantize, with QAT if PTQ loses too much.

Fence 1 checks the hand-worked table against torch.quantize_per_tensor: same codes, exact zero, error at most \(s/2\). It proves the formula matches PyTorch's quint8 path, not that min/max is good calibration.

import torch

x = torch.tensor([-1.0, 0.0, 0.5, 2.0])
lo, hi = min(x.min().item(), 0.0), max(x.max().item(), 0.0)
s = (hi - lo) / 255                      # 3/255 = 0.011765
z = int(round(-lo / s))                  # 85: the code that means 0.0
manual = torch.clamp(torch.round(x / s) + z, 0, 255)
assert manual.tolist() == [0.0, 85.0, 127.0, 255.0]   # 0.5/s = 42.5 rounds half-to-even
q = torch.quantize_per_tensor(x, s, z, torch.quint8)
assert torch.equal(q.int_repr().to(torch.float32), manual)
assert q.q_scale() == s and q.q_zero_point() == z
x_hat = q.dequantize()
assert x_hat[1].item() == 0.0                         # zero is exactly representable
assert ((x_hat - x).abs() <= s / 2 + 1e-7).all()
print(f"s={s:.6f} z={z} q={q.int_repr().tolist()} x_hat={[round(v, 4) for v in x_hat.tolist()]}")

Fence 2 puts a ×50 outlier in row 3 of an (8, 64) weight. That row's error is identical under both schemes, so the assertion compares the seven ordinary rows (per-channel MSE over 10× lower) and checks torch.quantize_per_channel within \(s/2\) per row. One synthetic matrix, not a trained model.

import torch

torch.manual_seed(0)
w = torch.randn(8, 64) * 0.05
w[3] *= 50                                            # one outlier output channel

def fq_sym(w, s):
    return torch.clamp(torch.round(w / s), -127, 127) * s

s_t = w.abs().max() / 127                             # one scale, shape ()
s_c = w.abs().amax(dim=1, keepdim=True) / 127         # one scale per row, (8, 1)
mse_t = ((fq_sym(w, s_t) - w) ** 2).mean(dim=1)       # per-row MSE, (8,)
mse_c = ((fq_sym(w, s_c) - w) ** 2).mean(dim=1)
normal = torch.arange(8) != 3
torch.testing.assert_close(mse_c[3], mse_t[3])        # the outlier row: same scale either way
assert mse_c[normal].mean() < mse_t[normal].mean() / 10   # the other rows pay for it

q = torch.quantize_per_channel(w, s_c.squeeze(1).double(), torch.zeros(8, dtype=torch.int64),
                               axis=0, dtype=torch.qint8)
assert q.qscheme() == torch.per_channel_affine            # zero points all 0
torch.testing.assert_close(q.q_per_channel_scales().float(), s_c.squeeze(1), rtol=1e-6, atol=0)
err = (q.dequantize() - w).abs()
assert (err <= s_c / 2 + 1e-6).all()
print(f"ordinary-row MSE per-tensor {mse_t[normal].mean():.2e} vs per-channel {mse_c[normal].mean():.2e}")

Fence 3 round-trips 128 int4 codes through 64 bytes and pins the nibble order. It proves the layout is lossless, not that a kernel reads it.

import torch

def pack_int4(q):                       # q: int8 in [-8, 7], even numel
    u = (q.to(torch.int16) & 0xF).to(torch.uint8).reshape(-1, 2)
    return u[:, 0] | (u[:, 1] << 4)     # low nibble = even index

def unpack_int4(p):
    u = torch.stack([p & 0xF, p >> 4], dim=1).reshape(-1).to(torch.int16)
    return torch.where(u >= 8, u - 16, u).to(torch.int8)   # sign-extend 4 bits

g = torch.Generator().manual_seed(0)
q = torch.randint(-8, 8, (128,), generator=g, dtype=torch.int8)
p = pack_int4(q)
assert p.dtype == torch.uint8 and p.numel() == q.numel() // 2
assert torch.equal(unpack_int4(p), q)
assert pack_int4(torch.tensor([-1, 7], dtype=torch.int8)).item() == 0x7F   # 0xF | 0x7 << 4
print(f"{q.numel()} int4 codes -> {p.numel()} bytes; first byte 0x{p[0].item():02X} from {q[:2].tolist()}")

Fence 4 feeds values inside and outside \([-2.0, 1.75]\) to three fake-quantizers. Forwards agree; the built-in and the autograd.Function give gradient 1 inside and 0 outside; the bare detach idiom gives 1 everywhere; plain round gives 0. It proves gradient semantics, not that QAT converges.

import torch

s, z, qmin, qmax = 0.25, 0, -8, 7                     # representable range [-2.0, 1.75]
x = torch.tensor([-3.0, -1.1, 0.3, 1.6, 2.5], requires_grad=True)

def fq(t):
    return (torch.clamp(torch.round(t / s) + z, qmin, qmax) - z) * s

class FakeQuantSTE(torch.autograd.Function):
    @staticmethod
    def forward(ctx, t):
        ctx.save_for_backward(t)
        return fq(t)

    @staticmethod
    def backward(ctx, g):
        (t,) = ctx.saved_tensors
        inside = (t / s + z >= qmin) & (t / s + z <= qmax)
        return g * inside                                 # STE, masked by the clamp

y_trick = x + (fq(x) - x).detach()                        # unmasked STE
y_fn = FakeQuantSTE.apply(x)
y_builtin = torch.fake_quantize_per_tensor_affine(x, s, z, qmin, qmax)
torch.testing.assert_close(y_trick, y_builtin)
torch.testing.assert_close(y_fn, y_builtin)

g_trick, = torch.autograd.grad(y_trick.sum(), x)
g_fn, = torch.autograd.grad(y_fn.sum(), x)
g_builtin, = torch.autograd.grad(y_builtin.sum(), x)
g_round, = torch.autograd.grad(fq(x).sum(), x)
assert g_builtin.tolist() == [0.0, 1.0, 1.0, 1.0, 0.0]
assert torch.equal(g_fn, g_builtin) and g_trick.tolist() == [1.0] * 5
assert g_round.abs().sum().item() == 0.0                  # plain round: no signal at all
print(f"forward {y_builtin.tolist()}; grad builtin {g_builtin.tolist()}, trick {g_trick.tolist()}")

Fence 5 checks pruning's bookkeeping before and after prune.remove. It proves the reparametrization contract, not a speedup.

import torch
import torch.nn as nn
from torch.nn.utils import prune

torch.manual_seed(0)
lin = nn.Linear(10, 10)
prune.l1_unstructured(lin, "weight", amount=0.5)
assert int((lin.weight == 0).sum()) == 50
assert "weight_mask" in dict(lin.named_buffers())
assert "weight_orig" in dict(lin.named_parameters())
assert "weight" not in dict(lin.named_parameters())       # now a plain attribute
torch.testing.assert_close(lin.weight, lin.weight_orig * lin.weight_mask)
assert len(lin._forward_pre_hooks) == 1                   # recomputes weight each forward

prune.remove(lin, "weight")
assert isinstance(lin.weight, nn.Parameter)
assert "weight_mask" not in dict(lin.named_buffers())
assert "weight_orig" not in dict(lin.named_parameters())
assert int((lin.weight == 0).sum()) == 50 and len(lin._forward_pre_hooks) == 0
print(f"pruned zeros {int((lin.weight == 0).sum())}/100; params {sorted(n for n, _ in lin.named_parameters())}")

./code/29_quantization.py trains a 64 → 128 → 128 → 10 ReLU MLP (30 Adam epochs, seeded 10-class Gaussian mixture, 2,048 train and 1,024 test points), then compresses it from scratch: int8 weight-only per-tensor and per-channel (Int8WeightOnlyLinear), int4 group-32 nibble-packed (Int4GroupLinear), static int8 PTQ (MinMaxObserver on 8 batches feeding StaticInt8Linear's int32 GEMM), QAT (5 epochs of masked-STE fake quantization, then PTQ conversion), and a 2:4 mask with 5 epochs of masked fine-tuning. --with-ao-dynamic adds quantize_dynamic with an engine from supported_engines, restored afterwards, or records a skip. Bytes count codes, scales, zero points, row sums and fp32 biases.

python content/courses/pytorch/code/29_quantization.py
python site/test_pytorch_29_quantization.py

Tested with Python 3.11 and PyTorch 2.8.0 on CPU. One local PyTorch 2.8 CPU run on an Apple M2 Pro laptop (4 threads, seed 0, --with-ao-dynamic) produced:

Per-channel scales halve logit error for 1,052 extra bytes; int4 costs 180× more error but almost no accuracy on this easy task; quantized activations make PTQ's error 8× weight-only's. QAT and 2:4 fine-tune, so their logits legitimately move: judge them by accuracy. Across seeds 0–5 every contract held, with per-channel 1.5–2.8× better than per-tensor. Lab work took about 0.9 s, excluding import torch.

The test asserts contracts, not these values: byte ratios, per-channel ≤ per-tensor error, bounded accuracy drops, QAT ≥ PTQ − 0.01, exact 2:4 groups, a lossless int4 round trip, and a restored engine for --with-ao-dynamic.

Honest scope. Algorithms and bytes are real; speed is not measured (Python dequantize-then-matmul is slower than fp32 on CPU), and torchao kernels and GPU 2:4 are fragments above.