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

摘要

小语言模型通常是在延迟、成本和本地部署约束较为严格的情况下唯一可行的部署选择，但这类模型很少从头开始训练：压缩模型通常通过知识蒸馏（KD）来恢复。这一恢复步骤在很大程度上决定了最终质量，然而其成本却很高。我们提出了一项面向实践者的蒸馏训练效率研究，围绕两个系统贡献展开。第一，我们表明离线KD（一次性缓存教师模型的Top-K logits，并让学生模型针对该缓存进行训练）能够在几乎相同的训练损失下与在线蒸馏持平，同时将教师模型从内存中移除，每次迭代运行速度提升约29%，在单个H200 GPU上吞吐量最高可提升41%。第二，我们引入了一种融合式分块KL损失，它从不实例化完整的词表大小的logit张量，从而使峰值内存与序列长度呈线性关系。这消除了原本会限制上下文长度的内存峰值，使我们能够在单个GPU上以四倍的上下文长度（32,768个token）进行训练。另一个仅包含输出头的玩具基准单独测试了损失核，并确认了其在4K到256K token范围内的内存和迭代速率扩展特性。这些改进共同使大规模恢复和数百次消融实验变得负担得起。我们还报告了关于损失设计和序列打包的辅助消融实验。我们发布了分块损失的实现：https://github.com/CompactifAI/Full-Chunked-KL-Loss。

English

Small language models are often the only option for deployment under tight latency, cost, and on-premises constraints, but they are rarely trained from scratch: a compressed model is usually recovered through knowledge distillation (KD). This recovery step largely decides the final quality, yet it is expensive. We present a practitioner's study of how to make distillation training efficient, organised around two systems contributions. First, we show that offline KD (caching the teacher's top-K logits once and training the student against the cache) matches online distillation at near-identical training loss while removing the teacher from memory, running about 29\% faster per iteration, and reaching up to 41\% higher throughput on a single H200 GPU. Second, we introduce a fused, chunked KL loss that never materialises the full vocabulary-sized logit tensor, making peak memory linear in the sequence length. This removes the memory spike that otherwise caps context length and lets us train at four times the context (32{,}768 tokens) on a single GPU. A separate output-head-only toy benchmark isolates the loss kernel and confirms its memory and iteration-rate scaling from 4K to 256K tokens. Together these make large-scale healing and hundreds of ablations affordable. We also report supporting ablations on loss design and sequence packing. We release our chunked-loss implementation: https://github.com/CompactifAI/Full-Chunked-KL-Loss.