Module 14: Quantization

Introduction

You have trained a model (Modules 06–08). Every weight in it is a 32-bit float — 4 bytes each. A 7-billion-parameter model is therefore 28 GB just to store, and every matmul reads all of it from memory. Quantization is the single most effective way to shrink that: round each weight onto a small grid of integers and keep one floating-point scale to undo the rounding at run time.

Store 8-bit integers instead of 32-bit floats and the model is 4× smaller; store 4-bit integers and it is 8× smaller — often the difference between a model fitting on your GPU or not. The cost is a small, measurable rounding error, and this module builds the whole machinery from scratch so you can see exactly what that error is and how to keep it tiny.

What You’ll Learn

  • Why a trained weight has far more precision than inference needs
  • Symmetric (absmax) and affine (zero-point) quantization, from the formulas
  • How to quantize a tensor to int8 / int4 and measure the error in SNR
  • Why per-channel scales beat a single per-tensor scale (the outlier problem)
  • How to build a QuantizedLinear that replaces nn.Linear at a fraction of the memory

Prerequisites

Intuition: A Weight Is a Number With Too Many Digits

A trained weight might be 0.0417328.... During inference, does the model really need all those digits? Almost never. The weights in a given tensor span some range — say roughly [-0.2, 0.2] — and if you chop that range into a few hundred evenly spaced buckets, rounding each weight to the nearest bucket barely changes the model’s output.

That is the whole idea. Pick a grid of integer levels, store which bucket each weight fell into (a small integer), and remember one scale — the width of a bucket — to turn buckets back into real numbers:

\[w \approx \text{scale} \times q, \qquad q \in \{-127, \dots, 127\}\ \text{(int8)}\]

Eight bits gives \(2^8 = 256\) buckets; four bits gives just \(16\). Fewer buckets = a coarser grid = more rounding error. The art of quantization is choosing the grid — how wide, how centred, and how many independent grids to keep — so that the error stays negligible while the integers stay tiny.

NoteKey Insight

Quantization trades precision for size. A float carries ~7 decimal digits of precision that a trained, redundant network mostly does not use. Rounding weights onto a 256-level (int8) grid typically changes outputs by well under a percent, while cutting the model’s memory — and its memory-bandwidth cost, which usually dominates inference — by 4×.

The Math: Symmetric and Affine Quantization

Symmetric (absmax). The simplest scheme centres the grid on zero. For \(b\) bits we use the signed range \([-q_{\max}, q_{\max}]\) with \(q_{\max} = 2^{b-1}-1\) (so \(127\) for int8). One scale covers the whole tensor:

\[s = \frac{\max_i |w_i|}{q_{\max}}, \qquad q_i = \text{clamp}\!\left(\text{round}\!\left(\frac{w_i}{s}\right),\, -q_{\max},\, q_{\max}\right), \qquad \hat{w}_i = s\, q_i\]

Because the grid is symmetric, \(0.0\) maps exactly to the integer \(0\) — no offset needed. This is the default for weights, which are roughly zero-centred.

Affine (zero-point). When the values are one-sided — think the output of a ReLU, all \(\geq 0\) — a symmetric grid wastes half its levels on negatives that never occur. Affine quantization maps the true range \([w_{\min}, w_{\max}]\) onto the unsigned range \([0, 2^b - 1]\) using an integer zero-point \(z\) (the code that represents \(0.0\)):

\[s = \frac{w_{\max} - w_{\min}}{2^b - 1}, \qquad z = \text{round}\!\left(-\frac{w_{\min}}{s}\right), \qquad q_i = \text{clamp}\!\left(\text{round}\!\left(\tfrac{w_i}{s}\right) + z,\, 0,\, 2^b - 1\right)\]

and we recover \(\hat{w}_i = s\,(q_i - z)\). No levels are wasted, so affine is the better choice for activations.

NoteKey Insight

