Module 4 · Reading 4.4 · read before S26: Efficiency & Scale

The PDF posted in Canvas is the assigned version of this reading. It is built from this page, so the text is identical; anything here marked as deep reading is optional and is not in the PDF.

For most of this course, efficiency has been a footnote: our models fit in memory, train in seconds, and run on a laptop CPU. At modern scale the footnote becomes the headline. A 7-billion-parameter model (small by today’s standards) cannot even be loaded for standard training on any single consumer GPU, and the electricity and hardware budget of a large pretraining run is measured in millions of dollars. When scaling laws say that capability tracks compute, efficiency techniques are not housekeeping; they are the exchange rate between your budget and your model’s quality.

This reading covers the ideas you will meet everywhere: mixed-precision arithmetic, memory-partitioned distributed training (ZeRO/FSDP), the memory economics of serving (KV cache and batching), and the compression trio of quantization, pruning, and distillation. Most of them we can demonstrate with real numbers on a CPU — because at bottom they are about arithmetic and bookkeeping, not about owning a data center.

How to read this

Number formats: what half precision actually buys and costs

Every tensor you have created so far defaulted to float32: 1 sign bit, 8 exponent bits (range), 23 mantissa bits (precision). Halving storage to 16 bits doubles the parameters you can fit per gigabyte and, on GPU tensor cores, can more than double arithmetic throughput. But 16 bits must be split between range and precision, and the two standard splits make opposite choices. float16 (fp16) keeps 10 mantissa bits but only 5 exponent bits; bfloat16 (bf16), introduced for Google’s TPUs, keeps fp32’s full 8 exponent bits and only 7 mantissa bits. The consequences are not abstract:

import torch

torch.manual_seed(0)

for dt in [torch.float32, torch.float16, torch.bfloat16]:
    fi = torch.finfo(dt)
    print(f"{str(dt):15}  max {fi.max:.3e}  smallest normal {fi.tiny:.3e}  eps {fi.eps:.3e}")

print()
# fp16 overflows where bf16 does not (but bf16 rounds coarsely)
x = torch.tensor(70000.0)
print(f"70000 in fp16: {x.to(torch.float16).item()}")     # overflow -> inf
print(f"70000 in bf16: {x.to(torch.bfloat16).item()}")    # fits, but rounded

# precision: 1 + 2^-10 is representable in fp16, lost in bf16
y = torch.tensor(1.0 + 2**-10)
print(f"1 + 2^-10 in fp32: {y.item():.10f}")
print(f"1 + 2^-10 in fp16: {y.to(torch.float16).item():.10f}")
print(f"1 + 2^-10 in bf16: {y.to(torch.bfloat16).item():.10f}")

print()
# the weight-update hazard: a small update vanishes when added to a large fp16 weight
w32 = torch.tensor(2048.0); w16 = w32.to(torch.float16)
print(f"fp32: 2048 + 1 = {(w32 + 1).item()}")
print(f"fp16: 2048 + 1 = {(w16 + torch.tensor(1.0, dtype=torch.float16)).item()}")

# gradient underflow: tiny fp16 gradients flush to zero -> loss scaling exists for this
g = torch.tensor(1e-8)
print(f"gradient 1e-8 in fp16: {g.to(torch.float16).item()}")
print(f"after 1024x loss scaling: {(g * 1024).to(torch.float16).item():.3e}")

Running this prints:

torch.float32    max 3.403e+38  smallest normal 1.175e-38  eps 1.192e-07
torch.float16    max 6.550e+04  smallest normal 6.104e-05  eps 9.766e-04
torch.bfloat16   max 3.390e+38  smallest normal 1.175e-38  eps 7.812e-03

70000 in fp16: inf
70000 in bf16: 70144.0
1 + 2^-10 in fp32: 1.0009765625
1 + 2^-10 in fp16: 1.0009765625
1 + 2^-10 in bf16: 1.0000000000

fp32: 2048 + 1 = 2049.0
fp16: 2048 + 1 = 2048.0
gradient 1e-8 in fp16: 0.0
after 1024x loss scaling: 1.025e-05

Read the table first. fp16’s maximum value is 65,504 — so the perfectly ordinary number 70,000 becomes inf, and one inf in a loss or gradient poisons everything downstream into NaNs. bf16 stores 70,000 without complaint but as 70,144: with 7 mantissa bits, it carries only about 2–3 decimal digits, and 1 + 2^-10 (which fp16 represents exactly) rounds back to 1.0.

The last four lines show why pure half-precision training is fragile. In fp16, adding 1 to a weight of value 2048 does nothing. A gradient of 1e-8 also rounds to zero: it is below the smallest positive fp16 subnormal (about 5.96e-8), not merely below the smallest normal value 6.10e-5. Values between those thresholds can be represented subnormally, subject to hardware flush-to-zero behavior.

