LLMのための効率的な知識蒸留：オフラインTop-Kロジットと融合チャンク化KL損失

要旨

小規模言語モデルは、厳しいレイテンシ、コスト、オンプレミス環境の制約下でのデプロイメントにおいてしばしば唯一の選択肢となる。しかし、それらがゼロから学習されることはまれであり、通常は知識蒸留（KD）によって圧縮モデルが復元される。この復元ステップは最終的な品質を大きく決定するが、そのコストは高い。本稿では、2つのシステム上の貢献を中心に、蒸留学習を効率的にするための実務者向けの研究を提示する。第一に、オフラインKD（教師のtop-Kロジットを一度キャッシュし、そのキャッシュに対して生徒モデルを学習する手法）が、ほぼ同一の訓練損失でオンライン蒸留に匹敵し、教師モデルをメモリから取り除きつつ、1反復あたり約29%高速に動作し、単一のH200 GPU上で最大41%高いスループットを達成することを示す。第二に、語彙サイズ全体のロジットテンソルを実体化しない融合チャンク化KL損失を導入し、ピークメモリをシーケンス長に対して線形にする。これにより、そうでなければコンテキスト長を制限していたメモリスパイクが除去され、単一GPU上で4倍のコンテキスト（32,768トークン）で学習できるようになる。出力ヘッドのみの別の簡易ベンチマークは損失カーネルを分離し、4Kから256Kトークンまでのこの損失のメモリおよび反復レートのスケーリングを確認する。これらを組み合わせることで、大規模な修復と数百回のアブレーションが手頃なコストで実行可能になる。さらに、損失設計とシーケンスパッキングに関する裏付けとなるアブレーションも報告する。チャンク化損失の実装を公開する：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.