Symmetric needs only a scale; affine needs a scale and a zero-point but uses the full code range. Weights → symmetric (zero-centred, and a zero-point would slow the matmul). Activations → affine (often one-sided). Both are just a linear map between reals and integers; the only lossy step is round.

Code: Quantize a Tensor from Scratch

Everything lives in quantization.py. Load it and quantize a weight tensor to int8, then read back the reconstruction error.

import importlib.util
import sys
from pathlib import Path

import torch
import torch.nn as nn

spec = importlib.util.spec_from_file_location("quantization", Path("quantization.py").resolve())
q = importlib.util.module_from_spec(spec)
sys.modules["quantization"] = q
spec.loader.exec_module(q)

torch.manual_seed(0)
w = torch.randn(4, 4) * 0.1                       # a small weight matrix
qi, scale = q.absmax_quantize(w, num_bits=8)

print("original[0]:   ", w[0].tolist())
print("integers[0]:   ", qi[0].tolist(), f"  (dtype {qi.dtype})")
print(f"scale:          {scale.item():.6f}")
print("dequantized[0]:", q.dequantize(qi, scale)[0].round(decimals=4).tolist())
original[0]:    [-0.11258398741483688, -0.11523602157831192, -0.025057857856154442, -0.04338788241147995]
integers[0]:    [-68, -69, -15, -26]   (dtype torch.int8)
scale:          0.001666
dequantized[0]: [-0.11330000311136246, -0.11490000039339066, -0.02500000037252903, -0.043299999088048935]

The four floats became four small integers plus one shared scale. Now measure how much precision that cost, in int8 versus int4:

w = torch.randn(512, 512) * 0.05
for bits in (8, 4):
    qi, s = q.absmax_quantize(w, num_bits=bits)
    err = q.quantization_error(w, q.dequantize(qi, s))
    print(f"int{bits}: max error {err['max_abs']:.5f}   SNR {err['snr_db']:5.1f} dB")
int8: max error 0.00092   SNR  39.5 dB
int4: max error 0.01664   SNR  14.3 dB

int8 reconstructs the tensor at ~40 dB (rounding error ~10 000× smaller than the signal); int4, with only 16 levels, drops to the mid-teens. Signal-to-noise ratio is the honest way to compare: every extra bit adds ~6 dB, so int8 is worth about 24 dB more than int4.

Per-Tensor vs Per-Channel: The Outlier Problem

One scale for a whole matrix has a weakness. Neural-network weight matrices routinely contain a few outlier channels whose magnitudes dwarf the rest. The absmax scale is set by that single largest value, so every other channel is forced onto a needlessly coarse grid.

The fix is per-channel quantization: give each output channel (row) its own scale. One extra float per row buys back the precision the outlier stole. Watch it on a matrix where channel 0 is 20× louder than the rest:

torch.manual_seed(0)
w = torch.randn(24, 64) * 0.05
w[0] *= 20.0                                       # one loud output channel

qt, st = q.absmax_quantize(w, num_bits=8)          # one scale for everything
qc, sc = q.quantize_per_channel(w, num_bits=8, axis=0)   # one scale per row

# Per-channel mean-squared error under each scheme (for the visual below).
mse_tensor = ((w - q.dequantize(qt, st)) ** 2).mean(dim=1)
mse_chan = ((w - q.dequantize(qc, sc)) ** 2).mean(dim=1)

ojs_define(quantChan = {
    "perTensor": mse_tensor.tolist(),
    "perChannel": mse_chan.tolist(),
    "outlier": 0,
})
et = q.quantization_error(w, q.dequantize(qt, st))
ec = q.quantization_error(w, q.dequantize(qc, sc))
print(f"per-tensor  : SNR {et['snr_db']:.1f} dB")
print(f"per-channel : SNR {ec['snr_db']:.1f} dB   ← the quiet channels keep their precision")
per-tensor  : SNR 29.1 dB
per-channel : SNR 42.8 dB   ← the quiet channels keep their precision

