Knowledge distillation is the standard way to compress large open LLMs into smaller, deployable models, but it is often expensive because the teacher and student must be loaded together and the teacher's full-vocabulary outputs recomputed at every step. A new paper from Hugging Face, titled "Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss," presents two systems-level contributions that make distillation training significantly more efficient.
The first contribution is offline distillation: caching the teacher's top-K logits once and training the student against that cache. This removes the teacher from memory during training, runs about 29% faster per iteration, and achieves up to 41% higher throughput on a single H200 GPU, all while matching the training loss of online distillation.
The second contribution is a fused, chunked KL loss that never materializes the full vocabulary-sized logit tensor. By processing the sequence in slices, peak memory becomes linear in sequence length, eliminating the memory spike that caps context length. This enables training at four times the context length (32,768 tokens) on a single GPU. The authors also release their implementation at https://github.com/CompactifAI/Full-Chunked-KL-Loss.
The paper includes supporting ablations on loss design and sequence packing, and notes that for both offline and online KD, logits should be computed on data that is in-distribution for the teacher—ideally the teacher's original training data or synthetic data generated by the teacher. The authors also mention that their current loss assumes shared vocabulary between teacher and student, and they are interested in exploring cross-family distillation in the future.