提升大模型压缩效率：离线缓存+分块损失，训练快41%且上下文长4倍。
Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss
- 离线缓存教师模型前K个输出概率，避免重复计算
- 分块融合KL损失使内存仅随序列长度线性增长，支持32768词长上下文
- 适合资源受限场景下的大模型高效蒸馏，支持大规模实验
小语言模型常因延迟、成本和本地部署限制而被使用，但通常不从头训练：通过知识蒸馏（KD）压缩得到。该过程决定最终质量，但代价高昂。本文提出两项系统级改进：第一，离线蒸馏——提前缓存教师模型的top-K logits，训练时直接使用缓存数据，无需在内存中保留教师模型，单次迭代快29%，单卡H200最高吞吐提升41%；第二，提出融合分块的KL损失，避免生成完整词汇表大小的logit张量，峰值内存仅与序列长度线性相关，突破原有内存瓶颈，使单卡训练上下文长度达32,768词，是之前的4倍。独立输出头小规模测试验证了该损失在4K至256K词范围内的内存与迭代速率良好缩放。二者结合使大规模模型修复与数百次消融实验成为可能。我们开源了分块损失实现：https://github.com/CompactifAI/Full-Chunked-KL-Loss。
原文摘要 · Abstract (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 \emph{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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。