Making knowledge distillation practical at scale

Knowledge distillation, training a smaller student model to match the performance of a larger teacher, is a foundational technique in machine learning. As massive open-source models proliferate, compressing them into efficient student models has become standard practice. However, the distillation step itself remains computationally heavy and expensive.

The cost of traditional distillation

The standard setup, known as online distillation using Kullback-Leibler divergence loss, keeps both the teacher and the student loaded simultaneously. At every training step, the teacher runs a full forward pass to produce an output distribution across its entire vocabulary. For models with hundreds of thousands of vocabulary tokens and long sequence lengths, the resulting probability tensors consume enormous amounts of VRAM.

For instance, a model with a vocabulary of over 200,000 tokens evaluated at a sequence length of 32K and a batch size of 4 can see a single training iteration peak at roughly 250GB of VRAM. This exceeds the capacity of many modern single-GPU setups and typically demands clusters of hundreds of GPUs using careful tensor-parallelism strategies.

Two systems changes for efficiency

To address these memory constraints, researchers have introduced two primary systems modifications: offline distillation and a fused, chunked KL loss.

Offline distillation computes the teacher's output just once, caching the top-100 most likely tokens per position. The teacher model is then removed from memory entirely during student training, and the same cache can be reused across multiple experimental runs.

The second change involves how the KL loss is computed. Standard implementations build a massive grid comparing every vocabulary entry against every sequence position before generating a single number. The new fused chunked KL loss processes data in manageable slices. It integrates the model's output projection directly into the loss computation, discarding each sequence chunk after processing rather than storing a full vocabulary-by-sequence matrix.

Performance at long contexts

By avoiding the massive memory spikes associated with dense loss grids, the fused chunked approach scales effectively as context lengths grow. At 32K tokens, peak memory drops significantly compared to traditional dense loss methods, and the approach remains stable at lengths where dense losses fail outright. In practical tests distilling a 20B model at a 32,768-token context, memory reductions allowed the workload to shrink from four GPU nodes down to a single device, substantially accelerating step times and increasing throughput.