让知识蒸馏便宜到足以大规模运行

Why distillation recovery is expensive

The standard setup, online distillation using the Kullback-Leibler divergence loss (KL loss), keeps both the teacher and the student loaded at the same time. At every training step, the teacher runs a full forward pass to produce its output distribution, and the student is trained to match it. This is the most expressive setup, since the full teacher distribution is available, but it is also the most memory- and compute-intensive: two full-vocabulary tensors have to be held per token position, and the teacher has to be recomputed on every single step even though its behavior does not change across a training run.

As a practical example, gpt-oss-120b 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's already about 50GB of VRAM for a single tensor. Add gradients, activations, model weights, and optimizer states, and a single training iteration of distillation can peak at roughly 250GB of VRAM, more than even an H200 or B200 GPU can provide. In this post, we show that reformulating the KL loss to process the data in chunks reduces this cost to almost nothing.

Dense KL spikes to roughly 250GB, above a single H200's 141GB capacity. The fused chunked loss never forms that spike and peaks at about 128GB. Source: paper Figure 1.

Two systems changes

Offline distillation. Instead of recomputing the teacher at every step, we compute its output once, cache the top-100 most likely tokens per position, and train the student against that cache. The teacher never has to sit in memory during training and does not need to be run again once the cache exists, so the same cache can be reused across many ablations.

A fused, chunked KL loss. To see why the loss itself is expensive, picture what it actually builds: for every token position in a sequence and every word in the vocabulary, the loss needs a number describing how much the student's prediction disagrees with the teacher's. Laid out as a grid, that's one row per vocabulary entry and one column per sequence position, for a vocabulary of 100K+ words and a long sequence, that grid is enormous, and the default way of computing a KL loss builds the whole thing before it can produce a single number.

We compare three ways of computing this same loss, all mathematically equivalent:

Dense KL is the textbook approach. It rebuilds a full, dense teacher-probability grid from the cached top-100 logits and compares it against the student's own dense grid of log-probabilities. This is the version closest to how online distillation already works, so we use it as our correctness baseline, but it holds the full vocabulary × sequence grid in memory, twice over.

Forward-chunked KL keeps the teacher sparse (only its cached top-100 logits per position, never expanded into a dense grid) and computes the loss piece by piece, one slice of sequence positions at a time. This removes the dense teacher and the dense comparison, and turns out to be the fastest of the three methods in our benchmarks. It still has one blind spot, though: the student's own logits, the grid produced by the model's output layer, are still computed in full and kept around for the backward pass, so memory still grows steeply with sequence length.

Fused chunked KL , our main contribution, goes a step further and fuses the model's output projection directly into the loss computation. It never produces the student's full logits grid at all: it processes one chunk of the sequence at a time end to end, projecting hidden states to logits for that chunk, folding the result into the running loss, and discarding the chunk before moving to the next one. The backward pass recomputes each chunk on the fly instead of storing it. The cost is doing that projection twice, once forward, once in backward, but in exchange, peak memory grows only linearly with sequence length instead of spiking with the full vocabulary × sequence size.

The GIF below shows the difference between the dense and fused-chunked approaches: one builds the whole comparison grid and holds onto all of it, the other builds and discards one slice at a time, so memory never grows beyond a single chunk.

We have open-sourced the chunked-loss implementation: github.com/CompactifAI/Full-Chunked-KL-Loss

What this changes in practice

The table below puts all four setups head to head: online distillation, and the three offline loss implementations just described. Comparing them on a single H200 GPU with Llama 3.1 8B Instruct as teacher and a 3.2B Llama model as student at an 8K token context, all four reach near-identical training loss, even though the offline runs train against only the cached top-100 logits per token.

The loss curves overlap almost exactly across all four methods, confirming offline distillation with top-100 cached logits is lossless relative to online distillation. Source: paper Figure 2. At this sequence length, the fused chunked loss is not yet the fastest option, its extra backward-pass projection costs a bit of speed, but its real advantage only shows up as context length grows, which the next section demonstrates.

Scaling to long context lengths

To see the scaling pattern more starkly, we ran an isolated benchmark on a toy output-projection network (no transformer body, just the loss kernel). 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, and the dense loss fails outright from 64K tokens onward. At 256K tokens, the fully chunked loss uses 11.6 GiB against 134.2 GiB for the next-best chunked variant, and is about 3.3× faster per iteration at that length.

Distilling a GPT-OSS 20B model at a 32,768-token context, the memory freed by the fused loss let the setup shrink from four GPU nodes down to one. Step time fell from 57.0 to 12.23 seconds, about 5× faster, and throughput per GPU rose from 74.2 to 345.7 TFLOP/s.

The resulting student

The efficient offline setup is what made a large-scale distillation campaign affordable in the first place. The resulting compact student, distilled from Llama 3.1 8B Instruct down to about 3.2B parameters, retains most of the teacher's accuracy on BoolQ and HellaSwag, stays within about nine points of it on MMLU, at less than half the parameter count.

The student retains most of the teacher's short-context accuracy at less than half the size. Source: paper Figure 6.

This work is part of Multiverse Computing's ongoing research into making distillation and healing practical to run at scale , not just as a one-off recipe, but as something teams can iterate on cheaply. The paper also covers additional ablations, such as how the choice of loss function and sequence packing affect recovery quality.

We have also open-sourced the chunked-loss implementation: github.com/CompactifAI/Full-Chunked-KL-Loss

Models mentioned in this article 6

MultiverseComputingCAI/Hypernova-60B-2605

Text Generation • 59B • Updated Aug 7 • 1.71k • 21 meta-llama/Llama-3.1-8B-Instruct Text Generation • 8B • Updated Sep 25, 2024 • 6.2M • 8.13k moonshotai/Kimi-K3 Image-Text-to-Text • 2.8T • Updated 30 days ago • 1.27M • 11.6k nvidia/NVIDIA-Nemotron-Labs-3-Puzzle-75B-A9B-NVFP4 Text Generation • 45B • Updated Jul 7 • 10k • 130 openai/gpt-oss-120b Text Generation • 117B • Updated Aug 26, 2025 • 4.48M • 5.34k zai-org/GLM-5.2 Text Generation • 753B • Updated about 1 month ago • 731k • 5.14k Papers mentioned in this article 1 Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss Paper • 2608.03796 • Published Aug 4 • 19 Collections mentioned in this article 1 Qwen3.5 Collection 21 items • Updated Mar 9 • 1.79k More from this author

15 September 29, 2026 Pruning LLMs Like a Physicist: Block Removal as an Ising Optimization Problem 33 September 21, 2026 Community Edit Preview Upload images, audio, and videos by dragging in the text input, pasting, or clicking here . Tap or paste here to upload images Comment · Sign up or log in to comment

+30

Models mentioned in this article 6

MultiverseComputingCAI/Hypernova-60B-2605

meta-llama/Llama-3.1-8B-Instruct

Papers mentioned in this article 1

Collections mentioned in this article 1

Qwen3.5