Efficient Offline Top-K Logits and Fused Chunked KL Loss Enable Scalable LLM Distillation

Efficient Offline Top-K Logits and Fused Chunked KL Loss Enable Scalable LLM Distillation Multiverse Computing has introduced a memory-efficient knowledge distillation framework that reduces VRAM requirements by caching the teacher model's top-K logits offline and utilizing a fused, chunked Kullback-Leibler (KL) divergence loss. The approach eliminates the need to keep both teacher and student models in memory simultaneously and avoids materializing full vocabulary-by-sequence tensor grids. Online distillation causes massive memory spikes, such as roughly 250GB VRAM for a 120B model at 32K context, exceeding single H200 or B200 capacities. Caching the teacher's top-100 logits per token offline allows the teacher to be excluded from training memory and enables cache reuse across experiments. The fused chunked KL loss computes the output projection and loss slice-by-slice, recomputing chunks during the backward pass to prevent memory spikes. In isolated loss kernel benchmarks at a 32K context, peak memory dropped from 85.2 GiB to 5.45 GiB (a 15.6× reduction), scaling up to 256K tokens where dense methods fail. Distilling a GPT-OSS 20B model at 32K context reduced hardware requirements from four GPU nodes to one, increasing per-GPU throughput from 74.2 to 345.7 TFLOP/s. The Memory Bottleneck of Online Distillation Standard online knowledge distillation requires keeping both the teacher and student models loaded in memory while the teacher executes a full forward pass at every step. Because LLMs feature expansive vocabularies—such as gpt-oss-120b with 201,088 tokens—materializing probabilities over long sequences generates massive intermediate tensors. At a sequence length of 32K with a batch size of 4, the teacher's probability tensor alone consumes about 50GB of VRAM in bfloat16, with full training iterations spiking to roughly 250GB once gradients, optimizer states, activations, and weights are included. Offline Caching and Fused Chunked KL Architecture The proposed method introduces two systems-level adjustments. First, offline distillation precomputes the teacher's outputs once and caches only the top-100 most likely logits per position, entirely removing the teacher from VRAM during training. Second, the authors designed a fused chunked KL loss to replace the standard dense calculation. Unlike dense KL (which reconstructs full vocabulary grids) or forward-chunked KL (which avoids dense teacher grids but still keeps student logits for backward passes), the fused chunked method integrates the model's output projection into the loss kernel. It projects hidden states to logits chunk by chunk, accumulates the running loss, discards each chunk, and recomputes the projection on the fly during the backward pass. Empirical Scaling and Context Length Benchmarks Across an 8K-token context on a single H200 GPU distilling Llama 3.1 8B Instruct into a 3.2B student, all offline methods showed nearly identical training loss curves to online distillation. Offline dense KL peaked at 78.3 GB, forward-chunked KL at 61.8 GB, and fused chunked KL at 58.3 GB, compared to 102.8 GB for online distillation. Memory savings scale significantly as context lengths increase: in an isolated kernel benchmark at 32K tokens, fused chunked loss reduced peak memory by 15.6× compared to dense KL (5.45 GiB versus 85.2 GiB). At 256K tokens, where dense KL fails entirely, the fused loss operated at 11.6 GiB with a 3.3× speed advantage over other chunked variants. For a GPT-OSS 20B distillation at 32K context, the technique reduced infrastructure needs from four nodes to one and cut step time from 57.0 to 12.23 seconds. Model Quality and Open-Source Release Evaluating the distilled 3.2B parameter student derived from Llama 3.1 8B Instruct showed strong quality retention despite having less than half the parameters. The model preserved most of the teacher's short-context accuracy on BoolQ and HellaSwag, while remaining within roughly nine points on MMLU. The implementation has been released openly in the CompactifAI Full-Chunked-KL-Loss repository on GitHub.

