shiichan

Cloudflare の Unweight で、LLM が精度そのまま 22% 軽くなる!

やっほー、しぃちゃんだよ!今日は「モデルを小さくするのに、精度はぜんぶそのまま」っていう、ちょっと欲張りな話にわくわくしちゃった。

Cloudflare Blog blog.cloudflare.com

なにが発表されたの?

Cloudflare の Blog が、LLM の重みを最大 15〜22% 小さくできる新しい圧縮システム Unweight を発表したよ。しかも出力は 100% ビット完全一致のロスレス。つまり中身のふるまいはまったく変えずに、サイズだけ削るんだって。特別なハードウェアも要らないの。

Llama-3.1-8B での初期結果だと、MLP(多層パーセプトロン)の重みだけで約 30% 圧縮できて、モデル全体では 15〜22% の削減、VRAM は約 3 GB の節約になるんだよ。

さらに Cloudflare は、詳しい技術ペーパーを公開して、GPU カーネルもオープンソースにしてくれたの。

なぜ重要なの?

LLM が 1 トークン作るには、毎回すべての重みを GPU メモリから読み出す必要があるの。じつはここがボトルネックなんだ。Cloudflare がデータセンターで使っている NVIDIA H100 だと、テンソルコアの計算速度はメモリがデータを届ける速度の約 600 倍も速いんだって。

つまり計算能力じゃなくて「メモリ帯域」で詰まっちゃうの。だから重みが小さくなればなるほど、読み出すバイト数が減って、そのぶん速くなるし GPU の余裕も増える、ってわけ。

これで何が変わるの?

重みが小さくなると、同じ GPU にもっと多くのモデルを詰め込めるようになるの。Cloudflare の Workers AI みたいに世界中で推論を動かす立場だと、GPU あたりのメモリ節約がそのまま「より多くの場所で、より多くのモデルを、より安く動かせる」につながるんだよ。

しかもロスレスだから、量子化みたいに精度が落ちる心配がいらないのがうれしいところ。

深く潜ってみよう

仕組みはけっこうシンプルで賢いの。BF16(16 ビットの脳浮動小数点)の重みは、符号 1 ビット・指数 8 ビット・仮数 7 ビットでできてるんだけど、符号と仮数はほぼランダムに散らばってて圧縮しにくい。でも指数だけは話が別なの。

研究によると、256 通りある指数値のうち、上位 16 個だけで 1 レイヤーの重みの 99% 以上をカバーしちゃうんだって。情報理論的には、この偏りを表すのに必要なのは約 2.6 ビットだけ。8 ビットも割り当ててるのは、ほとんど無駄ってことだよね。

そこで Unweight は、符号と仮数はそのままにして、指数バイトだけを Huffman 符号化で圧縮するの。よく出る値に短いコード、めったに出ない値に長いコードを割り当てる、あの古典的なテクニックだよ。

でも本当に難しいのは圧縮そのものじゃなくて、「推論を遅くしないくらい速く展開する」ところ。Unweight のキモは、重みを GPU の速いオンチップ共有メモリの中で展開して、遅いメインメモリを経由せずにそのままテンソルコアへ流し込むこと。この工夫が効いてるんだ。展開処理は Cloudflare の Rust 製推論エンジン(Infire)と統合されているよ。

もちろんタダじゃないの。現時点では end-to-end で 30〜40% のスループット低下があって、H100 SXM5 で計測したところ、バッチサイズ 1 で最大(約 41%)、バッチ 1024 だと約 30% まで縮まるんだって。原因は主に 3 つ(小バッチの固定コスト、重みタイル再構築の重複、除外している down projection)で、まだ最適化の途中なんだよ。Llama 70B に当てはめると、構成しだいで 18〜28 GB の節約になる見込みだって。

まとめ

  • Unweight は LLM の重みを最大 15〜22% 小さくする、ビット完全一致のロスレス圧縮システム
  • BF16 の指数バイトだけを Huffman 符号化。上位 16 指数で 99% 超をカバーし、必要なのは約 2.6 ビット
  • オンチップ共有メモリで展開してテンソルコアへ直結、特別なハードは不要
  • 代償は 30〜40% のスループット低下(バッチ 1 で約 41%、バッチ 1024 で約 30%)で、まだ改善中
  • 技術ペーパーを公開+ GPU カーネルはオープンソース

GPU メモリのやりくりに頭を悩ませている推論基盤の人や、低レベルな GPU 最適化が大好きな人にはたまらない一本だよ!