Multiverse Computing introduces a memory-efficient approach to large language model knowledge distillation by combining offline top-K logits caching with a fused chunked Kullback-Leibler (KL) divergence loss. This method eliminates the need to keep both teacher and student models in memory simultaneously, drastically reducing VRAM requirements.
- The system caches the teacher's top-100 logits once, allowing the teacher to be removed from memory during training.
- A fused chunked KL loss processes sequences in chunks, avoiding the materialization of full vocabulary-sequence matrices.
- Peak VRAM for a 32K token context drops from 85.2 GiB to 5.45 GiB, a 15.6x reduction.
- Distilling GPT-OSS 20B at 32,768 tokens reduces hardware needs from four GPU nodes to one and increases throughput to 345.7 TFLOP/s.
These changes make long-context distillation practical on a single GPU, enabling affordable large-scale experimentation and model compression.