利用离线 Top-K Logits 和融合分块 KL 损失提高 LLM 知识蒸馏效率
一项实践研究提出了两种提高大语言模型知识蒸馏效率的系统技术。首先，预先缓存教师模型的 Top-K logits，可在不将教师模型保留于 GPU 内存中的情况下进行训练，同时保持与在线蒸馏几乎相同的训练损失。在单张 H200 GPU 上，每次迭代速度提升约 29%，吞吐量最高提升 41%。其次，融合式分块 KL 损失避免生成覆盖完整词表的 logits 张量，从而降低峰值内存使用，并支持在单张 GPU 上训练最长 32,768 个 token 的上下文。额外基准测试还评估了最高 256K token 下的内存和迭代速率扩展性。相关实现已开源。
本文来源：Hugging Face，仅供学习参考，版权归原作者所有。