Mixed-precision training (Micikevicius et al., ICLR 2018) is the package of fixes that made 16-bit training standard: keep an fp32 master copy of the weights and apply updates there (fixing 2048 + 1 = 2048); multiply the loss by a scale factor like 1024 before backprop so gradients shift into fp16’s representable range, then unscale before the update (our last line shows the rescued gradient); and run the forward/backward math in 16-bit for speed. PyTorch wraps the whole recipe in torch.autocast plus GradScaler. With bf16, the range problems disappear (no loss scaling needed) at the cost of coarser precision — which is why bf16 has become the default on hardware that supports it.

Historical note. The mixed-precision recipe was published by NVIDIA and Baidu researchers in 2018, in the same period NVIDIA’s Volta GPUs shipped “tensor cores” — units that multiply 16-bit matrices far faster than the general-purpose pipeline. The technique and the hardware were designed for each other, and the pattern has repeated since: bf16 support followed on later GPUs and TPUs, then 8-bit float formats on H100-class hardware. Numeric formats are one of the few places where a software paper and a chip roadmap are two halves of one design.

Where training memory actually goes

Here is a useful napkin layout for one conventional mixed-precision Adam recipe: fp16 weights (\(2N\)), fp16 gradients (\(2N\)), fp32 master weights (\(4N\)), and two fp32 moments (\(8N\)) total \(16N\) bytes. Implementations may keep gradients in fp32, omit master weights for bf16, quantize moments, or shard/offload state, so always state the layout behind the estimate.

Plain data parallelism (the baseline distributed strategy) replicates all of this on every GPU: each GPU takes a slice of the batch, computes gradients, and the GPUs average them (an all-reduce) before stepping. Replication buys throughput but saves no memory. The ZeRO family of optimizations (Rajbhandari et al., 2020, the basis of Microsoft’s DeepSpeed) starts from the observation that the replicas are redundant: every GPU holds an identical copy of state that could be partitioned. ZeRO stage 1 shards the optimizer state across GPUs (each GPU updates only its shard); stage 2 also shards gradients; stage 3 also shards the weights themselves, gathering each layer’s parameters just in time for its forward and backward pass. PyTorch’s built-in FSDP (Fully Sharded Data Parallel) implements the stage-3 idea. The arithmetic is a script, not a cluster:

GB = 1024**3

def training_memory(n_params, n_gpus=1, zero_stage=0):
    """Bytes per GPU for mixed-precision Adam training (ZeRO-style sharding)."""
    weights_fp16 = 2 * n_params
    grads_fp16 = 2 * n_params
    # optimizer states kept in fp32: master weights + Adam m + Adam v
    opt_fp32 = (4 + 4 + 4) * n_params
    if zero_stage >= 1:
        opt_fp32 /= n_gpus
    if zero_stage >= 2:
        grads_fp16 /= n_gpus
    if zero_stage >= 3:
        weights_fp16 /= n_gpus
    return weights_fp16 + grads_fp16 + opt_fp32

N = 7_000_000_000
print(f"model: {N/1e9:.0f}B parameters, Adam, fp16/fp32 mixed precision")
print(f"inference only, fp16 weights:      {2*N/GB:6.1f} GiB")
print(f"training, single GPU:              {training_memory(N)/GB:6.1f} GiB")
for stage in [1, 2, 3]:
    m = training_memory(N, n_gpus=8, zero_stage=stage)
    print(f"training, 8 GPUs, ZeRO stage {stage}:    {m/GB:6.1f} GiB per GPU")
print("(activations, comm buffers, and fragmentation come on top of this)")

Running this prints:

model: 7B parameters, Adam, fp16/fp32 mixed precision
inference only, fp16 weights:        13.0 GiB
training, single GPU:               104.3 GiB
training, 8 GPUs, ZeRO stage 1:      35.9 GiB per GPU
training, 8 GPUs, ZeRO stage 2:      24.4 GiB per GPU
training, 8 GPUs, ZeRO stage 3:      13.0 GiB per GPU
(activations, comm buffers, and fragmentation come on top of this)

Pitfall. The napkin arithmetic above is a lower bound, and treating it as a prediction is a classic way to have a job die at 3 a.m. Activations scale with batch size and sequence length, not parameter count, and can rival the model state for long sequences; communication buffers, CUDA context, and allocator fragmentation add more. If the arithmetic says a job fits with 10% headroom, it does not fit. Budget the arithmetic plus activations (or use activation checkpointing, which recomputes them during the backward pass to trade compute for memory) and leave real margin.

The punchlines: training takes eight times the memory of inference before a single activation is stored (104 GiB vs. 13 GiB for a 7B model, beyond even an 80 GB flagship GPU), and sharding recovers almost all of the redundancy: with ZeRO-3 on 8 GPUs, per-GPU cost drops to 13 GB, the same as inference. The price is communication (stage 3 must gather weights layer by layer, every step), so the stages form a ladder you climb only as far as memory forces you to. This arithmetic also explains why LoRA is so effective for fine-tuning: with 99% of parameters frozen, the \(12N\) of optimizer state applies only to the adapters, and the whole ladder becomes unnecessary.