降低大模型知识蒸馏成本：离线Top-K与融合分块KL损失技术方案 Multiverse Computing团队针对大语言模型在知识蒸馏中显存占用极高、难以低成本扩展的问题，提出了结合离线Top-K Logits缓存与融合分块KL损失（Fused Chunked KL Loss）的优化系统，大幅削减显存峰值并提升长上下文训练效率。 在线蒸馏需同时常驻教师与学生模型并计算全词表矩阵，导致显存极易突破单卡极限（如32K上下文下可峰值达250GB）。 采用离线蒸馏策略，单次预计算并缓存教师模型的Top-100 Logits，使教师模型无需在训练中常驻显存且缓存可复用。 融合分块KL损失将模型输出投影直接融合到损失计算中，按序列分块处理并在反向传播时重算投影，避免实例化完整的词表×序列长度矩阵。 在32K上下文的玩具测试中，融合分块损失使显存峰值从85.2 GiB降至5.45 GiB；在GPT-OSS 20B模型蒸馏中将集群需求从4节点压减至单节点。 传统蒸馏的显存瓶颈与离线Top-K设计 标准的在线知识蒸馏需要在每个训练步同时加载教师和学生模型，并对词表上的完整概率分布计算KL散度。以gpt-oss-120b为例，在32K序列长度和批次大小为4时，单是教师概率张量在bfloat16精度下就占用约50GB显存，结合激活值与优化器状态后峰值可达250GB，远超单个H200或B200 GPU容量。研究人员通过离线蒸馏解耦两者：预先计算教师输出并仅缓存每位置最可能的Top-100 Logits，教师模型在训练时无需加载，且实验证明该缓存方式训练损失几乎无损于在线全分布蒸馏。 融合分块KL损失的技术机制 为了避免在内存中生成庞大的“词表大小×序列长度”全量网格，研究对比了三种KL损失实现。稠密KL和前向分块KL仍会保留学生模型的完整Logits网格，显存随序列长度急剧增加。核心贡献“融合分块KL损失”将输出投影层与损失计算直接融合，逐块将隐藏状态投影为Logits、折叠计入运行损失后立刻丢弃该块，在反向传播中按需重新计算投影。虽然增加了计算投影的开销，但使峰值显存随序列长度仅呈线性增长，消除了显存尖峰。 长上下文扩展性能与模型恢复成果 在独立输出投影基准中，32K上下文下融合分块损失相较稠密损失实现显存降低15.6倍（5.45 GiB对85.2 GiB），并在256K长上下文下保持仅11.6 GiB的显存使用与3.3倍的速度优势。在32,768上下文蒸馏GPT-OSS 20B模型时，硬件环境从4个GPU节点精简至单节点，单步时间从57.0秒降至12.23秒，单卡算力利用率由74.2提升至345.7 TFLOP/s。基于该方案，从Llama 3.1 8B Instruct蒸馏出的3.2B学生模型在保留大部分短上下文能力的同时，体积减少了一半以上。相关实现已在GitHub开源。

