Hugging Face、LLM.int8() 8ビット行列乗算を Transformers と Accelerate に統合

TL;DR

Hugging Face は、8ビット LLM.int8() 量子化手法が transformers と accelerate ライブラリに完全に統合されたことを発表しました。これにより、BLOOM‑176B のような巨大モデルでも、メモリ使用量を約半分に抑えつつ、精度の測定可能な低下なしに推論できるようになります。

なぜ大規模言語モデルに 8 ビット量子化が重要なのか

大規模言語モデル（LLM）は現在、数千億パラメータを超えています（例：PaLM 540B、OPT 176B、BLOOM 176B）。フル精度 FP32 でモデルを保存すると、重みあたり 4 バイトが必要となり、数百ギガバイトというメモリ要件になります——ほとんどの GPU の容量をはるかに超えます。半精度（FP16/BF16）に精度を落とすとメモリは半減しますが、BLOOM 176B でも約 350 GB が必要です。8 ビット整数（INT8）量子化を使えばさらに 2 倍の削減が可能ですが、従来の素朴な量子化は特に 6 B パラメータを超えるモデルで精度が低下しがちでした。

LLM.int8() の核心アイデア：精度低下なしの行列乗算

LLM.int8() は アウトライヤ（外れ値）を別扱いすることで精度低下を防ぎます。

  1. アウトライヤ抽出 – 大きさがしきい値（≈6）を超える値を、隠れ状態行列の列ごとに特定します。
  2. 混合精度 matmul – アウトライヤは FP16 で乗算し、残りの大部分は INT8 に量子化してベクトル単位（アクティベーションは行単位、重みは列単位）の量子化で乗算します。
  3. デ量子化と集約 – INT8 の結果を FP16 にデ量子化し、アウトライヤの FP16 結果と加算して最終的な FP16 出力を得ます。

この 3 ステップのプロセスにより、元の FP16/BF16 モデルと同等の推論品質を保ちつつ、メモリ使用量を 1/4 に削減できます。

量子化の仕組み：ゼロポイント vs. absmax

  • ゼロポイント量子化 は浮動小数点範囲（例：[-1, 1]）を INT8 範囲 [-127, 127] にスケーリングし、各値を丸めます。逆スケーリングで元の値の近似を復元します。
  • absmax 量子化 は各テンソルをその絶対最大値で割り、127 を掛けて丸めます。ベクトル [1.2, ‑0.5, ‑4.3, …, 5.4] の場合、スケーリング係数は 127/5.4 ≈ 23.5 となり、整数値は [-127, 127] に収まります。

どちらの方式も行単位または列単位で適用でき、スケールの大きい行列乗算を正確に行うために重要です。

ゼロ の精度低下を示す実証結果

lm‑eval‑harness を用いた OPT‑175B と BLOOM‑176B のベンチマークでは、INT8 と FP16/BF16 のスコア差がすべてのタスクで標準誤差以下であることが示されています（例：HellaSwag 正確度 0.7849 vs. 0.7849、Lambada パープレキシティ 3.0142 vs. 3.0152）。あるケース（BLOOM‑176B の Lambada）では、INT8 モデルがわずかに優れた結果を示しました。論文 LLM.int8(): 8‑bit Matrix Multiplication for Transformers at Scale に完全な評価が掲載されています。

スピードのトレードオフ

メモリ削減は、最大規模モデルでは若干の遅延を伴います：BLOOM‑176B は INT8 で FP16 に比べて 15 %–23 % 遅くなります。小規模モデル（例：T5‑3B、T5‑11B）は当初は大きな遅延がありましたが、最近の最適化によりトークンあたりのレイテンシが 312 ms から 173 ms（T5‑3B）、45 ms から 25 ms（T5‑11B）へと改善されました。将来のリリースでさらにギャップを埋めることが目指されています。