Quantization: inference in 8 bits (or fewer)

After training, the constraints flip: no optimizer state, no gradients — just weights and the desire to serve them cheaply. Quantization maps float weights onto a small grid of integers. The simplest scheme, symmetric int8, picks a scale \(s = \max \lvert W \rvert / 127\) and stores \(\operatorname{round}(W / s)\), an integer in \([-127, 127]\); dequantizing multiplies back. Storage drops 4x versus fp32, and int8 matrix multiplies are fast on most hardware. The question that decides everything: how much error does the rounding inject, and what makes it worse?

import torch

torch.manual_seed(0)

def quantize_int8(w, per_channel=False):
    """Symmetric int8 quantization; returns dequantized weights."""
    if per_channel:
        scale = w.abs().max(dim=1, keepdim=True).values / 127
    else:
        scale = w.abs().max() / 127
    wq = (w / scale).round().clamp(-127, 127)
    return wq * scale

def report(name, w, x):
    y = w @ x
    for mode, pc in [("per-tensor", False), ("per-channel", True)]:
        wdq = quantize_int8(w, per_channel=pc)
        w_err = (wdq - w).abs().max()
        y_err = (wdq @ x - y).norm() / y.norm()
        print(f"{name}, {mode:11}:  max weight err {w_err.item():.5f}  "
              f"relative matmul err {y_err.item():.5f}")

w = torch.randn(256, 256) * 0.05          # typical trained-weight scale
x = torch.randn(256, 64)
report("well-behaved W", w, x)

w_out = w.clone()
w_out[0, 0] = 5.0                         # one outlier weight, 100x the rest
report("W with outlier", w_out, x)

Running this prints:

well-behaved W, per-tensor :  max weight err 0.00090  relative matmul err 0.01032
well-behaved W, per-channel:  max weight err 0.00090  relative matmul err 0.00687
W with outlier, per-tensor :  max weight err 0.01968  relative matmul err 0.21665
W with outlier, per-channel:  max weight err 0.01964  relative matmul err 0.01491

For a well-behaved Gaussian weight matrix, int8 is nearly free: about 1% relative error in the matmul output. The second experiment is the important one. Change a single weight out of 65,536 to an outlier value and per-tensor quantization falls apart (a 21.7% matmul error) because the scale must stretch to cover the outlier, leaving every ordinary weight with a coarser grid. Per-channel quantization (one scale per output row) confines the damage to the outlier’s own row and restores 1.5% error. This tiny experiment is the shape of a real discovery: Dettmers et al. (NeurIPS 2022) showed that large transformers systematically develop outlier features that wreck naive int8, and that handling the outliers separately preserves accuracy — the insight behind the widely used LLM.int8() and, later, the 4-bit methods (GPTQ, QLoRA’s NF4) that let a 7B model run in about 4 GB.

People. Tim Dettmers did much of the quantization work that hobbyists and labs alike now depend on (LLM.int8(), the 4-bit QLoRA fine-tuning method, and the bitsandbytes library that packages them), largely during his PhD at the University of Washington. His stated through-line is making large models usable on accessible hardware, and it is hard to name work from this era that did more toward that goal in practice.

Pruning: removing weights outright

Quantization keeps every weight and shrinks its representation; pruning deletes weights entirely. The empirical license for this comes from a fact you can verify in seconds: trained networks are heavily redundant, and a large fraction of their weights can be set to exactly zero with little accuracy cost. The modern recipe (rank weights by magnitude, zero the smallest, then briefly retrain the survivors) was established at scale by Han, Pool, Tran, and Dally (NeurIPS 2015), whose paper reports cutting the weights of AlexNet-class networks roughly 9- to 13-fold without accuracy loss, using iterative rounds of prune-and-retrain.

The experiment below runs the whole idea on a small classifier, and adds the one design choice that dominates the outcome: where the threshold is computed. Global pruning ranks all weights together; per-layer pruning removes the same fraction from each layer. They sound interchangeable. They are not, and the failure mode is worth seeing in numbers before you meet it in a real model.

import torch
import torch.nn as nn

torch.manual_seed(0)

# Classification data: two interleaved moons
def moons(n, gen=None):
    t = torch.rand(n, generator=gen) * torch.pi
    x = torch.where(
        (torch.arange(n) % 2 == 0).unsqueeze(1),
        torch.stack([torch.cos(t), torch.sin(t)], dim=1),
        torch.stack([1 - torch.cos(t), 0.5 - torch.sin(t)], dim=1),
    ) + 0.15 * torch.randn(n, 2, generator=gen)
    y = (torch.arange(n) % 2).long()
    return x, y

