Xeon 上の 🤗 Optimum Intel で StarCoder を高速化：Q8/Q4 と投機的デコーディング

記事一覧に戻る

この記事は中国語版もご覧いただけます 简体中文.

はじめに

最近、コード生成モデルが非常に人気を集めています。特に BigCode の StarCoder や Meta AI の Code Llama などの最先端のオープンソースモデルがリリースされたことが背景にあります。Large Language Models (LLMs) をより最適化し、誰でも利用できるようにするための研究が活発に行われています。本記事では、人気のコード生成 LLM である StarCoder を Intel Xeon 上で最適化した最新の成果をお届けします。

StarCoder モデルは、コード補完、バグ修正、コード要約、自然言語記述からのコードスニペット生成など、さまざまなコーディングタスクを支援するために設計された最先端の LLM です。StarCoder モデルは StarCoder ファミリーに属し、StarCoderBase バリアントも含まれます。これらのコード用大規模言語モデル (Code LLMs) は、GitHub の許容ライセンスデータで学習されており、80 以上のプログラミング言語、Git コミット、GitHub の issue、Jupyter ノートブックなどが含まれています。本研究では、8 ビットおよび 4 ビットの量子化と assisted generation を統合することで、Intel 第 4 世代 Xeon 上で StarCoder-15B モデルの推論を 7 倍以上高速化できることを示します。

第 4 世代 Intel Xeon Scalable プロセッサで動作する デモ を Hugging Face Spaces でお試しください。

ステップ 1：ベースラインと評価

ベースラインとして、StarCoder (15B) を PyTorch と Intel Extension for PyTorch (IPEX) と組み合わせて使用します。自動コード補完の品質を評価するためのデータセットは複数存在します。本研究では、人気の HumanEval データセットを使用してモデルの品質と性能を評価します。HumanEval は 164 のプログラミング問題から構成されており、関数シグネチャと docstring の形式で提供され、モデルが関数のコードを完成させます。プロンプトの平均長は 139 です。品質は Bigcode Evaluation Harness を用いて pass@1 メトリックで測定します。性能は HumanEval テストセットで Time To First Token (TTFT) と Time Per Output Token (TPOT) を測定し、平均 TTFT と TPOT を報告します。 第 4 世代 Intel Xeon プロセッサは、Intel® Advanced Matrix Extensions (Intel® AMX) と呼ばれる AI 向けアクセラレーション機能を備えています。具体的には、ディープラーニングのトレーニングおよび推論ワークロードを高速化するために、各コアに BFloat16 (BF16) と Int8 GEMM アクセラレータが内蔵されています。AMX を活用した推論は、PyTorch 2.0 および Intel Extension for PyTorch (IPEX) を通じて導入され、LLM 推論で一般的に使用される各種演算子（例: レイヤー正規化、SoftMax、スケールド・ドット・プロダクト）向けの最適化も含まれています。 ベースラインとして、PyTorch と IPEX に搭載されている既存の最適化を利用し、BF16 モデルで推論を実行します。図 1 にベースラインモデルのレイテンシを示し、表 1 と表 2 にレイテンシと精度を示します。

図 1. ベースラインモデルのレイテンシ。

LLM 量子化

LLM におけるテキスト生成は自己回帰的に行われるため、新しいトークンを生成するたびにモデル全体をメモリから CPU に読み込む必要があります。オフチップメモリ (DRAM) と CPU の間の帯域幅が、トークン生成プロセスにおける最大のボトルネックであることがわかりました。量子化は、この問題を軽減するための一般的な手法です。モデルサイズを削減し、モデル重みの読み込み時間を短縮します。

本研究では、以下の 2 種類の量子化に焦点を当てます：

1. Weight Only Quantization (WOQ) - モデルの重みのみを量子化し、活性化は量子化せず、計算はより高い精度（例: BF16）で行います（逆量子化が必要）。
2. Static Quantization (SQ) - 重みと活性化の両方を量子化します。この量子化プロセスには、キャリブレーションステップを通じて量子化パラメータを事前計算する処理が含まれており、計算をより低い精度（例: INT8）で実行できます。図 2 に INT8 静的量子化の計算プロセスを示します。

