Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and Fused Chunked KL Loss

Multiverse Computing has introduced a new method for knowledge distillation that significantly reduces VRAM usage and training costs, enabling long-context distillation on a single GPU. By combining offline logit caching with a fused chunked KL-divergence loss, the researchers can distill large teacher models into smaller students without the massive memory overhead typically associated with full-vocabulary probability distributions.

The High Cost of Standard Knowledge Distillation

Traditional "online" knowledge distillation using Kullback-Leibler (KL) divergence loss is computationally expensive because it requires both the teacher and student models to be loaded into VRAM simultaneously. At every training step, the teacher must perform a full forward pass to generate a probability distribution over the entire vocabulary for every token position.

For large models, this creates a massive memory bottleneck. For example, using a model like gpt-oss-120b with a vocabulary of 201,088 tokens, a sequence length of 32K, and a batch size of 4, the teacher-probability tensor alone requires approximately 50GB of VRAM in bfloat16. When including gradients, activations, and optimizer states, a single iteration can peak at roughly 250GB of VRAM, exceeding the capacity of even high-end GPUs like the H200 or B200.

Two Systems Changes for Scalable Distillation

To solve these memory constraints, Multiverse Computing implemented two primary technical changes:

1. Offline Distillation via Top-K Logit Caching

Instead of recomputing teacher outputs at every step, the system computes the teacher's output once and caches the top-100 most likely tokens per position. The student is then trained against this cache. This removes the need to keep the teacher model in memory during the training process and allows the same cache to be reused across different experimental ablations.

2. Fused Chunked KL Loss

Standard KL loss implementations build a full vocabulary-by-sequence grid to calculate disagreement between the student and teacher. Multiverse Computing's "fused chunked KL loss" optimizes this by fusing the model's output projection directly into the loss computation.

Rather than materializing the full logits grid, the process works as follows:

• It processes one chunk of the sequence at a time.
• It projects hidden states to logits only for that specific chunk.
• It folds the result into the running loss and immediately discards the chunk.
• The backward pass recomputes each chunk on the fly.

This approach ensures that peak memory grows linearly with sequence length rather than spiking based on the full vocabulary size.

Performance and Memory Benchmarks

Comparing different loss implementations on a single H200 GPU using Llama 3.1 8B Instruct (teacher) and a 3.2B Llama model (student) at an 8K token context, the researchers found that all methods reached near-identical training loss. This confirms that offline distillation using only the top-100 cached logits is effectively lossless relative to online distillation.

Method (8K context, single H200)	Peak memory	Iteration time	Throughput
Online distillation	102.8 GB	25.9 s	237 TFLOP/s
Offline, dense KL	78.3 GB	18.5 s	331 TFLOP/s
Offline, forward-chunked KL	61.8 GB	18.4 s	335 TFLOP/s
Offline, fused chunked KL	58.3 GB	20.2 s	304 TFLOP/s

Scaling to Long Contexts

The advantages of the fused chunked loss become more pronounced as context length increases. In isolated benchmarks using a toy output-projection network:

• At 32K tokens: Peak memory dropped from 85.2 GiB (dense loss) to 5.45 GiB (fully chunked), a 15.6× reduction.
• At 64K tokens: The dense loss implementation failed entirely.
• At 256K tokens: The fully chunked loss used 11.6 GiB, compared to 134.2 GiB for the next-best chunked variant, and was 3.3× faster per iteration.

In a real-world scenario distilling a GPT-OSS 20B model at a 32,768-token context, the fused loss reduced the required hardware from four GPU nodes down to one. This resulted in a 5× speedup in step time (from 57.0 to 12.23 seconds) and an increase in throughput per GPU from 74.2 to 345.7 TFLOP/s.

Student Model Capabilities

The efficiency of this pipeline allowed for a large-scale distillation campaign that produced a 3.2B parameter student model distilled from Llama 3.1 8B Instruct. This compact student retains most of the teacher's accuracy on BoolQ and HellaSwag, and remains within approximately nine points of the teacher on the MMLU benchmark, despite having less than half the parameter count.