Multiverse Computing Cuts LLM Distillation Memory 15.6 Times With Fused KL Loss
Multiverse Computing cuts LLM distillation memory using cached top-100 logits and fused KL loss, putting 32K GPT-OSS 20B training on a single H200 GPU.
Summary
On August 10, 2026, Hugging Face published Multiverse Computing researchers Antonio Tiene, Iker García-Ferrero, Ali Hashemi and Bakbergen Ryskulov’s paper, “Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss.” It targets open gpt-oss, Qwen, GLM and Kimi models, including 2.8 trillion-parameter Kimi-K3, which needs roughly 3TB VRAM, and compression work such as Nvidia’s Nemotron 3 Puzzle 75B and Multiverse’s Hypernova 60B.
The method caches each teacher’s top-100 logits once, then fuses output projection into chunked KL, removing the resident teacher and full vocabulary by sequence matrices created by default PyTorch and NVIDIA Megatron-Bridge implementations. With gpt-oss-120b’s 201,088-token vocabulary, 32K sequence and batch 4, one bfloat16 teacher tensor occupies about 50GB; dense distillation peaks near 250GB, beyond an H200 or B200, while fused peaks about 128GB against H200’s 141GB. At 8K on one H200, Llama 3.1 8B Instruct teaching a 3.2B Llama delivered near-identical loss across online, offline dense, forward-chunked and fused methods: respectively 102.8, 78.3, 61.8 and 58.3GB; 25.9, 18.5, 18.4 and 20.2 seconds; and 237, 331, 335 and 304 TFLOP/s.
In a projection-only test at 32K, fused cut 85.2 GiB to 5.45 GiB, 15.6 times; dense failed from 64K. At 256K, fused used 11.6 GiB versus 134.2 GiB for forward chunking and ran about 3.3 times faster. GPT-OSS 20B distillation at 32,768 tokens shrank from four GPU nodes to one, cutting steps from 57.0 to 12.23 seconds, about 5 times, while per-GPU throughput rose from 74.2 to 345.7 TFLOP/s. The approximately 3.2B student kept most BoolQ and HellaSwag accuracy and trailed its teacher by about nine MMLU points at under half the parameters. Code is open at github.com/CompactifAI/Full-Chunked-KL-Loss; the paper also tests loss choice and sequence packing.
Positives
- Cached top-100 teacher logits produced near-identical 8K training loss to online distillation while eliminating repeated teacher inference.
- Fused chunked KL reduced 32K projection-loss memory from 85.2 GiB to 5.45 GiB, a 15.6-fold improvement.
- GPT-OSS 20B training at 32,768 tokens fell from four GPU nodes to one and became about five times faster.
- Per-GPU throughput for GPT-OSS 20B rose from 74.2 to 345.7 TFLOP/s after adopting the fused loss.
- The 3.2B student retained most BoolQ and HellaSwag accuracy at less than half its teacher’s parameter count.
- Multiverse Computing released the chunked-loss implementation through the CompactifAI GitHub repository.
Risks & concerns
- Dense distillation can peak near 250GB, exceeding the memory capacity of a single H200 or B200 GPU.
- Fused chunked KL took 20.2 seconds per 8K iteration, slower than forward-chunked KL at 18.4 seconds.
- Backward propagation recomputes each output-projection chunk, requiring the projection to run twice.
- Dense KL failed from 64K tokens onward in the projection-only benchmark.
- The 3.2B student remained about nine points behind Llama 3.1 8B Instruct on MMLU.