Knowledge distillation , training a smaller student model to match the performance of a larger teacher, is a well-known technique in Machine Learning. With the recent wave of open-source Large Language Models, such as gpt-oss , Qwen , GLM , or Kimi , it has become a mainstream research topic again. Deploying these very large models is expensive: the recent Kimi-K3 model has 2.8 trillion parameters and needs roughly 3TB of VRAM just to load. Compressing them into smaller models and recovering the original capabilities through knowledge distillation has therefore become standard practice, with companies like Nvidia ( Nemotron 3 Puzzle 75B ) or Multiverse Computing ( Hypernova 60B ) recently releasing high-quality compressed models. The distillation step is what decides most of the final quality, but it's also usually the most expensive part of the pipeline. Keeping both the teacher and student loaded, and producing a probability distribution over the entire vocabulary for every token, requires enormous amounts of VRAM, typically feasible only with hundreds of GPUs and careful tensor-parallelism strategies. Our latest paper, Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss , tackles this with two systems changes: caching the teacher's top-K logits once so the teacher never has to sit in memory alongside the student, and a new, memory-efficient KL-divergence loss that avoids ever materializing the full vocabulary-size × sequence-length matrix, cutting VRAM use far below what the default implementations in libraries like PyTorch or NVIDIA Megatron-Bridge achieve. Together, these two changes cut training cost enough to make long-context healing possible on a single GPU, and cheap enough to make large-scale experimentation practical. Why distillation recovery is expensive The standard setup, online distillation using the Kullback-Leibler divergence loss (KL loss), keeps both the teacher and the student loaded at the same time. At every training step, the teacher runs a full forward pass to produce its output distribution, and the student is trained to match it. This is the most expressive setup, since the full teacher distribution is available, but it is also the most memory- and compute-intensive: two full-vocabulary tensors have to be held per token position, and the teacher has to be recomputed on every single step even though its behavior does not change across a training run. As a practical example, gpt-oss-120b has a vocabulary of 201,088 tokens. At a sequence length of 32K and batch size 4, the teacher-probability tensor alone has shape 4 × 201,088 × 32,768 ; in bfloat16, that's already about 50GB of VRAM for a single tensor. Add gradients, activations, model weights, and optimizer states, and a single training iteration of distillation can peak at roughly 250GB of VRAM, more than even an H200 or B200 GPU can provide. In this post, we show that reformulating the KL loss to process the data in chunks reduces this cost to almost nothing. Dense KL spikes to roughly 250GB, above a single H200's 141GB capacity. The fused chunked loss never forms that spike and peaks at about 128GB. Source: paper Figure 1. Two systems changes Offline distillation. Instead of recomputing the teacher at every step, we compute its output once, cache the top-100 most likely tokens per position, and train the student against that cache. The teacher never has to sit in memory during training and does not need to be run again once the cache exists, so the same cache can be reused across many ablations. A fused, chunked KL loss. To see why the loss itself is expensive, picture what it actually builds: for every token position in a sequence and every word in the vocabulary, the loss needs a number describing how much the student's prediction disagrees with the teacher's. Laid out as a grid, that's one row per vocabulary entry and one column per sequence position, for a vocabulary of 100K+ words and a long sequence, that grid is enormous, and the default way of computing a KL loss builds the whole thing before it can produce a single number. We compare three ways of computing this same loss, all mathematically equivalent: Dense KL is the textbook approach. It rebuilds a full, dense teacher-probability grid from the cached top-100 logits and compares it against the student's own dense grid of log-probabilities. This is the version closest to how online distillation already works, so we use it as our correctness baseline, but it holds the full vocabulary × sequence grid in memory, twice over. Forward-chunked KL keeps the teacher sparse (only its cached top-100 logits per position, never expanded into a dense grid) and computes the loss piece by piece, one slice of sequence positions at a time. This removes the dense teacher and the dense comparison, and turns out to be the fastest of the three methods in our benchmarks. It still has one blind spot, though: the student's own logits, the grid produced by the model's output layer, are still computed in full and kept around for the backward pass, so memory still grows steeply with sequence length. Fused chunked KL , our main contribution, goes a step further and fuses the model's output projection directly into the loss computation. It never produces the student's full logits grid at all: it processes one chunk of the sequence at a time end to end, projecting hidden states to logits for that chunk, folding the result into the running loss, and discarding the chunk before moving to the next one. The backward pass recomputes each chunk on the fly instead of storing it. The cost is doing that projection twice, once forward, once in backward, but in exchange, peak memory grows only linearly with sequence length instead of spiking with the full vocabulary × sequence size. The GIF below shows the difference between the dense and fused-chunked approaches: one builds the whole comparison grid and holds onto all of it, the other builds and discards one slice at a time, so memory never grows beyond a single chunk. We have open-sourced the chunked-loss implementation: github.com/CompactifAI/Full-Chunked-KL-Loss What this changes in practice The table below puts all four setups head to head: online distillation, and the three offline loss implementations just described. Comparing them on a single H200 GPU with Llama 3.1 8B Instruct as teacher and a 3.2B Llama model as student at an 8K token context, all four reach near-identical training loss, even though the offline runs train against only the cached top-100 logits per token. Method (8K context, single H200) Peak memory Iteration time Throughput Online distillation 102.8 GB 25.9 s 237 TFLOP/s Offline, dense KL 78.3 GB 18.5 s 331 TFLOP/s Offline, forward-chunked KL 61.8 GB 18.4 s 335 TFLOP/s Offline, fused chunked KL 58.3 GB 20.2 s 304 TFLOP/s The loss curves overlap almost exactly across all four methods, confirming offline distillation with top-100 cached logits is lossless relative to online distillation. Source: paper Figure 2. At this sequence length, the fused chunked loss is not yet the fastest option, its extra backward-pass projection costs a bit of speed, but its real advantage only shows up as context length grows, which the next section demonstrates. Scaling to long context lengths To see the scaling pattern more starkly, we ran an isolated benchmark on a toy output-projection network (no transformer body, just the loss kernel). 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.6× 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 is about 3.3× faster per iteration at that length. Distilling a GPT-OSS 20B model at a 32,768-token context, the memory freed by the fused loss let the setup shrink from four GPU nodes down to one. Step time fell from 57.0 to 12.23 seconds, about 5× faster, and throughput per GPU rose from 74.2 to 345.7 TFLOP/s. The resulting student The efficient offline setup is what made a large-scale distillation campaign affordable in the first place. The resulting compact student, distilled from Llama 3.1 8B Instruct down to about 3.2B parameters, retains most of the teacher's accuracy on BoolQ and HellaSwag, stays within about nine points of it on MMLU, at less than half the parameter count. The student retains most of the teacher's short-context accuracy at less than half the size. Source: paper Figure 6. This work is part of Multiverse Computing's ongoing research into making distillation and healing practical to run at scale , not just as a one-off recipe, but as something teams can iterate on cheaply. The paper also covers additional ablations, such as how the choice of loss function and sequence packing affect recovery quality. Want the full technical details, including the closed-form gradient behind the fused chunked loss and the complete training configuration? Read the full paper , or get in touch with our team to talk about applying this to your own distillation pipelines. We have also open-sourced the chunked-loss implementation: github.com/CompactifAI/Full-Chunked-KL-Loss