モデル 精度 GPU トークン / ms (バッチ 1)
BLOOM‑176B BF16 8 × A100 80GB 239
BLOOM‑176B INT8 4 × A100 80GB 282
T5‑11B FP16 2 × T4 15GB 11.7
T5‑11B INT8 1 × T4 15GB 43.5

transformers への統合

主要コンポーネントは bitsandbytes.nn.Linear8bitLt で、torch.nn.Linear のドロップイン置換として機能します。最小限の変換ワークフローは次のとおりです。

import torch, bitsandbytes as bnb
from bnb.nn import Linear8bitLt

# FP16 モデルを定義し、重みを保存
fp16 = torch.nn.Sequential(torch.nn.Linear(64, 64), torch.nn.Linear(64, 64))
torch.save(fp16.state_dict(), "model.pt")

# INT8 バージョンを構築
int8 = torch.nn.Sequential(
    Linear8bitLt(64, 64, has_fp16_weights=False),
    Linear8bitLt(64, 64, has_fp16_weights=False),
)
int8.load_state_dict(torch.load("model.pt"))
int8 = int8.to(0)   # GPU 上で量子化が実行される

.to 呼び出し後、重みは [-127, 127] の範囲の int8 テンソルとして格納されます。元の FP16 値は (weight.CB * weight.SCB) / 127 で復元できます。

accelerate を使ったゼロメモリモデル構築

accelerate.init_empty_weights() は meta デバイス上にモデルを作成し、RAM を割り当てません。この統合により、パラメータが meta デバイスから外れたときにカスタムクラス（Int8Params）が保持されます。再帰的ヘルパーはすべての nn.Linear を Linear8bitLt に置き換え、lm_head のようにフル精度で残すべきモジュールはそのままにします。

from accelerate import init_empty_weights
import torch.nn as nn, bitsandbytes as bnb

def replace_8bit_linear(model, threshold=6.0, exclude="lm_head"):
    for name, module in model.named_children():
        if list(module.children()):
            replace_8bit_linear(module, threshold, exclude)
        if isinstance(module, nn.Linear) and name != exclude:
            with init_empty_weights():
                model._modules[name] = bnb.nn.Linear8bitLt(
                    module.in_features,
                    module.out_features,
                    module.bias is not None,
                    has_fp16_weights=False,
                    threshold=threshold,
                )
    return model

accelerate への 2 つの PR により、各 INT8 テンソルに対して set_module_tensor_to_device が正確に 1 回だけ呼び出され、二重量子化バグが回避されます。

ハードウェアとインストール要件

  • GPU サポート – INT8 テンソルコアが必要です（NVIDIA Turing、Ampere、RTX 20/30、A40‑A100、T4）。CPU や旧世代 Kepler GPU はネイティブサポートがありません。
  • インストール – Python ≥ 3.8 が前提です。
pip install accelerate bitsandbytes
pip install git+https://github.com/huggingface/transformers.git

デモンストレーション

Google Colab ノートブックでは、T5‑11B（元は FP32 で 42 GB）を INT8 でわずか 11 GB に抑えて実行できる例や、単一 T4 で快適に動作する BLOOM‑3B デモが紹介されています。

今後の課題と制限事項

  • 小規模モデルの速度 – 6 B 以下のモデルで INT8 のレイテンシを FP16 と同等にする作業が進行中です。
  • Kepler GPU のサポート – INT8 テンソルコアを持たない GPU（例：GTX 1080）向けに別途ソフトウェアスタックを追加する計画があります。
  • state‑dict の永続化 – 現在の INT8 チェックポイントは量子化統計（CB、SCB）を含まないため、Hub から直接ロードできません。メタデータの追加が優先課題です。
  • CPU 実行 – CPU には 8 ビットテンソルコアがなく、将来的にソフトウェア経路で対応できればアクセシビリティが向上します。
  • テキスト以外への展開 – 大規模ビジョン、オーディオ、マルチモーダルモデルへの適用は未踏の研究領域です。

Credits: Younes B., Tim Dettmers, and contributors listed in the original blog post.

Sources