g = torch.Generator().manual_seed(0)
x_tr, y_tr = moons(2000, g)
x_te, y_te = moons(1000, g)

net = nn.Sequential(nn.Linear(2, 64), nn.ReLU(), nn.Linear(64, 64), nn.ReLU(),
                    nn.Linear(64, 2))
opt = torch.optim.Adam(net.parameters(), lr=1e-2)
for step in range(1500):
    idx = torch.randint(0, len(x_tr), (128,))
    loss = nn.functional.cross_entropy(net(x_tr[idx]), y_tr[idx])
    opt.zero_grad(); loss.backward(); opt.step()

def accuracy():
    with torch.no_grad():
        return (net(x_te).argmax(1) == y_te).float().mean().item()

weights = [m.weight for m in net.modules() if isinstance(m, nn.Linear)]
original = [w.detach().clone() for w in weights]
print(f"trained baseline: {sum(w.numel() for w in weights)} prunable weights, "
      f"test acc {accuracy():.4f}")

def prune_and_eval(sparsity, per_layer):
    with torch.no_grad():
        for w, orig in zip(weights, original):
            w.copy_(orig)
        if per_layer:
            masks = [(w.abs() > w.abs().flatten().quantile(sparsity)).float()
                     for w in weights]
        else:  # one global threshold across all layers
            th = torch.cat([w.abs().flatten() for w in weights]).quantile(sparsity)
            masks = [(w.abs() > th).float() for w in weights]
        for w, m in zip(weights, masks):
            w.mul_(m)
    acc = accuracy()
    out_kept = masks[-1].mean().item()   # surviving fraction in the output layer
    # brief fine-tune with the masks enforced after each step
    opt_ft = torch.optim.Adam(net.parameters(), lr=1e-3)
    for step in range(200):
        idx = torch.randint(0, len(x_tr), (128,))
        loss = nn.functional.cross_entropy(net(x_tr[idx]), y_tr[idx])
        opt_ft.zero_grad(); loss.backward(); opt_ft.step()
        with torch.no_grad():
            for w, m in zip(weights, masks):
                w.mul_(m)
    return acc, accuracy(), out_kept

print(f"{'sparsity':>9} {'scheme':>10} {'acc pruned':>11} {'after tune':>11} "
      f"{'output layer kept':>18}")
for sparsity in [0.5, 0.8, 0.9, 0.95]:
    for per_layer in [False, True]:
        acc, acc_ft, kept = prune_and_eval(sparsity, per_layer)
        scheme = "per-layer" if per_layer else "global"
        print(f"{sparsity:>9.0%} {scheme:>10} {acc:>11.4f} {acc_ft:>11.4f} {kept:>17.0%}")

Running this prints:

trained baseline: 4352 prunable weights, test acc 0.9900
 sparsity     scheme  acc pruned  after tune  output layer kept
      50%     global      0.9890      0.9880               65%
      50%  per-layer      0.9910      0.9900               50%
      80%     global      0.9620      0.9840                4%
      80%  per-layer      0.7230      0.9890               20%
      90%     global      0.5000      0.5000                0%
      90%  per-layer      0.6020      0.8660               10%
      95%     global      0.5000      0.5000                0%
      95%  per-layer      0.7540      0.8880                5%

Four lessons, one per band of the table. At 50% sparsity, half the network was free: accuracy is unchanged before any retuning, under either scheme — that is the redundancy pruning lives on. At 80%, the raw post-prune accuracies wobble (0.96 global, 0.72 per-layer; single numbers this close to a cliff are noisy), but 200 steps of fine-tuning restores both to within half a point of baseline: retraining after pruning is not optional polish, it is the step that does the work. At 90%, the schemes separate for a structural reason, visible in the last column: the global threshold has deleted 100% of the output layer (its 128 weights are small compared to hidden-layer weights, so a global ranking removes them first), and a network with no output layer is a coin flip (0.5000) that no amount of masked fine-tuning can revive. Per-layer pruning, which by construction cannot annihilate a layer, degrades gracefully and fine-tunes back to 0.87. And at 95%, even per-layer pruning has passed this small network’s redundancy budget: 0.89 after tuning, a real loss. Every number generalizes qualitatively to large models: substantial free sparsity, retraining essential, allocation across layers a first-order decision, and a budget beyond which the function genuinely degrades.

One caution keeps this honest: a tensor full of zeros is not automatically faster. Dense matrix-multiply kernels spend the same time on a zero as on any other number, so unstructured sparsity (arbitrary zero positions, like ours) mostly saves storage unless specialized sparse kernels apply. Hardware-friendly formats impose structure: removing whole neurons or channels, or the 2:4 pattern (two zeros in every four consecutive weights) that recent NVIDIA GPUs accelerate natively. Structured pruning at matched sparsity usually costs more accuracy, because the constraint binds. When a paper or a vendor quotes a sparsity number, the first question is always: structured how, and measured on what hardware?

