知識蒸留のコストを削減するオフラインtop-Kキャッシュとチャンク化KL損失

蒸留コストの課題

大規模言語モデルを小さいモデルに圧縮する知識蒸留は、推論コストを下げるための標準的な手法になっている。たとえば、最近登場したKimi-K3は2.8兆パラメータで、読み込むだけで約3TBのVRAMを必要とする。こうした巨大モデルの能力を損なわずに扱いやすくするため、蒸留によって小さいモデルを作り、元の性能を回復させるというアプローチが広く使われている。NvidiaのNemotron 3 Puzzle 75BやMultiverse ComputingのHypernova 60Bなど、圧縮モデルの公開も相次いでいる。

しかし、蒸留ステップ自体がパイプラインの中で最もコストのかかる部分になりがちだ。従来のオンライン蒸留では、教師と生徒の両方を同時にメモリに保持し、全語彙に対する確率分布をトークンごとに計算する必要がある。具体例として、gpt-oss-120bの語彙サイズ201,088トークン、シーケンス長32K、バッチサイズ4の場合、教師の確率テンソルだけで約50GBになり、訓練イテレーション全体のピーク時には約250GBのVRAMが必要になる。これはH200の141GBを超え、単一GPUでは実行できない。

2つのシステム変更

この問題に対して、Multiverse Computingの論文は2つのシステム変更を提案している。1つ目はオフライン蒸留で、教師モデルの出力を一度だけ計算し、各位置で尤度の高い上位100件のロジットをキャッシュする。訓練中は教師をメモリに保持する必要がなくなり、キャッシュを再利用して様々な実験を行える。

2つ目は、融合型チャンク化KL損失（fused chunked KL loss）と呼ばれる新しい損失計算だ。従来のKL損失は、語彙数×シーケンス長の巨大な行列を構築してから損失を計算するため、メモリを大量に消費する。提案手法では、シーケンスをチャンクに分割し、出力射影と損失計算を融合して、チャンクごとに処理して破棄する。バックワード時にはチャンクを再計算するため、ピークメモリが語彙数に依存しなくなる。

性能評価の結果

トイネットワークによるベンチマークでは、32Kトークンのコンテキストでもピークメモリが85.2GiBから5.45GiBまで低下し、15.6倍の削減を実現した。従来のdense KL損失は64Kトークン以降でメモリ不足になり実行不能になるが、提案手法は256Kトークンでも11.6GiBで動作し、次に良いチャンク化手法の134.2GiBと比べて約3.3倍高速だった。

実際のGPT-OSS 20Bモデルの蒸留では、32,768トークンのコンテキストで、必要だったGPUノードが4つから1つに減った。ステップ時間は57.0秒から12.23秒へと約5倍速くなり、GPUあたりのスループットも74.2TFLOP/sから345.7TFLOP/sに向上した。

蒸留された生徒モデル

この効率的なオフライン蒸留は、大規模な蒸留キャンペーンを手頃なコストで実行するための基盤になる。論文では、Llama 3.1 8B Instructを教師に、約3.2Bパラメータの生徒モデルを蒸留している。結果として、BoolQとHellaSwagでは教師にほぼ匹敵する精度を維持し、MMLUでは約9ポイント差に収まった。パラメータ数は半分以下である。

この研究はMultiverse Computingが進める蒸留と「ヒーリング」を実用規模で実行する研究の一部で、損失関数の選択やシーケンスのパッキングが回復品質に与える影響なども報告している。実装はオープンソースとして公開されており、誰でも既存の蒸留パイプラインに組み込める。

ただし、トップ100ロジットへの制限が常に十分かどうかは、タスクやモデルに依存する可能性がある。論文ではオンライン蒸留とほぼ同一の損失曲線が得られたと報告されているが、あくまで特定の条件下での結果である。また、オープンソースの実装は特定のフレームワーク向けに提供されており、既存のパイプラインへの統合には追加の作業が必要かもしれない。