Per-tensor lets one loud channel drag every quiet channel’s error up by the same factor it is louder — here ~20×. Per-channel absorbs the outlier into its own row and leaves the rest untouched, recovering well over 10 dB of SNR. This is exactly why production weight quantization is always per-channel (and why methods like LLM.int8() go further and pull genuine outlier features out into fp16).

Code: A Quantized Linear Layer

Now assemble the piece that actually saves memory in a transformer: a drop-in for nn.Linear whose weight is stored per-channel int8. QuantizedLinear.from_linear quantizes a trained layer; its forward dequantizes the weight and runs an ordinary linear — the faithful, readable stand-in for what an integer kernel does (integer matmul, then a per-channel rescale).

torch.manual_seed(0)
linear = nn.Linear(512, 512)
qlinear = q.QuantizedLinear.from_linear(linear, num_bits=8)

x = torch.randn(8, 512)
ref = linear(x)
out = qlinear(x)

rel = (out - ref).norm() / ref.norm()
print(f"output relative error: {rel.item():.4%}")     # well under 1%

mem = qlinear.memory_bytes()
print(f"fp32 weight bytes: {mem['fp32_weight_bytes']:,}")
print(f"int8 stored bytes: {mem['weight_bytes'] + mem['scale_bytes']:,.0f}"
      f"  ({mem['compression']:.2f}× smaller)")
output relative error: 0.3824%
fp32 weight bytes: 1,048,576
int8 stored bytes: 264,192  (3.97× smaller)

The quantized layer’s output is within a fraction of a percent of the original, at ~4× less memory. Scale that up to a whole model:

for params, name in [(7e9, "7B"), (13e9, "13B"), (70e9, "70B")]:
    line = f"{name:>4}:  "
    for bits in (32, 16, 8, 4):
        line += f"int{bits}={q.model_footprint(int(params), bits)/1e9:6.1f}GB  "
    print(line)
  7B:  int32=  28.0GB  int16=  14.0GB  int8=   7.0GB  int4=   3.5GB  
 13B:  int32=  52.0GB  int16=  26.0GB  int8=  13.0GB  int4=   6.5GB  
 70B:  int32= 280.0GB  int16= 140.0GB  int8=  70.0GB  int4=  35.0GB  

A 70B model is 280 GB in fp32 — many GPUs — but 35 GB in int4, within reach of a single high-memory card. That table is why quantization is not optional at the frontier.

Interactive Exploration

Step Through the Pipeline

Every quantization is the same six steps: measure the range, form a scale, divide, round, store the integers, and dequantize on the way back out. Step through them and watch where the lossy round sits.

Drive the Grid Yourself

Slide the bit-width and switch schemes. Fewer bits coarsen the grid; the arrows show each weight snapping to its nearest level, and the readouts track the error. Below it, the per-channel chart (built from the tested code above) shows the outlier channel’s damage — and how per-channel scales contain it.

TipTry This
  1. Watch int4 break. In the grid explorer, drag bits from 8 down to 2. At int2 there are only 4 levels — the snapped dots barely resemble the weights, and the SNR collapses. This is why naive int4 needs help (grouping, GPTQ).
  2. Re-centre with affine. At low bit-widths, switch from symmetric to affine. The grid slides to hug the data’s actual range instead of straddling zero, and the error drops — the payoff of spending a zero-point.
  3. Read the per-channel chart. The red bars (per-tensor) tower over the green (per-channel) for every quiet channel, because the outlier set their shared scale. Only the outlier channel itself is a near-tie.

