Unweight：品質を犠牲にせずにLLMを22%圧縮する方法
Cloudflareは、GPUメモリ帯域幅をより効率的に使用するために、Unweightを開発しました。これは、モデルの重みを最大22%削減し、ビット単位で正確な出力を維持するロスレス推論時圧縮システムです。BF16指数バイトの冗長性を利用し、ハフマン符号化でMLP重みを圧縮し、高速なオンチップメモリで展開してテンソルコアに直接供給することで、GPUメモリ帯域幅のボトルネックを軽減します。Llama 3.1 8Bでは、推論時に約13%、配布時に約22%のモデルサイズ削減を実現しますが、現在のスループットオーバーヘッドは30〜40%です。
Cloudflareは、世界のインターネットユーザーの95%に対して50ms以内で推論を実行するために、GPUメモリを徹底的に効率的に使用する必要があります。昨年、Rustベースの推論エンジンInfireとモデルスケジューリングプラットフォームOmniを導入しましたが、今度は推論プラットフォームの次の大きなボトルネックであるモデル重みに取り組んでいます。
大規模言語モデルから1つのトークンを生成するには、GPUメモリからすべてのモデル重みを読み取る必要があります。当社が多くのデータセンターで使用するNVIDIA H100 GPUでは、テンソルコアはメモリがデータを提供できる速度よりも約600倍速くデータを処理できるため、ボトルネックは計算ではなくメモリ帯域幅にあります。Unweightはこの問題を解決します。これは、モデル重みを15〜22%縮小しながら、ビット単位で正確な出力を維持するロスレス圧縮システムであり、特別なハードウェアに依存しません。重要なブレークスルーは、高速なオンチップメモリで重みを展開し、テンソルコアに直接供給することで、低速なメインメモリへの余分なラウンドトリップを回避することです。
圧縮がなぜ難しいのか
一般的な量子化技術はロッシー圧縮であり、モデルの品質に予測不能な影響を与えます。Unweightはロスレスを追求し、正確なモデル動作を保持します。ZipNN、Huff-LLM、ZipServなどの既存システムはそれぞれ異なる問題に焦点を当てていますが、どれもCloudflareのニーズ（Hopper GPUでの推論時のロスレス展開とRust推論エンジンへの統合）に適合しません。核心的な課題は、圧縮そのもの（BF16重みの指数バイトは高度に冗長であり、エントロピー符号化が効果的）ではなく、展開速度を推論を遅くしないほど十分に速くすることです。
モデル重みを効果的に圧縮する方法
AIモデルの各数値は16ビットの「ブレインフロート」（BF16）で格納され、符号ビット、指数ビット、仮数ビットから構成されます。符号と仮数は予測不能に変動しますが、指数には規則性があります。訓練済みLLMでは、256の可能な指数値のうち、上位16の値が全重みの99%以上を占め、理論的には約2.6ビットで表現できます。Unweightはこの冗長性を利用し、符号と仮数はそのままにして、指数バイトのみにハフマン符号化を適用し、指数ストリームで約30%の圧縮を実現します。この技術は、MLP重み行列（ゲート、アップ、ダウンプロジェクション）に選択的に適用され、モデルパラメータの約3分の2を占め、トークン生成時のメモリトラフィックを支配します。アテンション重み、埋め込み、レイヤー正規化は圧縮されません。
GPUメモリのボトルネックと4つの実行パイプライン
NVIDIA H100 GPUには2種類のメモリがあります。高帯域幅メモリ（HBM）と共有メモリ（SMEM）です。推論中、各トークンを生成するにはHBMから完全な重み行列を読み取る必要があり、メモリバスがパフォーマンスのボトルネックとなります。Unweightは4つの圧縮実行パイプラインを提供します。
- 完全デコード：元のBF16重みを完全に再構築し、NVIDIA cuBLASライブラリに渡して標準的な行列乗算を実行します。シンプルですが、メインメモリへの書き込みが最も多く、小バッチに適しています。
- 指数のみデコード：指数バイトのみをデコードし、前処理トラフィックを半減させ、再構成型行列乗算を使用します。
- パレットトランスコード：実行時に重みをコンパクトな4ビットインデックス形式にトランスコードし、前処理トラフィックを4分の1に削減します。
- ダイレクトパレット：前処理をスキップし、重みはロード時に4ビット形式に事前トランスコードされ、行列乗算カーネルがリアルタイムでBF16値を再構築します。
単一のパイプラインが常に最適とは限りません。Unweightは自動チューナーを使用してエンドツーエンドのスループットを測定し、重み行列とバッチサイズごとに最適なパイプラインを選択します。
再構成型行列乗算の仕組み
3つのパイプラインは、展開と計算を融合したカスタム行列乗算カーネルを使用します。このカーネルはHBMから圧縮データをロードし、共有メモリ内で元のBF16値を再構築し、テンソルコアに直接供給します。再構築された重みはメインメモリに存在しません。GPUスレッドグループは、プロデューサー（TMAを使用してデータをロード）とコンシューマー（再構築して計算）に分割され、循環バッファを介してパイプライン化されます。
自動チューニングと結果
自動チューナーは候補構成（パイプライン、行列乗算バリアント、SM割り当て）をスイープし、各プロジェクションとバッチサイズに最適な構成を決定します。Llama 3.1 8Bでは、Unweightは推論時に約13%（ゲート/アッププロジェクションのみ圧縮）、配布時に約22%（すべてのMLPプロジェクションを圧縮）のモデルサイズ削減を実現し、すべて100%ビット単位でロスレスです。現在のスループットオーバーヘッドは約30〜40%で、主な原因は小バッチ固定コスト、冗長な重みタイル再構築、未圧縮のダウンプロジェクションであり、現在最適化中です。
今後の方向性
Cloudflareは、ダウンプロジェクションの圧縮（圧縮可能な重みの約3分の1を占める）、カーネルの最適化、Unweightの大規模モデルや混合エキスパートモデルへの拡張を計画しています。Unweightはオープンソース化されており、GPU効率と圧縮研究への貢献を目指しています。