Knowledge distillation made cheap enough for single GPU

Knowledge distillation is a well-established machine learning technique in which a smaller model (the "student") is trained to match the behavior of a larger, more capable model (the "teacher"). With the recent wave of open-source large language models—including gpt-oss, Qwen, GLM, and Kimi—this approach has become mainstream research again. The motivation is clear: deploying very large models is expensive. The recent Kimi-K3 model has 2.8 trillion parameters and requires roughly 3TB of VRAM just to load. Companies like Nvidia (Nemotron 3 Puzzle 75B) and Multiverse Computing (Hypernova 60B) have released high-quality compressed models through distillation, making it standard practice for production deployment.

However, the distillation step itself has remained the most expensive part of the pipeline, and in many cases the limiting factor for experimentation. The standard approach, called online distillation using Kullback-Leibler divergence loss (KL loss), keeps both teacher and student models in memory at the same time. At each training step, the teacher runs a full forward pass to produce its output distribution across the entire vocabulary, and the student is trained to match it. This is highly expressive—the full teacher distribution is available—but also extremely memory-intensive. As a concrete example, the gpt-oss-120b model has a vocabulary of 201,088 tokens. At a sequence length of 32K and batch size 4, the teacher-probability tensor alone has shape 4 × 201,088 × 32,768; in bfloat16, that single tensor requires about 50GB of VRAM. Adding gradients, activations, model weights, and optimizer states, a single training iteration can peak at roughly 250GB—more than even an H200 or B200 GPU can provide. This has made large-scale distillation experimentation feasible only for teams with access to hundreds of GPUs and the expertise to manage complex tensor-parallelism strategies.

Researchers at Multiverse Computing, in a paper titled "Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss," propose two systems-level changes to address this bottleneck. The first is offline distillation: instead of recomputing the teacher at every training step, the teacher is run once, and its output—specifically, the top-100 most likely tokens per position—is cached. The student is then trained against this cache. The teacher never needs to sit in memory during training and does not need to be run again once the cache exists. The same cache can be reused across many ablations and experiments, eliminating redundant computation.

The second innovation is a fused, chunked KL loss. The standard KL loss requires materializing a full grid where each row represents a vocabulary entry and each column represents a sequence position. For a 100K+ word vocabulary and a long sequence, this grid is enormous, and the default implementations in libraries like PyTorch and NVIDIA Megatron-Bridge build the entire grid before producing a single loss number. The new approach reformulates this computation in three progressive stages: Dense KL (the baseline, still building the full grid for validation), Forward-chunked KL (keeping the teacher sparse and processing in slices of the sequence), and Fused chunked KL (the main contribution, which fuses the model's output projection directly into the loss computation, processing one chunk of the sequence end-to-end, folding the result into the running loss, and discarding the chunk before moving to the next). The backward pass recomputes each chunk on the fly instead of storing it, doing the projection twice but ensuring peak memory grows only linearly with sequence length.

On a single H200 GPU at 8K token context, using Llama 3.1 8B Instruct as teacher and a 3.2B Llama model as student, the methods show these characteristics: Online distillation peaks at 102.8 GB with 237 TFLOP/s throughput; Offline Dense KL drops to 78.3 GB with 331 TFLOP/s; Forward-chunked KL reaches 61.8 GB with 335 TFLOP/s; and Fused chunked KL uses 58.3 GB with 304 TFLOP/s. All four methods achieve nearly identical training loss, confirming that offline distillation with cached top-100 logits is lossless relative to online distillation. At this context length, the fused chunked loss is not yet fastest due to its extra backward-pass projection cost, but its real advantage emerges as context length grows.

At longer sequences, the gains are dramatic. In an isolated benchmark on just the output-projection network at 32K tokens, peak memory falls from 85.2 GiB with the dense loss to 5.45 GiB with the fully chunked version—a 15.6× reduction. The dense loss fails outright from 64K tokens onward. At 256K tokens, the fully chunked loss uses 11.6 GiB versus 134.2 GiB for the next-best variant and is about 3.3× faster per iteration at that length. When distilling a GPT-OSS 20B model at 32,768-token context, the memory savings allow the setup to shrink from four GPU nodes down to one. Step time falls from 57.0 to 12.23 seconds—about 5× faster—and throughput per GPU rises from 74.2 to 345.7 TFLOP/s.

The resulting student model, a 3.2B parameter model distilled from Llama 3.1 8B Instruct, retains most of the teacher's accuracy on BoolQ and HellaSwag benchmarks and stays within about nine points of it on MMLU, achieving this at less than half the parameter count. The researchers have open-sourced the chunked-loss implementation at github.com/CompactifAI/Full-Chunked-KL-Loss and describe this work as part of a broader effort to make distillation and model healing practical to run at scale—not as a one-off recipe requiring specialized infrastructure, but as something teams can iterate on cheaply.