Historical note. Pruning is one of the oldest ideas in the field’s toolbox. At the close of the 1980s, LeCun, Denker, and Solla’s “Optimal Brain Damage” used second-derivative information to decide which weights of a trained network mattered least, with the explicit goals of speed and generalization that motivate pruning today. The idea then waited two and a half decades for networks large enough to make compression urgent. A modern echo: the lottery-ticket hypothesis (Frankle & Carbin, ICLR 2019) reports that dense networks contain sparse subnetworks that could have trained to full accuracy from the start given the right initialization — suggesting the redundancy pruning removes was partly a scaffold for optimization rather than part of the final function.

Distillation: compressing the function, not the numbers

Quantization shrinks the numbers; distillation shrinks the function. Hinton, Vinyals, and Dean (2015) formalized the recipe: train a small “student” network to match a large “teacher” network’s full output distribution (its soft targets) rather than only the hard labels. The soft targets carry information the labels lack: on an image of a 7, a good teacher assigns a 2 some probability and a truck essentially none, and that relational structure (what Hinton has called dark knowledge) tells the student how classes relate. In practice both teacher and student logits are softened by a temperature \(T \gt 1\) in the softmax before comparing, so the small probabilities become visible to the loss.

Distillation is the third leg of the deployment stool because it composes with everything else: DistilBERT (Sanh et al., 2019) is the canonical demonstration — the paper reports a student 40% smaller and 60% faster than BERT while retaining about 97% of its language-understanding performance. The pattern now appears wherever a big model exists and a cheap one is wanted: instruction-following students trained on frontier-model outputs, and the step-distilled diffusion samplers mentioned in the diffusion section. A useful mental model so far: mixed precision and ZeRO make training affordable; quantization, pruning, and distillation make the model cheap. What remains is the cost that appears only when the model meets its users (the cost of serving requests), and it has its own arithmetic, which the next two sections work out.

Serving transformers: the KV cache

When a transformer generates token by token (decoding), each new token’s attention must look back at the keys and values of every earlier position. Recomputing those from scratch at every step would make generation quadratic in practice, so every serving stack caches them: the KV cache stores, for each layer and each attention head, the key and value vectors of every token generated or ingested so far. The cache is pure bookkeeping (no approximation), but it is state that grows linearly with context, and at modern context lengths it competes with the weights themselves for memory. The per-token cost is one multiplication: 2 (K and V) * layers * kv_heads * head_dim * bytes. Everything about long-context serving follows from where that number lands.

The script compares two real architectural choices at 7–8B scale: classic multi-head attention, where all 32 heads store their own K and V, and grouped-query attention (GQA), where groups of query heads share one of just 8 KV heads — the design Llama-3-class models adopted (following Ainslie et al., EMNLP 2023, and the earlier single-KV-head “multi-query” proposal by Shazeer, 2019).

GB = 1024**3                                  # GiB (binary); hardware sizes like '80 GB' are names
MB = 1024**2

def kv_bytes_per_token(n_layers, n_kv_heads, head_dim, bytes_per_val=2):
    # K and V, per layer, per kv head, per head dimension
    return 2 * n_layers * n_kv_heads * head_dim * bytes_per_val

# Two real architectures (both ~7-8B parameters, fp16 weights ~13-15 GB):
# multi-head attention (every head stores K/V) vs grouped-query attention (8 shared KV heads)
configs = [
    ("7B, MHA (Llama-2-7B-like)", 32, 32, 128),
    ("8B, GQA (Llama-3-8B-like)", 32, 8, 128),
]

print(f"{'model':28} {'KV/token':>9} {'4k ctx':>8} {'32k ctx':>8} {'128k ctx':>9}")
for name, L, H, D in configs:
    per_tok = kv_bytes_per_token(L, H, D)
    row = [per_tok * n / GB for n in (4096, 32768, 131072)]
    print(f"{name:28} {per_tok/1024:>7.0f}KB {row[0]:>7.2f}G {row[1]:>7.2f}G {row[2]:>8.2f}G")

# Serving arithmetic: one 80 GB GPU, fp16 weights loaded, rest is KV budget
weights = {"7B, MHA (Llama-2-7B-like)": 2 * 7e9, "8B, GQA (Llama-3-8B-like)": 2 * 8e9}
print()
print("concurrent 8k-token sequences that fit in 80 GB alongside the weights:")
for name, L, H, D in configs:
    kv_per_seq = kv_bytes_per_token(L, H, D) * 8192
    budget = 80 * GB - weights[name]
    print(f"{name:28} {int(budget // kv_per_seq):>4d} sequences "
          f"({kv_per_seq/GB:.2f} GiB per sequence)")

Running this prints:

