LLM を 1.58 ビットにファインチューニング：BitNet による極端な量子化

Hugging Face は、BitNet アーキテクチャを用いて既存の大規模言語モデル（LLM）を 1.58 ビット精度にファインチューニングする手法を開発しました。このアプローチにより、モデルはパラメータを -1、0、1 の 3 つの値だけで表現でき、計算コストとエネルギーコストを大幅に削減し、1 ビットモデルをゼロから事前学習する際に通常必要となる膨大な予算を必要としません。

BitNet アーキテクチャと 1.58 ビット量子化

BitNet は、マルチヘッドアテンションとフィードフォワードネットワークの標準 Linear 層を BitLinear 層に置き換えます。これらの層は、重みには三値精度、アクティベーションには 8 ビット精度を使用します。

計算パラダイム

標準的な LLM（例：Llama）が FP16 の加算と乗算に依存しているのとは異なり、BitNet b1.58 は行列乗算に INT8 加算を使用します。この計算の変更により、理論的には Llama のベースラインと比較して行列乗算のエネルギー消費が 71.4 倍削減されます。

Straight-Through Estimator (STE) を用いたトレーニング

三値量子化に使用される round() 関数は微分不可能であるため、BitNet は Straight Through Estimator (STE) を採用しています。STE は丸め操作の勾配を 1 と近似し、まるで恒等関数であるかのように勾配が流れるようにすることで、標準的な勾配ベースの最適化を可能にします。

量子化の仕組み

  • 重み: 対称的な per-tensor 量子化を使用して量子化します。スケールは重み行列の絶対平均値の逆数です。重みはスケーリングされ、丸められ、-1 から 1 の間にクランプされ、再びスケーリングされます。
  • アクティベーション: absmax per-token 量子化を用いて 8 ビット精度に量子化し、値を [-128, 127] の範囲にスケーリングします。アクティベーション量子化の前に Layer Normalization (LN) を適用して出力分散を維持します。

既存モデルを 1.58 ビットにファインチューニング

Hugging Face は Llama 3 8B モデルを 1.58 ビット精度にうまくファインチューニングしました。初期の実験では、BitLinear 層を急に導入すると、モデルが事前学習された情報のほとんどを失い、損失が急上昇することが判明しました。

動的ウォームアップ量子化

事前知識の喪失を防ぐため、Hugging Face は量子化を段階的に導入する動的 $\lambda$（ラムダ）値を実装しました：

$$\lambda = \min\left(\frac{\text{training_step}}{\text{total_training_steps}}, 1\right)$$

元の値と量子化された値の差を $\lambda$ でスケーリングすることで、モデルはフル精度（$\lambda=0$）からフル量子化（$\lambda=1$）へと遷移します。この線形スケジューラにより、収束が改善され、TinyStories データセットでのパープレキシティは約 4 となりました。

スケーリングと汎化

モデルが一般的な知識を保持し、小規模データセットに過剰適合しないように、チームはトレーニングを FineWeb-edu データセットに拡大しました。学習率 1e-4、バッチサイズ 200 万トークンで 100 億トークンを使用した結果、WikiText のパープレキシティは 12.2 を達成しました。

さらに 1000 億トークンにスケーリングしたところ、いくつかの指標では元の Llama 3 8B に近い性能を示したものの、全精度ベースラインに比べて全体的にやや遅れを取っていることが分かりました。

パフォーマンスベンチマークと結果

1.58 ビットアーキテクチャでファインチューニングされたモデルは HF1BitLLM 組織の下でリリースされました。

主な発見

  • 競争力のある性能: 10 億トークンでファインチューニングした結果、1.58 ビット Llama 3 8B モデルは BitNet 7B モデル（100 億トークンで訓練）および FBI LLM（1.26 兆トークンで蒸留）を上回りました。
  • MMLU ベンチマーク: 開発された 8B モデルは MMLU ベンチマークで Llama 1 7B モデルを超えました。
  • モデルサイズ: 重みを int8 テンソルにパックすることで、パラメータ数は 8B から 2.8B に削減されました。

推論最適化とカスタムカーネル

1.58 ビット重みの速度とメモリの利点を実現するため、Hugging Face は行列乗算中に重みをオンザフライでアンパックするカスタム CUDA および Triton カーネルを実装しました。

タイル化行列乗算

メモリ帯域幅のボトルネックと冗長なデータアクセスを克服するため、チームは tiling を使用しました。この手法は行列を GPU の高速共有メモリに収まる小さなサブマトリックス（タイル）に分割し、遅いグローバルメモリアクセスの頻度を減らします。

カーネルベンチマーク

  • Triton vs. Torch: カスタム Triton カーネルは BF16 精度で @torch.compile とほぼ同等の性能を達成しました。
  • BitBlas: チームは、混合精度ソフトウェアライブラリである BitBlas が、カスタム Triton カーネルと Torch の matmul 関数の両方を低精度で上回ることを確認しました。ただし、カーネルのコンパイルに伴う初期ロード時間が長くなるという欠点があります。

Transformers との統合

統合は transformers ライブラリの新しい「bitnet」量子化手法で行われます。標準の Linear 層は専用の BitLinear 層に置き換えられます。API は変更されず、ユーザーは AutoModelForCausalLM.from_pretrained を使用してモデルをロードできます。

Sources