Distilling an LLM on One GPU: Offline Top-K Logits and a Fused Chunked KL Loss

Knowledge distillation, training a smaller student model to match a larger teacher, is back as a mainstream topic for a simple reason: the open-weight models everyone wants to deploy are enormous. Kimi K3 sits at 2.8 trillion parameters and needs roughly 3TB of VRAM just to load. Compressing these into smaller models, then recovering capability through distillation, is how teams like NVIDIA (Nemotron 3 Puzzle 75B) and Multiverse Computing (Hypernova 60B) ship usable open-weight releases.

The recovery step decides most of the final quality, and it is also the most expensive part of the pipeline. A new paper from Multiverse Computing's CompactifAI team (arXiv 2608.03796) attacks that cost with two systems changes, and the code is open-sourced. The headline: distilling a GPT-OSS 20B model at a 32,768-token context goes from four GPU nodes to one, with step time falling from 57.0 to 12.23 seconds, about 5x faster per step.

Why distillation recovery is expensive#

The standard setup is online distillation with a KL divergence loss: teacher and student sit in memory at the same time, the teacher runs a full forward pass every step, and the student trains against its full output distribution. It is the most expressive setup and the most memory-hungry one. The paper's concrete example: gpt-oss-120b has a vocabulary of 201,088 tokens. At a sequence length of 32K with batch size 4, the teacher probability tensor alone is 4 x 201,088 x 32,768, about 50GB in bfloat16 for a single tensor. Add gradients, activations, weights and optimizer states and a single training iteration peaks near 250GB of VRAM, more than an H200 or B200 can hold.

The two changes#

Offline distillation with cached top-K logits. Instead of recomputing the teacher every step, the paper computes its output once, caches the top-100 most likely tokens per position, and trains the student against that cache. The teacher never sits in memory during training and never runs again once the cache exists, and the same cache is reusable across hundreds of ablations. The surprising result is that this is lossless: at 8K context on a single H200, all four setups compared in the paper reach near-identical training loss, even though offline runs train against only 100 cached logits per position instead of a 201K-wide distribution. Offline distillation runs about 29% faster per iteration and reaches up to 41% higher throughput.

A fused, chunked KL loss. The second problem is the loss itself. Computed naively, KL divergence builds a grid of one row per vocabulary entry and one column per sequence position, which is enormous at 100K+ vocabularies and long sequences. The paper's fused chunked loss never produces the student's full logits grid at all: it processes one chunk of the sequence end to end, projects hidden states to logits for that chunk, folds the result into the running loss, and discards the chunk. The backward pass recomputes each chunk on the fly. The cost is running the output projection twice, but peak memory grows only linearly with sequence length instead of spiking with the full vocab-by-sequence size.

The numbers scale cleanly. On an isolated output-projection benchmark, peak memory at 32K tokens 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 at 64K tokens. At 256K tokens the chunked loss uses 11.6 GiB against 134.2 GiB for the next-best variant, and is about 3.3x faster per iteration.

What this changes in practice#

For the 8K single-GPU case the fused loss is not the fastest option, the extra backward-pass projection costs a bit of speed (20.2s per iteration vs 18.4s for forward-chunked). Its advantage appears as context grows. The GPT-OSS 20B distillation at 32,768-token context is the demonstration: the memory freed by the fused loss shrank the setup from four GPU nodes to one, step time fell from 57.0 to 12.23 seconds, and throughput per GPU rose from 74.2 to 345.7 TFLOP/s.

The resulting student in the paper, distilled from Llama 3.1 8B Instruct down to about 3.2B parameters, keeps most of the teacher's accuracy on BoolQ and HellaSwag and stays within about nine points on MMLU, at less than half the parameter count. The recovery quality bar is not the point; the point is that this quality can now be reached on a single GPU and iterated on cheaply.

Why this matters to developers#

Three takeaways:

1. Distillation is now a systems problem, not an algorithm problem. The two wins here are both about memory layout and when to materialize tensors, not about a new loss function or a cleverer objective. The paper's own framing is a practitioner's study of training efficiency. That is the direction the field is moving: the open-weight ecosystem has the models, and the bottleneck is who can afford the expensive recovery step.

2. Full-vocabulary distributions are overkill for distillation. Offline training against cached top-100 logits matches online training loss curves almost exactly. That is a practical license to build cheaper distillation pipelines: one teacher pass, a small cache per position, and the teacher hardware is free for other work. For teams that already run open-weight serving (see our break-even math for self-hosting), the cache also survives across experiments, so ablation runs stop multiplying teacher passes.

3. Long-context distillation on commodity hardware changes the economics of local models. The 5x step-time win at 32K context is exactly the regime that matters for agentic and RAG workloads, and it lands on one GPU instead of a four-node cluster. Combined with the quantization and serving work we have covered on GLM 5.2 and AMD MI355X, the path to a capable small model is becoming: distill on one GPU, quantize, serve locally.

The paper is early and the evaluation is limited, the accuracy recovery study is on short-context benchmarks, and the long-context numbers come from an output-head-only toy benchmark rather than a full training run. But the open-source implementation (github.com/CompactifAI/Full-Chunked-KL-Loss) means any team with a PyTorch stack can try the fused chunked loss against their own pipeline today. Expect offline distillation with cached logits to become the default recipe in open-weight training guides within a quarter.

Continue Reading#

• Where to Access Kimi K3 and What It Costs - the 2.8T open-weight model that makes distillation recovery necessary
• GLM 5.2 Cost Math: Open-Weights Coding Models - what running large open-weight models actually costs
• Self-Hosting Open-Weights Models: The Break-Even Math - when local serving beats API calls
• The $500 RL Fine-Tune That Beats Frontier Models - what small-scale training on commodity hardware can achieve
• AMD MI355X vs NVIDIA B200/B300 for Open-Weights Serving - the hardware side of the same cost question

Sources#

• Multiverse Computing blog post: Making Knowledge Distillation Cheap Enough to Run at Scale (Hugging Face blog, August 10, 2026)
• Paper: Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss - arXiv 2608.03796 (fetched August 10, 2026)
• Code: CompactifAI/Full-Chunked-KL-Loss on GitHub