model                         KV/token   4k ctx  32k ctx  128k ctx
7B, MHA (Llama-2-7B-like)        512KB    2.00G   16.00G    64.00G
8B, GQA (Llama-3-8B-like)        128KB    0.50G    4.00G    16.00G

concurrent 8k-token sequences that fit in 80 GB alongside the weights:
7B, MHA (Llama-2-7B-like)      16 sequences (4.00 GiB per sequence)
8B, GQA (Llama-3-8B-like)      65 sequences (1.00 GiB per sequence)

Half a megabyte per token sounds small until you multiply: at a 128k-token context, the MHA model’s cache is 64 GiB (five times the size of the model’s own weights) for a single conversation. That one row explains why long context windows arrived years after long-context research: the bottleneck was never the attention math alone, it was the state. It also explains why GQA swept through model families: cutting KV heads from 32 to 8 divides all cache numbers by 4 at nearly no quality cost (the GQA paper’s finding). The serving consequence is immediate: 65 concurrent 8k conversations per GPU instead of 16, which is the difference between a viable product and an underutilized one. When you read a model card, the KV-head count is a serving-cost disclosure hiding in an architecture table.

Batching: throughput against latency

The second serving fact is stranger: during generation, a big GPU is almost idle. Producing one token requires reading every weight from memory once but does only a couple of floating-point operations per weight read, and modern GPUs can perform hundreds of operations in the time one operand arrives from memory. Single-stream decoding is therefore memory-bandwidth-bound: the compute units wait on the memory system. The escape is batching. Serving many users at once reuses each weight read across the whole batch, multiplying throughput at almost no extra time per step — until the batch’s aggregate KV-cache traffic becomes the new bottleneck. The simulation below prices this out with the simplest defensible model: each decode step streams the weights plus every active sequence’s KV cache through memory at a fixed bandwidth.

GB = 1024**3

# A deliberately simple model of memory-bound decoding:
# each decode step must stream (weights + the batch's KV cache) from GPU memory.
# time per step ~= bytes moved / memory bandwidth.
weights_bytes = 2 * 8e9                     # 8B model, fp16
kv_per_seq = 131072 * 4096                  # 8k-token context, GQA (from the KV script)
bandwidth = 2.0e12                          # 2 TB/s HBM, an A100-80GB-class figure

print(f"{'batch':>6} {'step time':>10} {'per-user tok/s':>15} {'total tok/s':>12} "
      f"{'GPU-seconds/1M tok':>19}")
for batch in [1, 2, 4, 8, 16, 32, 64, 128, 256]:
    bytes_per_step = weights_bytes + batch * kv_per_seq
    step_time = bytes_per_step / bandwidth          # seconds per decode step
    per_user = 1 / step_time                        # each user gets 1 token per step
    total = batch / step_time
    cost = 1e6 / total
    print(f"{batch:>6} {step_time*1000:>8.2f}ms {per_user:>15.1f} {total:>12.0f} "
          f"{cost:>19.1f}")

Running this prints:

 batch  step time  per-user tok/s  total tok/s  GPU-seconds/1M tok
     1     8.27ms           120.9          121              8268.4
     2     8.54ms           117.1          234              4268.4
     4     9.07ms           110.2          441              2268.4
     8    10.15ms            98.5          788              1268.4
    16    12.29ms            81.3         1301               768.4
    32    16.59ms            60.3         1929               518.4
    64    25.18ms            39.7         2542               393.4
   128    42.36ms            23.6         3022               330.9
   256    76.72ms            13.0         3337               299.7

Read the first rows and the last rows as two different regimes. From batch 1 to 8, total throughput multiplies 6.5-fold while each user’s speed falls only 19% — the batch is riding along on weight reads that had to happen anyway, which is why batching is the closest thing to free money in model serving. By batch 256, the regime has flipped: KV traffic (256 sequences at 0.5 GB each) dwarfs the 16 GB of weights, so throughput is flattening (adding users no longer helps much) while per-user speed has fallen to 13 tokens/s, around the threshold where a chat interface starts to feel slow. The economics column makes the operator’s dilemma exact: serving cost per million tokens falls 28-fold across the table, so the provider wants the largest batch users will tolerate, and the users want the reverse. Every serving stack lives on this curve; the arithmetic is a design tool, not trivia. (The model is deliberately crude: it ignores prompt prefill, which is compute-bound, assumes perfect overlap, and holds context length fixed. But the two regimes and the knee between them are real, and this is how you locate them for any hardware and model by changing three constants.)

Pitfall. Latency and throughput are different quantities with different owners, and a benchmark that reports only one of them is answering a question you may not have asked. “Model X serves 3,300 tokens per second” (batch 256, in our table) and “Model X generates at 13 tokens per second per user” describe the same configuration. Vendor throughput numbers are usually quoted at large batch; your interactive application lives at the small-batch end of the curve. Always ask: at what batch size, context length, and per-user latency was this measured?

