Making Knowledge Distillation Cheap Enough to Run at Scale

让知识蒸馏变得足够廉价以实现规模化运行

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. 知识蒸馏（即训练一个较小的学生模型来匹配较大教师模型的性能）是机器学习中一种广为人知的技术。随着近期开源大语言模型（如 gpt-oss、Qwen、GLM 或 Kimi）的浪潮，它再次成为主流研究课题。

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. 部署这些超大规模模型成本高昂：近期的 Kimi-K3 模型拥有 2.8 万亿参数，仅加载就需要约 3TB 的显存。因此，通过知识蒸馏将其压缩为较小模型并恢复原始能力已成为行业标准做法，Nvidia（Nemotron 3 Puzzle 75B）和 Multiverse Computing（Hypernova 60B）等公司近期都发布了高质量的压缩模型。

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. 蒸馏步骤决定了最终模型的大部分质量，但它通常也是整个流程中最昂贵的部分。同时保持教师模型和学生模型在显存中，并为每个 token 生成覆盖整个词表的概率分布，需要巨大的显存空间，这通常只有在数百个 GPU 和精细的张量并行策略下才能实现。

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. 我们最新的论文《针对大语言模型的高效知识蒸馏：离线 Top-K Logits 与融合分块 KL 损失》通过两项系统性改进解决了这一问题：首先是一次性缓存教师模型的 Top-K Logits，从而无需让教师模型与学生模型同时驻留内存；其次是一种全新的、内存高效的 KL 散度损失函数，它避免了生成完整的“词表大小 × 序列长度”矩阵，将显存占用降低到了远低于 PyTorch 或 NVIDIA Megatron-Bridge 等库默认实现水平的程度。

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. 这两项改进共同降低了训练成本，使得在单张 GPU 上进行长上下文修复成为可能，并让大规模实验变得切实可行。

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. 标准的设置是使用 KL 散度损失（KL loss）进行在线蒸馏，这要求教师模型和学生模型同时加载。在每个训练步骤中，教师模型运行完整的前向传播以生成输出分布，学生模型则被训练以匹配该分布。这是表达能力最强的设置，因为可以获取完整的教师分布，但它也是内存和计算密集度最高的：每个 token 位置都需要保留两个完整的词表张量，且尽管教师模型的行为在整个训练过程中不会改变，它仍需在每一步被重新计算。

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. 以一个实际例子说明，gpt-oss-120b 的词表大小为 201,088 个 token。在序列长度为 32K、批次大小为 4 的情况下，仅教师概率张量的形状就达到了 4 × 201,088 × 32,768；使用 bfloat16 精度时，单个张量就需要约 50GB 显存。加上梯度、激活值、模型权重和优化器状态，单次蒸馏训练迭代的显存峰值可达约 250GB，这超过了单张 H200 或 B200 GPU 的显存容量。

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. 在这篇文章中，我们展示了通过重构 KL 损失函数以分块方式处理数据，可以将这种成本降低到几乎可以忽略不计。密集型 KL 损失的显存峰值约为 250GB，超过了单张 H200 的 141GB 容量。而融合分块损失函数则不会产生这种峰值，其峰值仅为约 128GB。

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. 离线蒸馏。 我们不再在每一步重新计算教师模型，而是计算一次其输出，缓存每个位置概率最高的 100 个 token，并基于该缓存训练学生模型。教师模型在训练期间无需驻留内存，且一旦缓存生成，无需再次运行，因此同一缓存可在多次消融实验中重复使用。

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. 融合分块 KL 损失。 要理解为什么损失函数本身很昂贵，可以想象它实际构建的内容：对于序列中的每个 token 位置和词表中的每个词，损失函数都需要一个数值来描述学生模型的预测与教师模型预测的差异程度。如果将其排列成网格，每一行对应一个词表条目，每一列对应一个序列位置。对于 10 万以上的词表和长序列，这个网格是巨大的，而计算 KL 损失的默认方式是在产生单个数值之前就构建出整个网格。

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. 密集型 KL 是教科书式的方法。它从缓存的 Top-100 Logits 重建一个完整的密集教师概率网格，并将其与学生模型自身的密集对数概率网格进行比较。这是最接近现有在线蒸馏工作方式的版本，因此我们将其作为正确性基准，但它在内存中保留了完整的“词表 × 序列”网格，且是双份的。

• 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. 前向分块 KL 保持教师模型为稀疏状态（仅保留每个位置缓存的 Top-100 Logits，从不展开为密集网格），并分块计算损失，每次处理序列位置的一个切片。这消除了密集教师网格和密集比较，在我们的基准测试中是三种方法中最快的一种。但它仍有一个盲点：学生模型自身的 Logits（由模型输出层生成的网格）仍然被完整计算并保留以用于反向传播，因此内存占用仍会随序列长度急剧增加。

• 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. 融合分块 KL 是我们的主要贡献，它更进一步，将模型的输出投影直接融合到损失计算中。它根本不会生成学生模型的完整 Logits 网格：它一次处理序列的一个分块，端到端地将隐藏状态投影为该分块的 Logits，将结果合并到运行中的损失中，并在移动到下一个分块前丢弃当前分块。反向传播过程会实时重新计算每个分块，而不是将其存储起来。