Hugging Face KV キャッシュ量子化

Hugging Face は Transformers ライブラリに KV キャッシュ量子化を導入し、キー・バリュー (KV) キャッシュのメモリフットプリントを削減することで、大規模言語モデル (LLM) がより長いテキストシーケンスを生成できるようにしました。この機能により、ユーザーは生成速度をわずかに犠牲にする代わりに、消費者向け GPU でコンテキスト長を拡張し、メモリ効率を大幅に向上させることができます。

自己回帰生成における KV キャッシュの役割

キー・バリュー (KV) キャッシュは、トークンごとにテキストを予測する自己回帰モデルの最適化に不可欠です。新しいトークンを予測するために、モデルはすべての過去トークンからの情報を必要とします。キャッシュがなければ、モデルは各ステップで過去のすべてのトークンに対して行列乗算を再計算しなければなりません。KV キャッシュはメモリバンクとして機能し、自己注意層からのキー・バリューのペアを過去に処理したトークンについて保存し、再計算せずに取得できるようにすることで、テキスト生成を大幅に高速化します。

しかし、コンテキスト長やバッチサイズが増加すると、KV キャッシュはメモリのボトルネックになります。コンテキスト長が 10,000 トークンの 7B Llama-2 モデルの場合、KV キャッシュは float16 精度で約 5GB のメモリを必要とし、これはモデルパラメータ自体が必要とするメモリの約 3 分の 1 に相当します。

KV キャッシュ量子化の技術的実装

KIVI 論文に触発され、Hugging Face の実装はアフィン量子化を利用して KV キャッシュを低精度フォーマットに圧縮します。

量子化手法

この実装はキーとバリューの両方に対してトークン単位の量子化を行います。生成ステップごとに量子化と逆量子化を行うことで生じる潜在的な遅延を軽減するため、Hugging Face は固定サイズの残差キャッシュを採用しています。この残差キャッシュは最新のキーとバリューを元の精度で保持し、キャッシュが最大容量に達した時点で保存された値を量子化し、キャッシュをクリアします。精度を保つためのベースラインとして、残差長さ 128 が使用されます。

サポートされているバックエンドと精度

  • Quanto: int2 と int4 の精度をサポート。
  • HQQ: int2、int4、int8 の精度をサポート。

パフォーマンスと品質のトレードオフ

KV キャッシュの量子化は、メモリ節約、生成速度、モデル品質の間のトレードオフを伴います。

モデル品質と精度

Llama2-7b-chat モデルを PG-19 データセットでテストした結果、int4 キャッシュ精度は fp16 精度とほぼ同等に動作することが示されました。ただし、int2 を使用すると品質が低下します。LongBench ベンチマークでは、Quanto バックエンドを介した int4 精度は、TREC、SAMSum、TriviaQA など複数のデータセットにおいて fp16 と同等かやや上回る性能を示しました。

メモリ効率とレイテンシ

キャッシュを int4 に量子化すると、約 2.5 倍のメモリ節約が得られます。メモリ圧迫は軽減されますが、特にバッチサイズが増加すると生成速度が低下する可能性があります。さらに、KV キャッシュ量子化と重み量子化を組み合わせると、生成速度が約 3 倍低下することがあります。

コンテキスト長の容量

80GB A100 GPU 上で Flash Attention と組み合わせると、KV キャッシュ量子化によりモデルは最大 128k トークンをサポートでき、半精度キャッシュ使用時の最大 40k トークンと比較して大幅に拡張されます。

Transformers との統合

KV キャッシュ量子化はデバイスに依存せず、CPU、GPU、MPS (Apple Silicon) で動作します。🤗 Transformers でこれを実装するには、ユーザーは quanto ライブラリをインストールし、generate 呼び出し時に cache_implementation="quantized" 引数と cache_config 辞書を指定する必要があります。

# Example usage
out = model.generate(
    **inputs, 
    do_sample=False, 
    max_new_tokens=20, 
    cache_implementation="quantized", 
    cache_config={"backend": "quanto", "nbits": 4}
)

Sources