Going deeper

Mixture of experts: buying parameters without proportional compute. In a sparse mixture-of-experts (MoE) Transformer, each feed-forward block is replaced by several parallel expert blocks plus a small router that sends each token to a subset of them, so total parameters scale with the number of experts while compute per token scales only with the number consulted. Mixtral 8x7B (eight experts per layer, two active per token) has about 47B total parameters and about 13B active per token. That arithmetic does not make it literally a “70B-capacity” model: capacity is not a scalar determined by total or active count, and comparisons with dense 70B models are empirical benchmark claims. The precise systems trade is 47B-parameter memory residency with roughly 13B-parameter per-token compute, plus routing and communication overhead.

Speculative decoding: spending idle compute to save time. The batching section showed single-stream decoding leaves the GPU’s arithmetic units starved. Speculative decoding (Leviathan, Kalman & Matias, ICML 2023; concurrently Chen et al., 2023) turns that idle compute into latency reduction: a small, fast draft model proposes several tokens ahead; the large model then scores all of them in one parallel forward pass (the parallel pass costs barely more than generating a single token, precisely because decoding is memory-bound) and accepts the longest prefix consistent with what it would have sampled itself. A rejection-sampling correction makes the output distribution provably identical to the large model’s own: this is a pure speed trick, not an approximation. Reported speedups in the original papers are around 2–3x on production-style models. It composes with everything else here (a quantized or distilled model often serves as the draft), and it echoes a theme worth noticing: the distillation student and the draft model are the same kind of object — a cheap approximation of an expensive function, deployed where its errors can be checked or tolerated.

Systems software as an efficiency layer. Two software ideas closed much of the gap between napkin arithmetic and real utilization. FlashAttention (Dao et al., NeurIPS 2022) computes exact attention in a tiling pattern that avoids materializing the full attention matrix in slow memory (an IO-aware reorganization, not an approximation) and became a default kernel almost immediately. PagedAttention, the core idea of the vLLM serving system (Kwon et al., SOSP 2023), manages the KV cache the way an operating system manages RAM: in small pages, allocated on demand, eliminating the fragmentation that wasted much of the cache budget our arithmetic assumed was usable. Neither changes a single model weight; both moved real-world cost more than many modeling papers. If the efficiency story of this section has one moral, it is that accounting (of bytes, of bandwidth, of idle compute) is where the wins come from, and the accounting is arithmetic you can now do.

The bitter arithmetic of formats, continued. The precision ladder keeps descending: fp8 formats shipped with H100-class hardware for training and inference, and 4-bit weight quantization is now routine for serving, with active research at 2–3 bits and on quantizing the KV cache itself (whose size, as you computed, rivals the weights at long context). Each step down doubles capacity per gigabyte and re-runs the same question the fp16 story answered: which tensors tolerate the coarser grid, and which need outlier handling or higher-precision islands. The methods change; the napkin stays.

Common pitfalls

  • Loss becomes NaN partway through fp16 training. Cause: usually overflow to inf followed by invalid arithmetic, or another unstable operation; underflow usually produces zeros and stalled updates rather than NaNs by itself. Fix: inspect finite values and gradient norms, use autocast/GradScaler where appropriate, or use bf16 on supported hardware.
  • A model trains in bf16 but a numerically sensitive metric drifts from the fp32 baseline. Cause: bf16 carries only about 2–3 decimal digits; reductions over many terms accumulate rounding. Fix: keep sensitive accumulations (losses, norms, softmax denominators) in fp32 — which is what autocast does for you.
  • “The model is 7B parameters so it needs 14 GB” — and the training job dies at 100+ GB. Cause: counting only low-precision weights. Fix: declare the optimizer/precision layout, estimate its weights, gradients, moments, master copies, and activations, then account for sharding/offload.
  • A quantized model’s average benchmark score looks fine, but specific behaviors break. Cause: outlier weights or activations stretching per-tensor scales, injecting error concentrated in particular channels. Fix: per-channel scales, outlier-aware schemes, and evaluation that looks at more than one aggregate number — the subject of the S27 reading.
  • A model is pruned to 90% sparsity, storage shrinks, and inference speed does not change. Cause: unstructured zeros run through dense kernels at dense speed. Fix: structured sparsity that your hardware actually accelerates (channel/neuron removal, 2:4 patterns), and benchmark wall-clock time, not just parameter counts.
  • An OOM appears in week three of serving, at the same traffic that ran fine in week one. Cause: KV-cache growth with longer conversations (memory that scales with context, not with request count) plus allocator fragmentation. Fix: do the KV arithmetic for your worst-case context and concurrency (not the average), and use paged cache management.

