Quantisation and acceleration

Two ways into the same idea

Before You Start

The Core Idea

Quantisation stores weights at lower precision. Here you start with bf16, quantise a matrix by hand, measure the effect on a real task, then load a quantised language model.

1Lower precision: bf16

Each weight in fp32 takes 4 bytes. bf16 takes 2 with almost no loss, and it is today's default for any model on a modern card.

Look at the picture: a 1024 by 1024 matrix at each precision. Each step down halves the memory.

2Quantisation: a scale and rounding

Each row of weights has a largest absolute value. We divide the row by a scale that makes that value 127 for 8 bits, and round to the nearest integer.

Look at the picture: at use time we multiply by the scale again. What is lost is the rounding fractions.

import numpy as np

rng = np.random.default_rng(0)
W = rng.normal(scale=0.05, size=(1024, 1024)).astype(np.float32)

def quantize(w, bits):
    levels = 2 ** (bits - 1) - 1                     # 127 for int8, 7 for int4
    scale = np.abs(w).max(axis=1, keepdims=True) / levels   # one scale per row
    q = np.round(w / scale).clip(-levels, levels)
    return q, scale

for bits in [8, 4]:
    q, scale = quantize(W, bits)
    restored = q * scale
    err = np.abs(W - restored).mean() / np.abs(W).mean()
    size = W.size * bits / 8 / 1e6
    print(f"int{bits}: {size:.2f} MB instead of {W.nbytes / 1e6:.2f} MB, mean relative error {err:.1%}")

8 bits: a quarter of the size with under 1% error. 4 bits: an eighth of the size with 15% weight error. But weight error is not what matters; the output error on your task does, and it can be much smaller because random errors partly cancel. That is what we measure in the next example.

3Measure on your task

Weight error is not what matters. The output error on your task does, and we measure it with two numbers: accuracy, and the share of answers matching the original model.

Look at the picture: 8 bits match everything, 4 bits nearly, and 2 bits collapse.

A trained digit classifier; we quantise its weights to 8, 4 and 2 bits, and measure accuracy and the share of answers matching the original.

import numpy as np
from sklearn.datasets import load_digits
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split

X, y = load_digits(return_X_y=True)
X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.3, random_state=0)
model = LogisticRegression(max_iter=2000).fit(X_tr / 16, y_tr)
W, b = model.coef_.copy(), model.intercept_

def predict(weights):
    return np.argmax(X_te / 16 @ weights.T + b, axis=1)

def quantize(w, bits):
    levels = 2 ** (bits - 1) - 1
    scale = np.abs(w).max(axis=1, keepdims=True) / levels
    return np.round(w / scale).clip(-levels, levels) * scale

base = predict(W)
print(f"float32: accuracy {np.mean(base == y_te):.3f}")
for bits in [8, 4, 2]:
    p = predict(quantize(W, bits))
    print(f"int{bits}:    accuracy {np.mean(p == y_te):.3f}   same answer as float32 on {np.mean(p == base):.1%} of inputs")

8 bits: the same accuracy and the same answer on every input. 4 bits: a 0.7-point loss and a different answer on about 2% of inputs. 2 bits: collapse. Notice the weight error at 4 bits was 15%, yet the task error is under a point. Notice the agreement too: even when accuracy looks similar, the differing 2% may be important cases, so read them.

Losses show first on hard tasks: arithmetic, long reasoning and less-represented languages. Simple tasks suffer less.

4A quantised language model in practice

With language models you do not write quantisation by hand. The bitsandbytes library loads the model in 4 bits directly, and 7 billion weights drop from about 15 GB to about 5.

from transformers import AutoModelForCausalLM, BitsAndBytesConfig

quant = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                           bnb_4bit_compute_dtype="bfloat16")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-7B-Instruct", quantization_config=quant)
print(model.get_memory_footprint() / 1e9, "GB")   # about 5 GB instead of about 15 GB in bf16

nf4 is a 4-bit type designed for neural network weight distributions, and compute_dtype makes the compute itself bf16 after decompression. Run the frozen set on this model and on the bf16 version, and put both numbers in a table before you decide.

How It Works, Step by Step: how a row of weights is quantised

Worked example

The idea: each row of weights has a largest absolute value. We divide the row by a scale that makes that value 127 (for 8 bits), and round to the nearest integer. At use time we multiply by the scale again. What is lost is the rounding fractions.

Try it without code

Myth vs Reality

Try it yourself

Quick review cards

Check Your Understanding