オフラインTop-K蒸留で常駐教師モデルを排除、単一H200のスループットが最大41%向上

CompactifAIチームは、大規模言語モデルの知識蒸留コストを削減する手法を公開した。この手法では、教師モデルの推論と学生モデルの学習を分離する。従来のオンライン蒸留では、教師モデルと学生モデルを同時にGPU上へ常駐させ、各学習ステップで教師分布を再生成する必要があった。新しいワークフローでは、教師モデルが各位置で出力するTop-K logitsを事前計算して保存し、学生モデルの学習時にキャッシュから直接読み込むことで、教師モデルが占有していたGPUメモリと計算リソースを解放する。

著者らの報告によると、オフライン蒸留はほぼ同等のtraining lossを達成しつつ、1イテレーション当たり約29%高速化し、単一のNVIDIA H200におけるスループットを最大41%向上させた。Top-Kキャッシュは追加のディスク容量を必要とするものの、複数回の学生モデル学習やハイパーパラメータ実験で再利用できるため、モデル圧縮、量子化後の能力回復、大規模なアブレーション実験に特に適している。また、高コストな教師モデル推論を、スケジューリング可能な一度限りのデータ作成処理へ変換できる。

2つ目の改良は、融合・チャンク化されたKL divergence lossである。一般的な実装では、「シーケンス長×語彙サイズ」に近い形状の完全な学生logitsを生成して教師分布と比較するため、長いコンテキストではGPUメモリ使用量が急増しやすい。新しいkernelは、出力投影、softmax、KLをチャンク単位で計算し、完全な語彙テンソルを実体化しない。これにより、ピークメモリはシーケンス長に対してより緩やかに増加する。この結果、単一GPUで学習コンテキストを約8,192 tokenから32,768 tokenへ拡張できた。さらに、出力ヘッドを用いたテストにより、4Kから256K tokenまでのメモリ使用量とイテレーション時間のスケーリングを検証した。

チームはFull-Chunked-KL-Lossの実装を公開しているが、論文には関連特許を出願中であると記載されている。現時点では、kernelのマイクロベンチマークを、そのまま完全なモデルのエンドツーエンド高速化と見なすことはできない。また、Top-Kによる打ち切りで、教師モデルが持つ低確率のシグナルが失われる可能性もある。エンジニアリングチームは今後、異なる語彙サイズ、モデルファミリー、K値、データキャッシュのI/Oに加え、マルチGPU学習において通信がGPUメモリに代わる新たなボトルネックになるかどうかにも注目すべきである。