Optimum Intel を使用した Intel Xeon 上の StarCoder の高速化

Hugging Face と Intel は、第4世代 Intel Xeon Scalable プロセッサ上で StarCoder-15B モデルに対して 7 倍以上の推論高速化を実現しました。この性能向上は、optimum-intel ライブラリを介した 8 ビットおよび 4 ビット量子化技術と支援生成(投機的デコード)の統合により実現されています。

推論性能ベンチマーク

量子化と支援生成の統合により、HumanEval データセットでのモデル精度を維持しながら、出力トークンあたりの時間 (TPOT) が大幅に短縮されます。

StarCoder 量子化 精度 HumanEval (pass@1) TTFT (ms) TTFT 速度向上 TPOT (ms) TPOT 速度向上
Baseline None A16W16 33.54 357.9 1.00x 181.0 1.00x
INT8 SmoothQuant A8W8 33.96 163.4 2.19x 82.4 2.20x
INT4 RTN (g128) A16W4 32.80 425.1 0.84x 54.0 3.35x
INT8 + AG SmoothQuant A8W8 33.96 183.6 1.95x 24.8 7.30x

Intel Xeon 用の量子化戦略

StarCoder を最適化するために、チームは自動回帰トークン生成における主要な制約であるメモリ帯域幅のボトルネックに対処するため、2 つの異なる量子化手法を使用しました。

8 ビット 静的量子化 (INT8)

SmoothQuant アルゴリズムを使用して、モデルは最小限の精度損失で INT8 に量子化されます。SmoothQuant は、活性化と重みの両方に平滑化スケーリング係数を適用し、活性化における大きな外れ値の影響を軽減します。このアプローチにより、TPOT が約 2.20 倍、Time To First Token (TTFT) が約 2.19 倍の速度向上が得られました。

4 ビット 重みのみ量子化 (INT4)

メモリフットプリントとロード時間をさらに削減するため、モデルの重みはグループ単位の Round To Nearest (RTN) 量子化(128 グループ)を使用して 4 ビットに量子化されました。これにより TPOT が 3.35 倍の速度向上を達成しましたが、計算前に 4 ビット重みを 16 ビットに逆量子化するための計算オーバーヘッドにより、TTFT が 0.84 倍遅くなりました。

支援生成 (投機的デコード)

支援生成 (AG) は、小さく高速なドラフトモデルを使用して候補トークンを予測し、より大きなターゲットモデルが並行してそれらを検証することで推論を高速化します。

  • ドラフトモデル: チームは bigcode/tiny_starcoder_py (164M パラメータ) を使用しました。これはターゲットの StarCoder-15B モデルより約 95 倍小さいです。
  • 最適構成: 最良の結果(TPOT が約 7.30 倍の速度向上)は、ドラフトモデルとターゲットモデルの両方に 8 ビット SmoothQuant を適用することで得られました。

支援生成における技術的トレードオフ

Assisted Generation において 8 ビット量子化を 4 ビットより選択する理由は、ボトルネックの変化によります:

  1. ドラフトモデル: 8 ビット量子化された 164M パラメータのモデルは、CPU キャッシュ内にほぼ収まり、メモリ帯域幅のボトルネックを排除し、4 ビット WOQ に伴う逆量子化オーバーヘッドを回避します。
  2. ターゲットモデル: ターゲットモデルは複数のトークンを並行して検証するため、ボトルネックがメモリ帯域幅から計算へとシフトします。8 ビット量子化モデルは、計算負荷の高い逆量子化プロセスを回避するため、4 ビットモデルよりも優れています。

Optimum Intel を使用した実装

ユーザーは、標準の AutoModelForCausalLM クラスを optimum-intel ライブラリの IPEXModelForCausalLM に置き換えることで、これらの最適化を実装できます。このライブラリは第4世代 Xeon プロセッサ上で Intel Extension for PyTorch (IPEX) と Intel Advanced Matrix Extensions (AMX) を活用します。

pip install --upgrade-strategy eager optimum[ipex]
from optimum.intel import IPEXModelForCausalLM
from transformers import AutoTokenizer, pipeline

model = IPEXModelForCausalLM.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

Sources