第6章 — プルーニングと知識蒸留

LLM Primer VI: AIシステムのスケーリング を章ごとに紹介していくウォークスルー、第6回。重みの数を直接攻める章 — まず重要でない重みをゼロにし、次に大きなモデルの振る舞いを小さなモデルへ転写する。

なぜこの章があるのか

量子化は重みあたりのバイト数を縮めます。プルーニングはそもそも掛け算される重みの数を縮めます。知識蒸留は大きな教師の振る舞いを、より小さく端から端まで安く走る生徒へ転写します。プルーニングは同じアーキテクチャ内での編集でオーバーヘッドを削る動き。蒸留はアーキテクチャの変更で、訓練時のコストを支払う代わりに恒久的な推論時の節約を得る動きです。第6章は両方を歩きます — それぞれが依存する機構、実務家が突き当たる失敗モード、そしてHopperとBlackwellで構造化プルーニングを再び経済的に興味深いものにしたハードウェア加速パターン(2:4スパース性)。

6.1 2:4スパース性はハードウェアが報いるプルーニングパターン

非構造化のマグニチュードプルーニング — 最小の50パーセントの重みをゼロにして短くファインチューン — は数値的には機能しますが加速しません。GPUは密なGEMMを実行するので、散らばったゼロも掛け算される。構造化プルーニングは、ハードウェアがスキップできる単位を丸ごと削除します。アテンションヘッド、MLPニューロン、あるいは(重要な例)重み行列内のN:Mパターン。HopperとBlackwellは2:4スパース性をネイティブ対応します — 入力次元に沿った4つの連続する重みの中に必ず2つのゼロがある。チップは各4重みグループを2つの非ゼロ値と2ビットマスクとして格納し、スパーステンソルコアはゼロの掛け算をスキップし、帯域とFLOPsが両方おおよそ半分に落ちます。測定されたデコードスループットで1.7〜1.9倍のスピードアップ。コツは、2:4がプルーニング時に課されなければならないこと。アルゴリズムは各行列を4つの塊で歩き、小さい2つをゼロに強制し、短くファインチューンするかSparseGPTを適用して補償します。マスクはメタデータで、ランタイムの節約は自動です。

6.2 蒸留は教師のargmaxではなく分布を転写する

訓練済みの教師LLMは各位置で語彙上の分布 — 128k次元のsoftmax — を生みます。argmaxは貪欲デコード時に出力するトークン。しかし分布は教師の不確実性、類義語間の選好構造、キャリブレーションを符号化しています。生徒をその分布に一致させるよう訓練すると、それら全てが転写される。損失は温度付きで柔らかくした教師と生徒の出力の間のKLダイバージェンスで、温度は勾配信号を低確率の裾まで広げる。得られる生徒は、同じデータでゼロから訓練された小さなモデルではなく、教師の小さく速い版として振る舞います。Patient Knowledge Distillationはレシピを中間層に拡張し、マップされた生徒-教師の層ペア間で隠れ状態を一致させる。MiniLLMは、より微妙な失敗 — 順KLの生徒が複数のもっともらしい教師の継続の間で及び腰になり、無難な出力を生むモード平均化問題 — に対処するため、逆KLに切り替え、教師の複数モードの緩やかな網羅より1つのモードとの鋭い一致を好みます。

6.3 3つの圧縮は特定の順序で積み上げる

先に蒸留、次にプルーニング、最後に量子化。蒸留は最も高価で最もアーキテクチャに侵襲的な操作で、他の選択が固定される前に密で高精度のモデル上で一度実行する意味があります。生徒をSparseGPTで2:4にプルーニングし、短い回復ファインチューンを行う。スパース化された生徒を llm-compressor でFP8に量子化する。動くスタック: 70B BF16教師 → 13B BF16蒸留生徒 → 2:4スパース化 → FP8量子化。順伝播あたりの重みバイトは140 GBからおよそ6.5 GBに落ち — 22倍の削減 — HopperのFP8と2:4のパスがそれをGPUあたり8〜10倍のスループット向上に変えます。このスタックが縮めないのはKVキャッシュで、これは重み数ではなく隠れ次元、ヘッド数、コンテキスト長に依存します。そこが第8章の主題です。そしてモデルレベルの圧縮は推論ごとの話で、GPUにヘッドルームを与えますが、そのヘッドルームを同時ユーザーで埋めるランタイム側のレバーはバッチングです。

第6章が敷いたもの

第5章と第6章でモデル側の圧縮を出し尽くしました — 重みあたりのバイト数を量子化し、重み数をプルーニングし、モデル全体を蒸留する。GPUは順伝播ごとにヘッドルームを持つ状態になりました。第7章はそのヘッドルームを有用な仕事で埋めるシステムレベルのレバー — バッチング — に移ります。静的バッチングはほぼ動きますが、最速完了問題で崩れる。反復単位でスケジューリングする連続バッチングこそが本番エンジンが走らせるもので、帯域律速のデコード局面を遊休チップから多数のユーザーを同時に捌くチップに変えるものです。

次回 — 第7章: 高度なバッチング戦略。 静的、動的、連続バッチング — そして連続バッチングがKVキャッシュに対して作り出す借金。