
Making Knowledge Distillation Cheap Enough to Run at Scale

Knowledge distillation is a standard way to compress large open-source LLMs, and recent releases such as gpt-oss, Qwen, GLM, and Kimi have brought it back to the foreground. Deployment is expensive: Kimi-K3 has 2.8 trillion parameters and requires roughly 3TB of VRAM just to load, so compressing and recovering capability through distillation has become common. The distillation step is usually the most expensive part of the pipeline. A new paper from Multiverse Computing, Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss, proposes two systems changes: cache the teacher’s top-K logits once so it can be removed from memory, and use a fused chunked KL loss that never materializes a full vocabulary-size × sequence-length matrix. The combined effect is to cut VRAM use below standard PyTorch or NVIDIA Megatron-Bridge implementations and make long-context healing practical on a single GPU.
Why is online distillation so expensive? The standard setup keeps teacher and student loaded simultaneously. Each step runs a full teacher forward pass and holds two full-vocabulary tensors per token position. With gpt-oss-120b, vocabulary is 201,088; at sequence length 32K and batch size 4, the teacher probability tensor alone is 4 × 201,088 × 32,768, about 50GB in bfloat16. Including gradients, activations, weights, and optimizer states, one iteration peaks near 250GB, beyond an H200’s 141GB. The dense KL baseline spikes to that level, while the fused chunked loss peaks around 128GB in the paper’s Figure 1.
The first change is offline distillation: compute teacher output once, cache the top-100 most likely tokens per position, and train against that cache. The teacher no longer sits in memory and need not be rerun, so one cache can be reused across ablations. The second change is the loss implementation. Dense KL rebuilds a full teacher-probability grid from the cached logits and compares it with the student’s dense grid; it is the correctness baseline but holds the full vocabulary × sequence grid twice. Forward-chunked KL keeps the teacher sparse and computes loss in slices, but still builds and stores the student’s full logit grid. Fused chunked KL, the main contribution, fuses the output projection into the loss. It processes one chunk of sequence positions end to end, projects hidden states to logits for that chunk, folds it into a running loss, and discards it before the next chunk; the backward pass recomputes chunks instead of storing them. The extra projection in the backward pass costs some speed, but peak memory grows only linearly with sequence length. The implementation is open-sourced at github.com/CompactifAI/Full-Chunked-KL-Loss.
Results on a single H200 with Llama 3.1 8B Instruct as teacher and a 3.2B Llama student at 8K context: online distillation peaks at 102.8GB and 25.9s per iteration; offline dense KL at 78.3GB and 18.5s; forward-chunked at 61.8GB and 18.4s; fused chunked at 58.3GB and 20.2s. Training losses overlap almost exactly across all four methods, so the offline top-100 approach is lossless relative to online distillation in this setup. On a toy output-projection benchmark at 32K tokens, peak memory drops from 85.2 GiB with dense loss to 5.45 GiB with fused chunking, a 15.6× reduction, and dense loss fails at 64K. At 256K, fused chunking uses 11.6 GiB versus 134.2 GiB for the next-best variant and is about 3.3× faster per iteration. Distilling GPT-OSS 20B at a 32,768-token context, the memory savings shrink the setup from four GPU nodes to one: step time falls from 57.0 to 12.23 seconds, about 5×, and throughput per GPU rises from 74.2 to 345.7 TFLOP/s.
The resulting compact student, distilled from Llama 3.1 8B Instruct to about 3.2B parameters, retains most of the teacher’s accuracy on BoolQ and HellaSwag and stays within about nine points on MMLU, at less than half the parameter count. The paper also covers other ablations on loss choice and sequence packing. The authors frame this as ongoing research at Multiverse Computing into making distillation and healing cheap enough to iterate on at scale.