ステップ 2：8 ビット量子化 (INT8)

SmoothQuant は、精度の低下を最小限に抑えつつ LLM を INT8 に量子化するための、学習後の量子化アルゴリズムです。静的量子化手法は、活性化の特定チャネルに存在する大きな外れ値のため、LLM では性能が低下することが示されています。活性化はトークン単位で量子化されるため、静的量子化では外れ値が切り捨てられるか、低い振幅の活性化がアンダーフローする結果となります。SmoothQuant アルゴリズムは、活性化と重みの両方に追加のスムージングスケーリングファクタを適用する前量子化フェーズを導入することでこの問題を解決し、活性化の外れ値を平滑化して量子化レベルの利用率を向上させます。

図 2. INT8 静的量子化の計算図。

IPEX を使用して StarCoder モデルに SmoothQuant を適用しました。キャリブレーションデータセットとして MBPP データセットのテスト分割を使用し、Q8-StarCoder を導入しました。評価の結果、Q8-StarCoder はベースラインに対して精度の低下がありませんでした（むしろわずかな改善が見られました）。性能面では、Q8-StarCoder は TTFT で 約 2.19 倍、TPOT で 約 2.20 倍 の高速化を達成しました。図 3 に BF16 ベースラインモデルと比較した Q8-StarCoder のレイテンシ (TPOT) を示します。

図 3. 8 ビット量子化モデルのレイテンシ高速化。

ステップ 3：4 ビット量子化 (INT4)

INT8 は BF16 と比較してモデルサイズを 2 倍削減しますが（重み 1 つあたり 8 ビット vs 16 ビット）、メモリ帯域幅は依然として最大のボトルネックです。メモリからのモデル読み込み時間をさらに短縮するため、モデルの重みを WOQ を用いて 4 ビットに量子化しました。なお、4 ビット WOQ は計算前に 16 ビットへの逆量子化を必要とするため（図 4）、計算オーバーヘッドが発生します。

図 4. INT4 に量子化されたモデルの計算図。

テンソル単位の非対称 Round To Nearest (RTN) 量子化という基本的な WOQ 手法は、精度低下を引き起こす課題を抱えていますが、文献 (Zhewei Yao, 2022) では、モデルの重みをグループ単位で量子化することで精度を維持できることが示されています。精度の劣化を避けるため、入力チャネルに沿って連続する値のグループ（例: 128）ごとに 4 ビット量子化を行い、各グループごとにスケーリングファクタを計算しました。グループ単位の 4 ビット RTN は、HumanEval データセットでの StarCoder の精度を維持するのに十分であることがわかりました。4 ビットモデルは BF16 ベースラインと比較して TPOT で 3.35 倍 の高速化を達成しましたが（図 5）、計算前に 4 ビットから 16 ビットへの逆量子化が必要なため、TTFT で 0.84 倍の速度低下が発生しました（表 1）。

図 5. 4 ビット量子化モデルのレイテンシ高速化。

最初のトークン生成と後続トークン生成の異なるボトルネック

最初のトークンを生成する初期ステップでは、入力プロンプト全体を並列処理するため、プロンプト長が長い場合には大きな計算リソースを必要とします。そのため、この段階では計算がボトルネックとなります。したがって、この処理を BF16 から INT8 精度に切り替えることで、ベースライン（および逆量子化に伴う計算オーバーヘッドが発生する 4 ビット WOQ）と比較して性能が向上します。一方、2 ステップ目以降では、生成された残りのトークンを自己回帰的に 1 つずつ生成するため、各新しいトークン生成時にモデルを繰り返しメモリから読み込む必要があります。その結果、ボトルネックは計算量（FLOPS）ではなくメモリ帯域幅となり、INT4 が INT8 や BF16 を上回ります。

