Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss
Bakbergen Ryskulov, Iker García-Ferrero, David Montero, David Jansen, Ali Hashemi, Jezabel R. Garcia, Antonio Tiene, Román Orús
cs.CL, cs.AI, cs.LG
2026-08-04
Distilling Llama-3.1-8B into a roughly 3.2B student, the authors cache the teacher's top-100 logits offline and add a fused chunked KL loss that never materializes the full vocabulary tensor. Per-iteration time drops 29% and throughput rises about 40%; a single H200 trains at 32K context, and the student reaches 60.6% MMLU and 67.5% GSM8K.
Small models are often the only option under tight latency, cost, and on-premises constraints, but they are rarely trained from scratch: knowledge distillation (KD) compresses a large model's capability into a small one, and this step largely decides final quality. It is expensive. In standard online distillation the teacher and student both sit in memory, and computing the KL divergence requires inflating logits to the full vocabulary size (often over a hundred thousand dimensions). That vocabulary-sized tensor grows with sequence length, and the memory spike it creates caps the context length you can train.
This is a practitioner's study of how to make distillation efficient, built on two systems-level changes. The authors are from CompactifAI, a team that works on quantized Llama compression, and the chunked-loss code is open source.
Two components, each usable on its own.
First, offline top-K distillation. Standard online distillation runs a teacher forward pass every batch and uses its full distribution as the target. Here the teacher's top-100 logits per position are cached once, and the student trains only against that sparse cache. The teacher leaves the training loop and frees its memory. Top-100 instead of the full vocabulary is what makes the engineering cheap.
Second, a fused chunked KL loss. Standard KL runs the student's hidden state through the output head and materializes a huge "sequence length x batch x vocabulary" logit tensor before computing the loss. Here the output projection is fused into the loss kernel and the sequence is processed in chunks that are discarded immediately, so peak memory goes from quadratic in sequence length to linear. The spike that capped context length is gone.
Online versus offline (8K context, single H200): training loss is near-identical; memory drops from 103GB to 78GB; per-iteration time falls from 25.9s to 18.5s (about 29% faster); throughput rises from 237 to 331 TFLOP/s (about 40%).
Adding the chunked KL loss: memory drops further from 78GB (dense) to 58GB (fused); a single GPU trains at 32K context, roughly four times what the dense approach fits.
Isolated loss-kernel benchmark (4K to 256K): at 32K the chunked loss uses 5.45GB versus 85.2GB dense, a 15.6x reduction; at 256K it runs 3.3x faster than forward-chunked (0.630 vs 0.190 iterations/s).
Model quality (teacher Llama-3.1-8B, student about 3.2B):
| Loss config | MMLU | GSM8K |
| Mid-layer feature loss only | about 28% | about 4% |
| Logit KL only | 59.9% | 65.9% |
| Logit KL + hidden-state feature loss | 60.6% | 67.5% |
At under half the teacher's parameter count the student keeps most of its short-context ability, landing within about 9 points on MMLU (with larger gaps on WinoGrande and GSM8K).
For anyone doing model compression or private deployment, this is a recipe you can lift directly: offline caching pulls the teacher out of the training loop, and long-context distillation no longer requires stacking GPUs. Making hundreds of ablations affordable is a side benefit of that efficiency. The takeaway is simple: logit KL is indispensable, and stacking a hidden-state feature loss on top gives a small, reliable gain.
The authors list several: only one teacher-student pair (8B to 3.2B) is tested, so transfer to other architectures is unverified; the 256K numbers come from an isolated toy network rather than end-to-end training; results are obtained with Megatron-Bridge plus ModelOpt on H200, and other frameworks are unvalidated. There is also a gap of more than 9 points between student and teacher on MMLU and GSM8K, which shows this efficient recipe mostly solves "can we train it at all" rather than "can we match the teacher." The limits of the compression ratio are not answered here.