SciTech Pulse
AI

New distillation trick cuts training memory below single-GPU limit for large language models

Multiverse Computing has published a paper that cuts the GPU memory needed to distil a large language model into a smaller one below the ceiling of a single H200 GPU, using two systems-level changes to the distillation l

Multiverse Computing has published a paper titled 'Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss', presenting two changes that cut the GPU memory needed to train a smaller student model to match a larger teacher model. The company says the changes make long-context distillation feasible on a single GPU and cheap enough for large-scale experimentation, where the standard approach typically requires hundreds of GPUs and careful tensor-parallelism strategies.

Knowledge distillation trains a smaller model to reproduce the output distribution of a larger one, a standard way to shrink very large language models without losing capability. Recent examples include NVIDIA's Nemotron 3 Puzzle 75B and Multiverse Computing's own Hypernova 60B. In the usual online setup, both the teacher and the student are loaded together and the teacher runs a full forward pass at every training step, producing two full-vocabulary probability tensors per token position — the most expensive part of the pipeline.

For a teacher with the vocabulary of gpt-oss-120b (201,088 tokens), at a sequence length of 32,000 and a batch size of 4, the teacher-probability tensor alone occupies about 50 GB of memory in bfloat16. Add gradients, activations, model weights and optimizer states, and a single training iteration can peak near 250 GB — above the 141 GB capacity of an H200 GPU.

The first change is offline distillation: the teacher is run once and its top-100 most-likely tokens per position are cached, so the student trains against that cache without the teacher ever having to sit in memory alongside it. The second is a fused, chunked KL-divergence loss. Instead of building the full vocabulary-by-sequence comparison grid at once, the loss is computed one slice of sequence positions at a time, and the model's output projection is folded directly into the loss so the student's own full logits grid is never materialised.

In the paper's benchmarks, the fused chunked loss peaks at about 128 GB — inside a single H200's memory — where the standard dense KL implementation spikes to roughly 250 GB. The backward pass recomputes each chunk on the fly rather than storing it, doubling the projection work but growing memory only linearly with sequence length.

#multiverse computing#knowledge distillation#llm#efficient training#memory optimization
Rate this story

Related stories