大语言模型的高效知识蒸馏：离线 Top-K Logits 与融合分块 KL 损失

Multiverse Computing 推出了一种全新的知识蒸馏方法，能够显著降低 VRAM 使用量和训练成本，从而实现在单个 GPU 上进行长上下文蒸馏。通过将离线 Logit 缓存与融合分块 KL 散度损失（fused chunked KL-divergence loss）相结合，研究人员可以在不产生通常与全词表概率分布相关的巨大内存开销的情况下，将大型教师模型蒸馏到较小的学生模型中。

标准知识蒸馏的高昂成本

使用 Kullback-Leibler (KL) 散度损失的传统“在线”知识蒸馏在计算上非常昂贵，因为它要求教师模型和学生模型同时加载到 VRAM 中。在每个训练步骤中，教师模型必须进行一次完整的正向传播，以生成每个 Token 位置在整个词表上的概率分布。

对于大型模型，这会造成巨大的内存瓶颈。例如，使用具有 201,088 个 Token 词表的 gpt-oss-120b 模型，在序列长度为 32K、Batch Size 为 4 的情况下，仅教师概率张量（teacher-probability tensor）在 bfloat16 精度下就需要约 50GB 的 VRAM。如果算上梯度、激活值和优化器状态，单次迭代的 VRAM 峰值可能达到约 250GB，甚至超过了 H200 或 B200 等高端 GPU 的容量。

实现可扩展蒸馏的两项系统性变革

为了解决这些内存限制，Multiverse Computing 实施了两个主要的工程技术变革：

1. 通过 Top-K Logit 缓存进行离线蒸馏

与其在每一步都重新计算教师模型的输出，该系统会计算一次教师模型的输出，并为每个位置缓存最有可能的前 100 个 Token。随后，学生模型针对该缓存进行训练。这消除了在训练过程中将教师模型保留在内存中的需求，并允许在不同的实验消融研究中重复使用相同的缓存。

2. 融合分块 KL 损失

标准的 KL 损失实现会构建一个完整的“词表 $\times$ 序列”网格来计算学生与教师之间的差异。Multiverse Computing 的“融合分块 KL 损失”通过将模型的输出投影直接融合到损失计算中来优化这一点。

该过程并非实例化完整的 Logits 网格，而是按如下方式运行：

• 它一次处理序列的一个分块（chunk）。
• 它仅为该特定分块将隐藏状态投影到 Logits。
• 它将结果折叠进运行中的损失值中，并立即丢弃该分块。
• 反向传播过程会实时重新计算每个分块。

这种方法确保了峰值内存随序列长度线性增长，而不是根据全词表大小出现激增。

性能与内存基准测试

在单个 H200 GPU 上，使用 Llama 3.1 8B Instruct（教师）和 3.2B Llama 模型（学生）在 8K Token 上下文下对比不同的损失实现，研究人员发现所有方法都达到了近乎相同的训练损失。这证实了仅使用前 100 个缓存的 Logits 进行离线蒸馏相对于在线蒸馏是几乎无损的。

方法 (8K 上下文, 单个 H200)	峰值内存	迭代时间	吞吐量
在线蒸馏	102.8 GB	25.9 s	237 TFLOP/s
离线, 密集型 KL	78.3 GB	18.5 s	331 TFLOP/s
离线, 前向分块 KL	61.8 GB	18.4 s	335 TFLOP/s
离线, 融合分块 KL	58.3 GB	20.2 s	304 TFLOP/s

扩展至长上下文

融合分块损失的优势随着上下文长度的增加而变得更加显著。在仅使用玩具输出投影网络的隔离基准测试中：

• 在 32K Token 时： 峰值内存从 85.2 GiB（密集型损失）降至 5.45 GiB（完全分块），降幅达 15.6 倍。
• 在 64K Token 时： 密集型损失实现完全失败。
• 在 256K Token 时： 完全分块损失仅使用 11.6 GiB，而次优的分块变体使用了 134.2 GiB，且单次迭代速度快了 3.3 倍。

在实际场景中，通过在 32,768 Token 上下文下蒸馏 GPT-OSS 20B 模型，融合损失将所需的硬件从四个 GPU 节点减少到了一个。这使得单步时间从 57.0 秒缩短至 12.23 秒，并使每个 GPU 的吞吐量从 74.2 提升至 345.7 TFLOP/s。

学生模型能力

该流水线的效率使得大规模蒸馏工作成为可能，并产生了一个从 Llama 3.1 8B Instruct 蒸馏而来的 3.2B 参数学生模型。尽管参数量不到教师模型的一半，但这个紧凑的学生模型在 BoolQ 和 HellaSwag 任务上保留了教师模型大部分的精度，并且在 MMLU 基准测试中与教师模型的差距保持在约 9 分以内。