Glonce

Multiverse Computing cuts LLM distillation VRAM by 15x with chunked KL loss

researchThe Glonce editorsHugging Face🔥71Glonce rating71Impact68Novelty68Relevance95Surprise45Players38Credibility87
Multiverse Computing cuts LLM distillation VRAM by 15x with chunked KL loss

Multiverse Computing published a paper and an accompanying Hugging Face blog post titled "Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss." Knowledge distillation trains a smaller student model to match a larger teacher's behavior, and it has become the standard way to compress today's largest open models: the post cites the Kimi-K3 model, which has 2.8 trillion parameters and needs roughly 3TB of VRAM just to load, and notes that Nvidia (Nemotron 3 Puzzle 75B) and Multiverse itself (Hypernova 60B) have both recently shipped compressed models built this way.

The standard, online version of distillation keeps both teacher and student loaded and has the teacher run a fresh forward pass at every training step, producing a full probability distribution over the vocabulary for every token. For gpt-oss-120b, whose vocabulary holds 201,088 tokens, a teacher-probability tensor at 32K sequence length and batch size 4 already has shape 4 x 201,088 x 32,768, about 50GB in bfloat16 for that one tensor alone. Add gradients, activations, weights and optimizer state, and a single training iteration can peak at roughly 250GB of VRAM, more than a single H200 (141GB) or B200 GPU provides.

Multiverse's fix has two parts. First, offline distillation: instead of recomputing the teacher every step, the team computes its output once, caches the top-100 most likely tokens per position, and trains the student against that cache, so the teacher never needs to sit in memory during training. Second, a fused, chunked KL-divergence loss. A naive KL loss has to build a full grid of one row per vocabulary entry and one column per sequence position before it can produce a single loss number; the fused version instead projects hidden states to logits one sequence chunk at a time, folds each chunk into the running loss, and discards it, recomputing chunks on the fly during the backward pass rather than storing them. The team compared this against a dense baseline and a simpler "forward-chunked" variant that keeps the teacher sparse but still materializes the student's full logits grid for the backward pass; the dense KL loss spikes to roughly 250GB against a single H200's 141GB capacity, while the fused chunked loss never forms that spike and peaks at about 128GB.

On a head-to-head run at 8K token context (Llama 3.1 8B Instruct as teacher, a 3.2B Llama model as student, single H200 GPU), all four setups, online distillation plus the three offline loss variants, reached near-identical training loss, which the authors say confirms offline distillation with top-100 cached logits is lossless relative to online distillation. At that shorter context the fused chunked loss was not the fastest of the three, since its extra backward-pass projection costs some speed; the authors state its advantage grows with context length instead.

A separate, isolated benchmark on a standalone output-projection network (no full transformer) demonstrates that scaling: at 32K tokens, peak memory falls from 85.2 GiB with the dense loss to 5.45 GiB with the fully chunked version, a 15.6x reduction, and the dense loss fails outright from 64K tokens onward. At 256K tokens the fully chunked loss uses 11.6 GiB against 134.2 GiB for the next-best chunked variant, and runs about 3.3x faster per iteration at that length. Applied to distilling a GPT-OSS 20B model at a 32,768-token context, the memory savings let the setup shrink from four GPU nodes to one; step time fell from 57.0 to 12.23 seconds, about 5x faster, and throughput per GPU rose from 74.2 to 345.7 TFLOP/s.

The efficiency gains are what made a larger distillation run affordable in the first place: the team distilled Llama 3.1 8B Instruct down to a roughly 3.2B-parameter student that retains most of the teacher's accuracy on BoolQ and HellaSwag and stays within about nine points of it on MMLU, at less than half the parameter count. Multiverse frames the work as part of an ongoing effort to make distillation and "healing" practical to iterate on rather than a one-off recipe, and says the paper covers further ablations on loss choice and sequence packing. The chunked-loss implementation is open-sourced at github.com/CompactifAI/Full-Chunked-KL-Loss.

Key facts

  • Multiverse Computing's paper combines offline top-100 logit caching (the teacher runs once, not every step) with a fused, chunked KL-divergence loss that never materializes the full vocabulary x sequence grid.
  • In a toy benchmark at 32K token context, peak memory drops from 85.2 GiB (dense loss) to 5.45 GiB (fully chunked), a 15.6x reduction; the dense loss fails outright from 64K tokens onward.
  • Distilling GPT-OSS 20B at a 32,768-token context, the fused loss let the setup shrink from four GPU nodes to one, cut step time from 57.0 to 12.23 seconds (about 5x), and raised per-GPU throughput from 74.2 to 345.7 TFLOP/s.
  • At 8K context, all four methods tested (online distillation and three offline loss variants) reached near-identical training loss, which the authors say shows offline distillation with top-100 cached logits is lossless relative to online.
  • The resulting distilled student, about 3.2B parameters from a Llama 3.1 8B Instruct teacher, stays within about nine points of the teacher on MMLU at less than half the parameter count; the chunked-loss code is open-sourced on GitHub.

Why it matters

Distillation is now the standard way labs compress huge open models like Kimi-K3 (2.8 trillion parameters, about 3TB of VRAM just to load) into deployable sizes, but the distillation step itself has needed hundreds of GPUs with careful tensor-parallelism, because keeping teacher and student loaded together and computing a full-vocabulary probability distribution per token is extremely memory-hungry. Multiverse Computing's two changes, caching the teacher's top-100 logits offline and fusing the KL loss computation into chunks so the full grid is never built, cut a single training iteration's peak VRAM from roughly 250GB to about 128GB, and in isolated long-context tests to a 15.6x reduction at 32K tokens and a working setup at 256K tokens where the dense approach fails outright past 64K.

Who it affects

Teams building or compressing large language models, particularly anyone distilling very large open teachers such as gpt-oss, Qwen, GLM or Kimi variants into smaller students, stand to benefit most directly. The GPT-OSS 20B example, where the fused loss shrank the required setup from four GPU nodes to one, is aimed at organizations without access to large multi-node clusters for this step.

How to use it

The chunked-loss implementation is open-sourced at github.com/CompactifAI/Full-Chunked-KL-Loss, and the full paper (linked from the Hugging Face post) covers the closed-form gradient behind the fused chunked loss plus the complete training configuration and further ablations on loss choice and sequence packing. Multiverse also invites teams to get in touch about applying the technique to their own distillation pipelines; no pricing or licensing terms for that engagement are stated in the post.

How solid is it

The claims rest on the team's own paper and blog post rather than independent replication. The core lossless-ness claim is backed by a controlled comparison: at 8K token context, online distillation and three offline loss variants (dense, forward-chunked, fused chunked) all converge to near-identical training loss on the same teacher-student pair (Llama 3.1 8B Instruct to a 3.2B Llama student) on a single H200 GPU. The larger memory and speed numbers come from two separate benchmarks, an isolated output-projection network used to test scaling to 256K tokens, and a full GPT-OSS 20B distillation run, so the headline 15.6x and 5x figures are drawn from different setups rather than one single test.

Risks and caveats

The post names no individual authors, crediting the work only to Multiverse Computing, and gives no publication date for the paper or the post. No dollar training costs are given, only VRAM, step-time and throughput figures, and the GPU type used for the GPT-OSS 20B four-nodes-to-one benchmark is not stated, unlike the H200 used elsewhere. No benchmark numbers are given for Nvidia's Nemotron 3 Puzzle 75B or Multiverse's own Hypernova 60B; they appear only as examples of compressed models, not as evaluations of this technique.