Vol. XXIV · Making Inference Cheap Plate 6 of 10distillation, small models, Gemma 2, deployment, memory bandwidth, tokens per second, MMLUCE

Smaller Students.

Every byte of weights is read for every token. A model a third the size is three times the speed.

Quantisation shrinks the bytes per parameter; the other lever is fewer parameters. A small model trained the ordinary way is much worse than a large one, but a small model trained to match a large teacher’s full predicted distribution (plate XII.7) closes part of the gap. Gemma 2’s technical report gives a clean case: its 2B and 9B models were trained by distillation, the 27B by next-token prediction, and in an ablation a 2B model trained on 500 billion tokens averaged clearly higher on three benchmarks when distilled from a 7B teacher than when trained from scratch.

The serving side is simple arithmetic: single-user decode reads every weight once per token, so tokens per second are bounded by bandwidth over model bytes (plate XXIII.1). The workshop places the three Gemma 2 sizes, at the chosen precision and device, on a chart of benchmark score against that bound. Scores are the report’s 5-shot MMLU; speeds are bandwidth bounds at 80% of the published peak, not measurements.

Below: MMLU score against the single-user decode bound for Gemma 2 2B, 9B and 27B on the chosen device and precision; readouts for the chosen size.

What you will learn
Workshop
Gemma 2 parameter counts (embedding included) and 5-shot MMLU from the Gemma 2 technical report, Tables 2 and 13. Bandwidth: H100 3.35 TB/s; Apple M4 Max 546 GB/s and M4 120 GB/s (Apple specifications). Bound: 80% of peak bandwidth divided by weight bytes; cache reads, compute and overheads ignored.

Most of the score at a third of the bytes

Gemma 2 9B scores 71.3 on MMLU against the 27B’s 75.2 — 95% of the score — with 2.9 times fewer parameters to read per token; the 2B scores 52.2. On an M4 Max laptop at 4 bits the 9B can decode at most about 95 tokens a second for one user and the 27B about 32; the 27B in BF16 on an H100, about 49. Distillation is part of why the small ones hold up: in the report’s ablation it added 7.4 points to a 2B model’s benchmark average at the same training tokens.

For deployment the levers multiply: a distilled model a third the size, quantised to a quarter of the bytes, reads a twelfth as much per token. What is lost is real but uneven — small models lose most on knowledge-heavy and multi-step tasks, least on narrow, well-specified ones — which is why deployed systems often route easy requests to a small model and hard ones to a large one, and why distilling a large model’s behaviour on one task into a small one (plate XVII.9) is a common production step.

A benchmark average hides where a small model falls short; before replacing a large model, evaluate the small one on the traffic it will actually serve (Vol. XXI).

Terms

Student
The smaller model trained to imitate the teacher.
Teacher
The larger model whose distribution is imitated.
Routing
Sending each request to a model sized for it.
Bandwidth bound
Bandwidth over model bytes: the single-user decode ceiling.
Smaller models are faster in proportion to their bytes; distillation is what lets them keep much of a large model’s quality.

Sources

  1. Gemma Team, “Gemma 2: Improving Open Language Models at a Practical Size”, arXiv:2408.00118, 2024, Tables 2, 6 and 13.
  2. Hinton, Vinyals and Dean, “Distilling the Knowledge in a Neural Network”, 2015.
  3. Apple, M4 Pro and M4 Max announcement, October 2024; NVIDIA H100 datasheet.