A KL loss tweak lets a single GPU distill a 120B model without melting

The dirty secret of LLM knowledge distillation isn’t the theory—it’s the electric bill. Training a smaller model to mimic a massive teacher like gpt-oss-120b typically requires keeping both models in VRAM simultaneously, generating a probability distribution over the entire vocabulary for every single token. For a 120-billion-parameter model with a 201,088-token vocabulary, a single training iteration can spike to roughly 250GB of VRAM. That’s not an engineering inconvenience; it’s a hard stop that pushes the work far beyond the reach of a single H200 or B200 GPU.

A new paper, “Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss,” attacks this problem not with better hardware, but with smarter math. The researchers propose two systemic changes. First, they pre-compute the teacher’s output just once, caching only the top 100 most likely tokens per position. The teacher is then shut off and never needs to sit in memory during the actual student training. This offline cache can be reused across countless experiments, which is a practical godsend for anyone iterating on data mixes or hyperparameters.

The second change is where the real magic happens, and it’s a bit counterintuitive. It targets the Kullback-Leibler divergence loss calculation itself. A naive, dense KL loss builds a monstrous grid in memory: vocabulary size by sequence length. Their fused chunked KL loss avoids materializing this entire grid. Instead, it keeps the teacher’s data sparse and computes the loss in manageable slices. The result is a memory profile that never forms the catastrophic 250GB spike, peaking instead around 128GB.

The speed implications are equally practical. Their forward-chunked method, which never expands the teacher’s sparse logits into a dense form, proved to be the fastest of the three mathematically equivalent approaches they benchmarked. This isn’t just an academic curiosity; it means long-context healing—a recovery process to fix a compressed model’s degraded ability to handle long sequences—becomes feasible on a single GPU. When companies like Nvidia and Multiverse Computing are regularly releasing distilled models like Nemotron-3-Puzzle-75B, making the distillation step itself cheap and accessible could shift who gets to participate in building the next generation of frontier models. It quietly challenges the assumption that you need a fleet of GPUs to do meaningful compression work.