Common Pitfalls

  1. Forgetting to clamp. After round(w/s) a value can land at \(\pm(q_{\max}+1)\) from floating-point error. Without clamp it overflows the integer type and corrupts the weight. Always clamp to the representable range.

  2. Symmetric quantization on skewed data. A symmetric grid centred on zero wastes half its levels on a range the data never visits (e.g. post-ReLU activations). Use affine — with a zero-point — whenever the values are one-sided.

  3. Per-tensor scales on weights. One outlier channel forces a coarse grid on the whole matrix. Per-channel (per-row) scales are nearly free — one float per channel — and are the standard for weight quantization.

  4. Expecting int4 to just work. Four bits is 16 levels; plain per-channel int4 usually loses too much. Real 4-bit uses smaller groups (a scale per 64 or 128 weights) and error-correcting methods like GPTQ/AWQ. This module builds the foundation those extend.

  5. Confusing storage bits with compute. Quantizing to int8 shrinks storage and memory bandwidth immediately. Getting the speedup also needs an integer matmul kernel; the QuantizedLinear here dequantizes-then-matmuls to stay readable, which is correct but not itself faster.

Exercises

Exercise 1: int4 vs int8 error

Quantize the same tensor to int8 and int4 and confirm int4’s SNR is roughly 24 dB lower (about 6 dB per bit).

x = torch.randn(10_000) * 0.1
q8, s8 = q.absmax_quantize(x, num_bits=8)
q4, s4 = q.absmax_quantize(x, num_bits=4)
snr8 = q.quantization_error(x, q.dequantize(q8, s8))["snr_db"]
snr4 = q.quantization_error(x, q.dequantize(q4, s4))["snr_db"]
print(f"int8 {snr8:.1f} dB,  int4 {snr4:.1f} dB,  gap {snr8 - snr4:.1f} dB")
int8 39.6 dB,  int4 14.5 dB,  gap 25.2 dB

Exercise 2: Build the per-channel win

Take a matrix, make one channel an outlier, and show per-channel MSE is far below per-tensor MSE. (Hint: quantize_per_channel(..., axis=0) vs absmax_quantize.)

# Your implementation here

Exercise 3: Footprint of a model that fits

What is the largest model (in billions of parameters) that fits in 24 GB of GPU memory at int4, leaving 4 GB for activations? Use model_footprint.

# budget = 20 GB for weights; solve params from model_footprint(params, 4)
budget_bytes = 20e9
params = budget_bytes / q.bytes_per_param(4)
print(f"~{params/1e9:.0f}B parameters fit in int4 within a 20 GB weight budget")
~40B parameters fit in int4 within a 20 GB weight budget

Summary

Key takeaways:

  1. Quantization stores integers, not floats. Round each weight onto a grid of \(2^b\) levels and keep a scale to undo it: \(w \approx s\,q\). int8 is 4× smaller than fp32, int4 is 8×.

  2. Two schemes. Symmetric (absmax) centres the grid on zero — best for weights. Affine adds a zero-point to use the full code range — best for one-sided activations.

  3. SNR is the honest metric. Every bit is worth ~6 dB; int8 reconstructs weights near 40 dB, int4 near the mid-teens.

  4. Per-channel beats per-tensor. A single outlier channel wrecks a per-tensor scale; one scale per row costs almost nothing and recovers the lost precision.

  5. QuantizedLinear is the payoff. Stored per-channel int8, it matches the fp32 layer to a fraction of a percent at ~4× less memory — the building block of every quantized transformer.

What’s Next

Quantization is the first lever of fast inference. From here the frontier adds speculative decoding (a small draft model proposes, the big model verifies), paged attention (vLLM’s KV-cache memory manager), and lower-bit methods (GPTQ, AWQ) that push past int4 with grouping and error correction — each building on the integer round-trip you just implemented.

Going Deeper

Core Papers:

  • Jacob et al., Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference (2017) — the affine/zero-point scheme. https://arxiv.org/abs/1712.05877
  • Dettmers et al., LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale (2022) — outlier features and why plain int8 breaks at scale. https://arxiv.org/abs/2208.07339
  • Frantar et al., GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers (2022) — accurate one-shot 4-bit weight quantization. https://arxiv.org/abs/2210.17323

Practical Resources: