Memory arithmetic and quantizing a model yourself

Whether a model fits on a given GPU is arithmetic you can do in a minute, and you will be asked to. This section is that arithmetic: parameters times bytes per parameter for the weights, and then the KV cache, which is charged per token and per sequence and is usually what runs out first when serving several users at long contexts. The second half is practical: quantizing a matrix by hand so you can see exactly what rounding costs, running that on a trained model, and reading the result, since the damage lands on particular abilities rather than showing up in casual chat.

Weights#

8 billion parameters
  16-bit (2 bytes each)   16 GB
  8-bit  (1 byte)          8 GB
  4-bit  (half a byte)     4 GB   plus a little for scales and some layers kept at higher precision

Parameters times bytes per parameter, and nothing else. The only subtlety is that a 4-bit checkpoint is never exactly half a byte per weight: every group of weights carries a scale, usually at 16 bits, and most methods keep a few sensitive tensors at higher precision. Budget five to ten per cent above the arithmetic and you will be close.

The KV cache, which people forget#

per token = 2 (key and value) x layers x kv_heads x head_dim x bytes

Example: 32 layers, 8 KV heads, head_dim 128, 16-bit
  2 x 32 x 8 x 128 x 2 bytes = 131,072 bytes = 128 KiB per token
  8,000-token context        = about 1 GiB per sequence
  batch of 16 at 8k tokens   = about 16 GiB, more than the 4-bit weights

For serving, the cache often decides how many requests fit, not the weights. That is why techniques that share key and value heads, and quantizing the cache itself, matter so much. Note which numbers appear and which do not: the number of query heads is irrelevant, because grouped-query attention shares one key and value head across several query heads, and only the shared ones are stored.

Does it fit, as a function rather than an argument#

This is the question an interviewer asks and the question you will answer weekly on the job, so write it down once. The example below is an 8B model on a 24 GB card, serving sixteen users at an 8k context, and it runs through the options in the order you would actually try them.

def fits(params_b, layers, kv_heads, head_dim, weight_bits, kv_bits,
         batch, ctx, gpu_gib):
    weights = params_b * 1e9 * weight_bits / 8
    per_token = 2 * layers * kv_heads * head_dim * kv_bits / 8
    cache = per_token * batch * ctx
    total = (weights + cache) / 2**30
    print("%2d-bit weights, %2d-bit cache, batch %2d at %5d ctx: "
          "weights %5.1f + cache %5.1f = %5.1f GiB  %s"
          % (weight_bits, kv_bits, batch, ctx, weights / 2**30,
             cache / 2**30, total, "fits" if total < gpu_gib else "does not fit"))

for wb, kb, batch in ((16, 16, 16), (8, 16, 16), (4, 16, 16), (4, 8, 16), (4, 8, 32)):
    fits(8, 32, 8, 128, wb, kb, batch, 8192, gpu_gib=24)

# 16-bit weights, 16-bit cache, batch 16 at  8192 ctx: weights  14.9 + cache  16.0 =  30.9 GiB  does not fit
#  8-bit weights, 16-bit cache, batch 16 at  8192 ctx: weights   7.5 + cache  16.0 =  23.5 GiB  fits
#  4-bit weights, 16-bit cache, batch 16 at  8192 ctx: weights   3.7 + cache  16.0 =  19.7 GiB  fits
#  4-bit weights,  8-bit cache, batch 16 at  8192 ctx: weights   3.7 + cache   8.0 =  11.7 GiB  fits
#  4-bit weights,  8-bit cache, batch 32 at  8192 ctx: weights   3.7 + cache  16.0 =  19.7 GiB  fits

The last two rows are the ones worth remembering. Halving the cache precision buys more headroom here than halving the weight precision did, because at batch 16 and 8k tokens the cache is larger than the weights. And the room you free goes straight into concurrency: the same 19.7 GiB serves sixteen users with a 16-bit cache or thirty-two with an 8-bit one.

Quantize a matrix by hand#

Quantization in its simplest form is three operations: divide by a scale, round to an integer, clamp to the range the integer can hold. Everything the methods argue about is how the scale is chosen and how many weights share one. Start with a weight matrix shaped like a real one, which means mostly small values and a couple of channels with a much wider range.

import torch

torch.manual_seed(0)

W = torch.randn(512, 512) * 0.02
W[7] *= 25.0                                     # two loud output channels
W[200] *= 12.0

def quantize(W, bits, per_channel):
    """Symmetric integer quantization, returned dequantized so the error is visible."""
    qmax = 2 ** (bits - 1) - 1                   # 127 at int8, 7 at int4
    amax = W.abs().max(dim=1, keepdim=True).values if per_channel else W.abs().max()
    scale = amax / qmax
    q = torch.clamp(torch.round(W / scale), -qmax - 1, qmax)
    return q * scale

def rel_err(W, Wq):
    return float((Wq - W).pow(2).mean().sqrt() / W.pow(2).mean().sqrt())

print("abs max %.4f   std %.4f   ratio %.0f"
      % (W.abs().max(), W.std(), W.abs().max() / W.std()))
Wq = quantize(W, 8, per_channel=False)
print("int8 per-tensor: scale %.3e  rel RMS err %.4f"
      % (W.abs().max() / 127, rel_err(W, Wq)))

# abs max 1.6916   std 0.0313   ratio 54
# int8 per-tensor: scale 1.332e-02  rel RMS err 0.1229

That is the whole problem in one number. The scale has to be large enough that the biggest weight does not clip, so it is set by the single loudest value in the tensor. Every ordinary weight is then rounded to a grid whose spacing was chosen for somebody else.

Where the error actually lands#

row_err = (Wq - W).pow(2).mean(dim=1).sqrt() / W.pow(2).mean(dim=1).sqrt()
loud = W.abs() > 0.2
print("weights above 0.2: %d of %d (%.2f%%)"
      % (int(loud.sum()), W.numel(), 100 * float(loud.float().mean())))
print("rel err: median row %.4f   row 0 %.4f   row 7 (the loud one) %.4f"
      % (row_err.median(), row_err[0], row_err[7]))

# weights above 0.2: 559 of 262144 (0.21%)
# rel err: median row 0.1922   row 0 0.1928   row 7 (the loud one) 0.0079

This is the result worth carrying away, because it is the opposite of what most people assume. Outliers do not get damaged by quantization. They get served perfectly, at the expense of the 99.79 per cent of weights that are nothing like them. Once you see it that way, every method in the field reads as a different answer to the same question: how do we stop a handful of values from setting the grid for everyone else.

Per-channel, and what it buys#

The cheapest answer is to stop sharing. Give each output channel its own scale, so a loud row sets the grid only for itself. The code above already supports it; the only change is which axis the maximum is taken over.

print("bits  scheme       rel RMS err   max abs err")
for bits in (8, 4, 3):
    for name, pc in (("per-tensor", False), ("per-channel", True)):
        Q = quantize(W, bits, per_channel=pc)
        print("%4d  %-11s  %11.4f   %11.4f"
              % (bits, name, rel_err(W, Q), (Q - W).abs().max()))

# bits  scheme       rel RMS err   max abs err
#    8  per-tensor        0.1229        0.0067
#    8  per-channel       0.0076        0.0067
#    4  per-tensor        0.6537        0.1207
#    4  per-channel       0.1383        0.1206
#    3  per-tensor        0.7146        0.2818
#    3  per-channel       0.3243        0.2817

Two things in that table repay a second look. The max absolute error column barely moves, because it is still dominated by the loud channel in both schemes; only the average improves, which is why RMS error is the number to watch rather than the worst case. And the gap widens as bits fall: at eight bits per-channel is a refinement, at three bits it is the difference between a usable tensor and a broken one. Group-wise quantization, which every serious 4-bit method uses, is the same idea taken further: one scale per 64 or 128 weights within a row rather than one per row.

n = W.numel()
print("bits  scales  bytes    vs 16-bit")
for bits, scales in ((16, 0), (8, 0), (8, 512), (4, 0), (4, 512)):
    total = n * bits / 8 + scales * 2             # scales stored at 16-bit
    print("%4d  %6d  %7d  %8.2fx" % (bits, scales, total, n * 2 / total))

# bits  scales  bytes    vs 16-bit
#   16       0   524288      1.00x
#    8       0   262144      2.00x
#    8     512   263168      1.99x
#    4       0   131072      4.00x
#    4     512   132096      3.97x

Now do it to a trained model#

Error on a matrix is a proxy. What you actually care about is whether the model still works, and the only way to know is to quantize the weights and rerun the eval. A three-layer classifier on the digits dataset trains in a few seconds on a CPU and is enough to show the shape of the answer.

import copy
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

X, y = load_digits(return_X_y=True)
Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=0.3, random_state=0, stratify=y)
Xtr = torch.tensor(Xtr, dtype=torch.float32) / 16.0
Xte = torch.tensor(Xte, dtype=torch.float32) / 16.0
ytr, yte = torch.tensor(ytr), torch.tensor(yte)

torch.manual_seed(0)
model = torch.nn.Sequential(
    torch.nn.Linear(64, 256), torch.nn.ReLU(),
    torch.nn.Linear(256, 256), torch.nn.ReLU(),
    torch.nn.Linear(256, 10),
)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
for _ in range(200):
    opt.zero_grad()
    torch.nn.functional.cross_entropy(model(Xtr), ytr).backward()
    opt.step()

def accuracy(m):
    with torch.no_grad():
        return float((m(Xte).argmax(1) == yte).float().mean())

def quantized_copy(m, bits, per_channel):
    c = copy.deepcopy(m)
    with torch.no_grad():
        for layer in c:
            if isinstance(layer, torch.nn.Linear):
                layer.weight.copy_(quantize(layer.weight, bits, per_channel))
    return c

with torch.no_grad():
    spread = [float(l.weight.abs().max(1).values.max() / l.weight.abs().max(1).values.median())
              for l in model if isinstance(l, torch.nn.Linear)]
print("channel range spread per layer:", [round(s, 2) for s in spread])
print("fp32 accuracy %.4f" % accuracy(model))
print("bits  per-tensor  per-channel")
for bits in (8, 4, 3, 2):
    print("%4d  %10.4f  %11.4f"
          % (bits, accuracy(quantized_copy(model, bits, False)),
             accuracy(quantized_copy(model, bits, True))))

# channel range spread per layer: [1.59, 1.47, 1.19]
# fp32 accuracy 0.9759
# bits  per-tensor  per-channel
#    8      0.9759       0.9759
#    4      0.9685       0.9722
#    3      0.9667       0.9704
#    2      0.1944       0.4741

The first line explains why per-channel helps so little here: the loudest channel in each layer is only 1.2 to 1.6 times the median, against a ratio of 54 in the synthetic matrix. A small well-conditioned network has no outlier problem to fix. Large language models do, and the gap between the two experiments is exactly the reason quantizing a transformer is a research area and quantizing this MLP is a one-liner.

The other thing to take from the table is the shape of the curve. Quality does not degrade smoothly as bits fall; it is flat, then it slopes, then it falls off a cliff between three and two bits. That cliff is why the field settled on four bits as the working floor for weights, and why claims about two-bit models always come with a method attached rather than a rounding rule.

Common approaches#

• Weight-only quantization stores weights in 8 or 4 bits and computes in higher precision. It is the most common for serving LLMs, because generation is limited by reading weights from memory.
• GPTQ and AWQ are widely used methods for 4-bit weights that use a small calibration set to decide how to round with less damage. Both are doing a smarter version of what you did above: GPTQ rounds one column at a time and corrects the remaining columns for the error it just made, AWQ finds the channels the calibration activations care about and scales them so they survive.
• GGUF is the file format used by llama.cpp for running quantized models on laptops and CPUs. Its Q4_K_M and similar names encode a bit width plus a grouping scheme, not a separate algorithm.
• bitsandbytes loads Hugging Face models in 8 or 4 bits with one argument, convenient for experiments and for QLoRA fine-tuning.

from transformers import AutoModelForCausalLM, BitsAndBytesConfig

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=BitsAndBytesConfig(load_in_4bit=True),
    device_map="auto",
)

Activation quantization is a separate and harder problem, and worth knowing the name of. Weights are fixed, so you can look at them all and pick a scale. Activations depend on the input, so the scale must either be calibrated on sample data and hoped to hold, or computed at runtime, and transformer activations carry outliers far more extreme than the weights do. That is why the common production answer is weight-only, and why the papers about activation outliers exist at all.

Measure what you lost#

Perplexity on held-out text is a quick check: a large rise means damage. It is not enough on its own, because it averages over everything and the damage is not evenly spread. Run your task eval on both versions, and look especially at long inputs, arithmetic, code and less common languages, which degrade first. The pattern from the matrix experiment repeats at the model level: the average looks fine while a specific capability has quietly gone.