LLM 的高效 知识蒸馏：离线 Top-K logits 和融合分块 KL 损失

小语言模型通常是在严格的延迟、成本和本地约束下部署的唯一选择，但它们很少从头开始训练：压缩模型通常通过 知识蒸馏 (KD) 恢复。这一恢复步骤在很大程度上决定了最终的质量，但其成本很高。我们提出了一项从业者关于如何提高 蒸馏 训练效率的研究，该研究围绕两个系统贡献进行组织。首先，我们展示了离线 KD（缓存教师的 top-$K$ logits 一次并根据缓存训练学生）以几乎相同的训练损失与在线 蒸馏 匹配，同时从内存中删除教师，每次迭代运行速度提高约 29\%，并且在单个 H200 GPU 上达到高出 41\% 的吞吐量。其次，我们引入了一个 \emph{融合的、分块的 KL 损失}，它永远不会实现完整词汇大小的 logits 张量，从而使峰值记忆在序列长度上呈线性。这消除了内存峰值，否则会限制上下文长度，并让我们可以在单个 GPU 上以上下文的四倍（32{,}768 个词元）进行训练。一个单独的仅输出头玩具基准隔离了损失内核，并确认其内存和迭代率从 4K 到 256K 词元的缩放。这些共同使得大规模治疗和数百次消融变得负担得起。我们还报告了对损失设计和序列打包的支持消融。我们发布了分块损失实现：https://github.com/CompactifAI/Full-Chunked-KL-Loss。