詳細を読む
AI企業のMultiverse Computingは2026年8月10日、大規模言語モデルの知識蒸留にかかるGPUメモリを大幅に削減する手法を論文とともに公開しました。教師モデルの出力を事前に保存するオフライン蒸留と、語彙全体の行列を作らない融合チャンク型KL損失という二つの改良により、これまで数百基のGPUを要した工程を1ノードで回せるようにしたのが要点です。
従来の蒸留は教師と生徒の両モデルを同時にメモリへ載せ、各トークンで語彙全体の確率分布を計算する必要がありました。gpt-oss-120bの語彙は20万1088トークンあり、系列長32K・バッチ4では教師側の確率テンソルだけでbfloat16換算で約50GBに達します。勾配や活性値、最適化器の状態まで含めると1回の学習ステップは約250GBに跳ね上がり、H200やB200の搭載量を超えてしまいます。
今回の手法はまず教師の出力を1度だけ計算し、位置ごとに上位100件のロジットだけを保存します。学習中に教師をメモリへ置く必要がなくなり、同じキャッシュを何度もの追試で使い回せます。もう一つの融合チャンク型KL損失は出力層の射影を損失計算に組み込み、系列を小さな塊ごとに処理しては捨てることで、巨大な比較行列を一度も作りません。
では実際にどれだけ軽くなるのでしょうか。出力層だけを取り出したベンチマークでは、32Kトークン時のピークメモリが85.2GiBから5.45GiBへと15.6分の1に下がり、従来手法が動かない64Kトークン以降でも学習を続けられました。GPT-OSS 20Bを32Kで蒸留した実測でも必要なGPUノードは4から1へ減り、1ステップは57.0秒から12.23秒へと約5倍速くなっています。
精度面でも、上位100件のロジットしか使わないオフライン学習の損失曲線は従来のオンライン蒸留とほぼ重なり、実質的な劣化は確認されていません。Llama 3.1 8B Instructから約32億パラメータへ圧縮した生徒モデルはBoolQやHellaSwagで教師の精度をほぼ保ち、MMLUでも差は約9ポイントにとどまります。同社はチャンク損失の実装をGitHubで公開しており、自社で蒸留を試す企業にとって検証の敷居は確実に下がりました。