Try it yourself

  1. Extend the quantization script to int4 (scale by \(\max \lvert W \rvert / 7\)). How do the two matmul errors change? At what bit width does the well-behaved matrix’s error exceed 5%?
  2. Extend the memory script to a 70B model and per-GPU memory of 80 GB. How many GPUs does each ZeRO stage need just to hold the training state (ignoring activations)?
  3. Simulate loss scaling end to end: generate 10,000 gradients from \(\mathcal{N}(0, (10^{-7})^2)\), cast to fp16 with and without a 1024x scale, and report what fraction survive as nonzero in each case.
  4. Compute the memory arithmetic for LoRA fine-tuning of a 7B model with 0.5% trainable parameters: fp16 frozen weights plus full Adam state for the trainable part only. Compare with full fine-tuning.
  5. In the pruning script, prune only the two hidden layers and leave the output layer untouched, then redo the global-pruning rows. How far does the 90% catastrophe move? What does this tell you about where the global scheme’s failure actually lived?
  6. Extend the batching simulation: hold the total memory at 80 GB and recompute the maximum feasible batch at context lengths 2k, 8k, 32k (batch limited by KV fitting in memory). Plot or tabulate achievable total throughput versus context length. Where does long context actually cost you?
  7. Redo the KV arithmetic for a 70B model (80 layers, 8 KV heads, head_dim 128) at 128k context. Does a single 80 GB GPU still hold even one sequence alongside int4-quantized weights (about 35 GB)?

Further reading

  • Micikevicius et al., “Mixed Precision Training,” ICLR 2018 — the master-weights/loss-scaling recipe that made 16-bit training standard. arxiv.org/abs/1710.03740
  • Rajbhandari et al., “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models,” SC 2020 — the sharding ladder behind our memory table and PyTorch’s FSDP. arxiv.org/abs/1910.02054
  • Dettmers et al., “LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale,” NeurIPS 2022 — the outlier-feature discovery our quantization experiment miniaturizes. arxiv.org/abs/2208.07339
  • Han et al., “Learning both Weights and Connections for Efficient Neural Networks,” NeurIPS 2015 — magnitude pruning with iterative retraining, at ImageNet scale. arxiv.org/abs/1506.02626
  • Hinton, Vinyals & Dean, “Distilling the Knowledge in a Neural Network,” NeurIPS 2014 Deep Learning Workshop — soft targets, temperature, and the dark-knowledge framing. arxiv.org/abs/1503.02531
  • Sanh et al., “DistilBERT, a distilled version of BERT,” 2019 — the canonical demonstration that distillation preserves most quality at a fraction of the cost. arxiv.org/abs/1910.01108
  • Kwon et al., “Efficient Memory Management for Large Language Model Serving with PagedAttention,” SOSP 2023 — the vLLM paper; OS-style paging for the KV cache whose size you just computed. arxiv.org/abs/2309.06180
  • Leviathan, Kalman & Matias, “Fast Inference from Transformers via Speculative Decoding,” ICML 2023 — draft-and-verify decoding with a proof that the output distribution is unchanged. arxiv.org/abs/2211.17192
  • Jiang et al., “Mixtral of Experts,” 2024 — an open, fully documented sparse MoE: 47B parameters, ~13B active per token. arxiv.org/abs/2401.04088
  • PyTorch documentation: Automatic Mixed Precision (docs.pytorch.org/docs/stable/amp.html) and FSDP (docs.pytorch.org/docs/stable/fsdp.html) — the APIs implementing the first half of this section.

Before class you should be able to

  • Contrast fp16 and bf16 by their exponent/mantissa split and predict which failure (overflow vs. coarse rounding) each risks.
  • Explain the three parts of the mixed-precision recipe: fp32 master weights, loss scaling, 16-bit forward/backward.
  • Reproduce the stated-layout napkin arithmetic that totals about 16 bytes per parameter, name variants that change it, and explain what ZeRO stages 1–3 shard.
  • Explain why one outlier weight breaks per-tensor int8 quantization and how per-channel scales confine the damage.
  • State why unstructured sparsity saves storage but usually not wall-clock time, and why retraining after pruning is essential.
  • Compute KV-cache bytes per token from layers, KV heads, head dimension, and bytes per value — and explain why GQA cuts serving cost roughly 4x.
  • Explain why single-stream decoding is memory-bandwidth-bound and why batching multiplies throughput at little cost until KV traffic dominates.

Connects to

  • This session is covered by Quiz 4 (Module 4: S23–S26), given at S27.
  • HW4 (evaluation + efficiency) goes out today, due at S29, the last class meeting. It is NumPy only: bootstrap CIs, a paired model comparison, stable perplexity, and int8 quantization with a before/after error, the last of which uses exactly this section’s tools.
  • The memory arithmetic closes loops from S18 (what compute budgets mean) and S19 (LoRA, whose optimizer-state savings you can now compute exactly).

← S25 · Diffusion Models · All readings · S27 · Evaluation and Safety →