Chapter 6 — Pruning and Knowledge Distillation

Published on: 2026-04-28 Last updated on: 2026-07-07 Version: 2
Chapter 6 — Pruning and Knowledge Distillation

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.

One line: Pruning and distillation attack the weight count; combined with quantization they compound into a 20× reduction in the bytes the chip streams per token.

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.

Worth holding onto: Compression is per-inference, not per-request. Every technique in this chapter makes one forward pass cheaper; it takes the batching and KV-management work of Chapters 7 and 8 to convert that cheapness into concurrent-user throughput.

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.

Want the full picture? The book chapter includes the runnable 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 →

SHO
SHO
CTO of Receipt Roller Inc., he builds innovative AI solutions and writes to make large language models more understandable, sharing both practical uses and behind-the-scenes insights.