离线 Top-K Logits 与融合分块 KL：降低大模型知识蒸馏显存开销

背景问题

知识蒸馏需要让较小的学生模型匹配较大的教师模型。在线蒸馏的常规做法是同时加载教师与学生，每一步都让教师做一次前向，产生完整的词表概率分布。作者指出，以 gpt-oss-120b 的 201,088 词表为例，在 32K 序列长度、batch size 4 时，仅教师概率张量在 bfloat16 下就约需 50GB 显存，叠加梯度、激活、权重与优化器状态后，单次迭代峰值可达约 250GB，超过单张 H200 的 141GB 容量。

两项改动

• 离线蒸馏：教师只运行一次，缓存每个位置 top-100 最可能 token 的 logits，学生针对这份缓存训练。教师无需在训练期间常驻显存，同一缓存可复用于多组消融实验。
• 融合分块 KL 损失：作者比较了三种数学等价实现。稠密 KL 会重建完整的词表 × 序列网格，保留两份；前向分块 KL 让教师保持稀疏、逐片计算，但仍完整保留学生的 logits 供反向使用；融合分块 KL 把输出投影直接并入损失计算，一次处理一个序列块，前向算出 logits 后并入累计损失即丢弃，反向时按需重算。代价是投影要算两次，换来峰值显存随序列长度线性增长。

实测结果

单张 H200，教师 Llama 3.1 8B Instruct、学生 3.2B Llama，8K 上下文，四种设置训练损失接近：

| 方法 | 峰值显存 | 迭代时间 | 吞吐 | | --- | --- | --- | --- | | 在线蒸馏 | 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 起直接失败；256K 时融合方案用 11.6 GiB，另一个分块变体为 134.2 GiB，前者每轮约快 3.3×。另以 GPT-OSS 20B、32,768 token 上下文蒸馏时，显存节省使配置从四个 GPU 节点缩到一个，单步从 57.0 秒降至 12.23 秒，每 GPU 吞吐从 74.2 升至 345.7 TFLOP/s。

学生模型效果

由 Llama 3.1 8B Instruct 蒸馏出的约 3.2B 学生在 BoolQ 与 HellaSwag 上保留了教师大部分准确率，MMLU 上差距约在九个点以内，参数不到教师一半。原文未披露完整的训练配置与消融数据，具体细节需参见论文。作者已开源分块损失实现。

注：文中数据为博客与论文自述，未经过独立复现验证。