ステップ 4：Assisted Generation (AG)

高い推論レイテンシを軽減し、メモリ帯域幅のボトルネック問題を緩和するもう 1 つの手法が Assisted generation (AG) で、speculative decoding の実用的な実装です。AG はメモリ操作と計算操作のバランスを改善することでこの問題に対処します。これは、小さくて高速なアシスタントドラフトモデルが、大きなターゲットモデルと同じトークンを生成することが多いという前提に基づいています。

AG は、小さくて高速なドラフトモデルを使用して K 個の候補トークンを貪欲に生成します。これらの出力トークンははるかに高速に生成されますが、元のターゲットモデルの出力トークンと一致しない場合があります。そこで次のステップで、ターゲットモデルが K 個の候補トークンすべてを単一のフォワードパスで並列に検証します。このプロセスにより、K 個のトークンを自己回帰的に生成するよりも、K 個のトークンを並列にデコードするレイテンシの方が小さいため、デコードが高速化されます。

StarCoder を高速化するため、ドラフトモデルとして bigcode/tiny_starcoder_py を使用しました。このモデルは StarCoder と同様のアーキテクチャを共有していますが、パラメータ数は 164M のみで、StarCoder よりも 約 95 倍 小さく、はるかに高速です。さらに高速化するため、ターゲットモデルの量子化に加えて、ドラフトモデルにも量子化を適用しました。ドラフトモデルとターゲットモデルの両方について 8 ビット SmoothQuant と 4 ビット WOQ の量子化を検討した結果、両方のモデルに 8 ビット SmoothQuant を適用した場合が最良の結果となり、TPOT で 約 7.30 倍 の高速化を達成しました（図 6）。

これらの量子化の選択は、以下の観測に基づいています：

1. ドラフトモデルの量子化：164M パラメータの 8 ビット量子化 StarCoder をドラフトモデルとして使用する場合、モデルはほぼ CPU キャッシュに収まります。その結果、各トークンごとにターゲットモデルをオフチップメモリから繰り返し読み込む必要がなくなり、メモリ帯域幅のボトルネックが緩和されます。この場合、メモリボトルネックは存在せず、4 ビット WOQ に量子化された StarCoder-164M と比較して、8 ビットに量子化された StarCoder-164M の方が高速化効果が大きくなります。なお、4 ビット WOQ はメモリフットプリントが小さいためメモリ帯域幅がボトルネックとなる場面で有利ですが、計算前に 4 ビットから 16 ビットへの逆量子化が必要なため計算オーバーヘッドが発生します。
2. ターゲットモデルの量子化：アシスト生成では、ターゲットモデルがドラフトモデルによって生成された K 個のトークンのシーケンスを処理します。K 個のトークンを「標準的な」逐次自己回帰処理ではなく、一度に（並列に）ターゲットモデルにフォワードすることで、ボトルネックはメモリ帯域幅から計算に移行します。そのため、4 ビットモデルを使用する場合の逆量子化に伴う追加の計算オーバーヘッドのため、8 ビット量子化されたターゲットモデルを使用した方が高い高速化効果が得られることがわかりました。

図 6. 最適化モデルのレイテンシ高速化。

表 1: Intel 第 4 世代 Xeon 上での StarCoder モデルの精度とレイテンシの測定結果

生成されたモデルを読み込んで推論を実行するには、optimum-intel の IPEXModelForXxx クラスで AutoModelForXxx クラスを置き換えるだけです。

開始する前に、必要なライブラリがすべてインストールされていることを確認してください：

pip install --upgrade-strategy eager optimum[ipex]

- from transformers import AutoModelForCausalLM
+ from optimum.intel import IPEXModelForCausalLM
  from transformers import AutoTokenizer, pipeline

- model = AutoModelForCausalLM.from_pretrained(model_id)
+ model = IPEXModelForCausalLM.from_pretrained(model_id)
  tokenizer = AutoTokenizer.from_pretrained(model_id)
  pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
  results = pipe("He's a dreadful magician and")