Chapter 6 — Pruning and Knowledge Distillation
Sixth post of the chapter-by-chapter walkthrough of LLM Primer VI: Scaling AI Systems. The chapter that attacks weight count directly — first by zeroing out weights that do not matter, then by transferring a big model's behavior into a smaller one.
Why this chapter exists
Quantization shrinks the bytes per weight. Pruning shrinks the count of weights that get multiplied at all. Knowledge distillation transfers the behavior of a large teacher into a smaller student that runs cheaper end to end. Pruning is a same-architecture edit that shaves overhead; distillation is an architecture change that accepts a training-time cost in exchange for a permanent inference-time saving. Chapter 6 walks both — the mechanics each relies on, the failure modes practitioners hit, and the hardware-accelerated pattern (2:4 sparsity) that has made structured pruning economically interesting again on Hopper and Blackwell.
6.1 2:4 sparsity is the pruning pattern that hardware rewards
Unstructured magnitude pruning — zero out the smallest 50 percent of weights and fine-tune briefly — works numerically but does not accelerate. GPUs execute dense GEMMs; scattered zeros are still multiplied. Structured pruning removes whole units the hardware can skip: attention heads, MLP neurons, or (the important case) N:M patterns within weight matrices. Hopper and Blackwell support 2:4 sparsity natively — exactly two zeros in every four consecutive weights along the input dimension. The chip stores each four-weight group as two non-zero values plus a two-bit mask, the sparse tensor core skips the zero multiplications, and both bandwidth and FLOPs drop by roughly a factor of two. In measured decoding throughput that is a 1.7–1.9× speedup. The trick is that 2:4 must be imposed during pruning: the algorithm walks each matrix in groups of four and forces the two smallest to zero, then either fine-tunes briefly or applies SparseGPT to compensate. The mask is metadata; the runtime saving is automatic.
6.2 Distillation transfers the teacher's distribution, not its argmax
A trained teacher LLM produces at each position a distribution over the vocabulary — a 128k-dimensional softmax. The argmax is the token it would emit at greedy decoding. But the distribution encodes the teacher's uncertainty, its preference structure across near-synonyms, and its calibration. Training a student to match that distribution transfers all of it: the loss is a KL divergence between softened teacher and student outputs, with a temperature that widens the gradient signal into the low-probability tail. The student that emerges behaves like a smaller, faster version of the teacher rather than a smaller model trained from scratch on the same data. Patient Knowledge Distillation extends the recipe to intermediate layers, matching hidden states between mapped student-teacher layer pairs. MiniLLM addresses a subtler failure — the mode-averaging problem, where a forward-KL student hedges across multiple plausible teacher continuations and produces a bland output — by switching to reverse-KL, which prefers a sharp match to one of the teacher's modes over a soft coverage of all of them.
6.3 The three compressions stack, in a specific order
Distill first, prune second, quantize last. Distillation is the most expensive and most architecturally invasive operation; it makes sense to do it once on a dense high-precision model before other choices lock in. Prune the student to 2:4 with SparseGPT and a brief recovery fine-tune. Quantize the sparsified student to FP8 with llm-compressor. A worked stack: 70B BF16 teacher → 13B BF16 distilled student → 2:4-sparsified → FP8-quantized. Weight bytes per forward pass drop from 140 GB to roughly 6.5 GB — a 22× reduction — and Hopper's FP8 and 2:4 paths turn that into an 8–10× throughput gain per GPU. What the stack does not reduce is the KV cache, which depends on hidden dimension, head count, and context length rather than weight count. That is the subject of Chapter 8. And model-level compression is per-inference; it gives the GPU headroom, but the runtime lever that fills the headroom with concurrent users is batching.
What Chapter 6 sets up
Chapters 5 and 6 exhaust the model-side compressions — quantize the bytes per weight, prune the count of weights, distill the whole model down. The GPU now has headroom on every forward pass. Chapter 7 turns to the system-level lever that fills that headroom with useful work: batching. Static batching almost works and collapses on the fastest-finisher problem. Continuous batching, with iteration-level scheduling, is what production engines run and what turns the bandwidth-bound decode phase from an idle chip into a chip serving many users concurrently.
Next — Chapter 7: Advanced Batching Strategies. Static, dynamic, and continuous batching — and the debt continuous batching creates for the KV cache.
torch.sparse code for 2:4 conversion, the full PyTorch distillation loop with KL loss and temperature, the PKD layer-mapping addition, the MiniLLM reverse-KL formulation, and the In Plain English sidebar on why students hedge and how to stop them. View LLM Primer VI on Amazon →