Multiverse Computing introduced a training system that reduces the memory required for knowledge distillation in large language models, according to a report published on Hugging Face. The technique allowed researchers distilling a 20-billion parameter model at a 32,768-token context length to collapse their hardware setup from four GPU nodes down to a single GPU node.
Standard online distillation requires running both the original teacher model and the smaller student model simultaneously. For a model such as gpt-oss-120b with a vocabulary of 201,088 tokens at a 32,000-token sequence length, the teacher probability tensor consumes roughly 50 gigabytes of VRAM, with overall training memory peaking near 250 gigabytes per iteration.
Offline Caching and Fused Losses
To eliminate teacher memory overhead during training, Multiverse Computing cached the teacher model's top-100 most likely tokens per position to disk. The student model then trains directly against this cache without requiring the teacher model to remain loaded in memory.
The researchers also created a fused chunked Kullback-Leibler loss implementation, which they open-sourced on GitHub under the CompactifAI repository. Instead of generating a full vocabulary-by-sequence matrix, the system projects hidden states to logits one slice of the sequence at a time, folds the result into a running loss total, and discards the chunk before moving to the next. The backward pass recomputes each chunk on the fly.
Scaling to Long Contexts
In benchmarks running Llama 3.1 8B Instruct as a teacher and a 3.2B Llama model as a student at an 8,000-token context on a single Nvidia H200 GPU, peak VRAM usage dropped from 102.8 gigabytes under online distillation to 58.3 gigabytes with the fused chunked loss. Step time on the 20-billion parameter distillation test fell from 57.0 to 12.23 seconds, while throughput per GPU increased from 74.2 to 345.7 TFLOP/s.
In isolated benchmarks on an output-projection network, dense loss calculations failed at context lengths of 64,000 tokens and above. The fully chunked loss kernel completed a 256,000-token context test using 11.6 gibibytes of peak memory compared to 134.2 gibibytes for the next-best chunked variant.
Tests distilling Llama 3.1 8B Instruct down to approximately 3.2 billion parameters showed that the compressed student model retained most short-context accuracy on BoolQ and HellaSwag benchmarks, while remaining within roughly nine points of the teacher on MMLU.
