Google Kubernetes Engine(GKE)で GPU を使用して大規模言語モデルの推論を最適化する際のベスト プラクティス


Google Kubernetes Engine(GKE)は、大規模言語モデル(LLM)推論をきめ細かく制御し、パフォーマンスと費用を最適化します。このガイドでは、vLLM および Text Generation Inference(TGI)サービング フレームワークを使用して、GKE の GPU でオープン LLM の推論とサービングを最適化する方法について説明します。

すべてのベスト プラクティスをまとめたチェックリストについては、チェックリストの概要をご覧ください。

目標

このガイドは、生成 AI をご利用のお客様、GKE の新規または既存のユーザー、ML エンジニア、LLMOps(DevOps)エンジニアで、Kubernetes で GPU を使用して LLM ワークロードを最適化することに関心のある方を対象としています。

このガイドでは、以下の内容を学習します。

  • 量子化、テンソル並列処理、メモリ最適化など、トレーニング後の LLM 最適化手法を選択する。
  • これらの最適化手法を検討する際に、大まかなトレードオフを検討する。
  • 最適化設定を有効にして、vLLM や TGI などのサービング フレームワークを使用して、オープン LLM モデルを GKE にデプロイする。

LLM サービング最適化手法の概要

AI 以外のワークロードとは異なり、LLM ワークロードは通常、行列乗算に依存するために、レイテンシが長く、スループットが低くなります。LLM 推論のパフォーマンスを向上させるには、専用のハードウェア アクセラレータ(GPU や TPU など)と最適化されたサービング フレームワークを使用できます。

次のベスト プラクティスの 1 つ以上を適用すると、LLM ワークロードのレイテンシを短縮しながら、スループットと費用対効果を改善できます。

このガイドに記載されている例では、Gemma 7B LLM と vLLM または TGI サービング フレームワークを使用して、これらのベスト プラクティスを適用しています。ただし、説明されているコンセプトと機能は、最も一般的なオープン LLM に適用できます。

始める前に

このガイドの例を試してみる前に、次の前提条件の作業を完了しておいてください。

  1. 次のガイドの手順に沿って Gemma モデルにアクセスし、環境を準備し、Google Cloud リソースを作成して構成します。

    必ず Hugging Face アクセス トークンを Kubernetes Secret に保存してください。

  2. https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/ サンプル リポジトリのクローンをローカル開発環境に作成します。

  3. 作業ディレクトリを /kubernetes-engine-samples/ai-ml/llm-serving-gemma/ に変更します。

ベスト プラクティス: 量子化

量子化とは、不可逆画像圧縮に似た手法で、重みを精度の低い形式(8 ビットまたは 4 ビット)で表すことでモデルサイズを削減する手法であり、これによりメモリ要件が削減されます。ただし画像圧縮と同様に、量子化にはトレードオフがあります。モデルサイズを小さくすると、精度が低下する可能性があります。

さまざまな量子化方法があり、それぞれに独自の長所と短所があります。AWQ や GPTQ などの一部の量子化方法は、事前量子化を必要とし、Hugging FaceKaggle などのプラットフォームで利用できます。たとえば、Llama-2 13B モデルに GPTQ を適用し、Gemma 7B モデルに AWQ を適用すると、量子化なしで 2 つの L4 GPU ではなく、1 つの L4 GPU でモデルを提供できます。

AutoAWQAutoGPTQ などのツールを使用して量子化を実行することもできます。これらの方法ではレイテンシとスループットを改善できます。一方、EETQbitsandbytes ライブラリを使用した量子化手法では、事前量子化されたモデルは必要ないため、事前量子化バージョンを使用できない場合に適しています。

使用する最適な量子化手法は、各自の目標と、使用するサービング フレームワークとその手法の互換性に応じて決まります。詳細については、Hugging Face の量子化ガイドをご覧ください。

次のいずれかのタブを選択して、TGI または vLLM フレームワークを使用して量子化を適用する例をご確認ください。

TGI

GKE は、TGI での次の量子化オプションをサポートしています。

  • awq
  • gptq
  • eetq
  • bitsandbytes
  • bitsandbytes-nf4
  • bitsandbytes-fp4

AWQ および GPTQ 量子化方法では事前量子化されたモデルが必要ですが、EETQ および bitsandbytes 量子化はあらゆるモデルに適用できます。これらのオプションの詳細については、Hugging Face の記事をご覧ください。

量子化を使用するには、モデルサーバーの起動時に -–quantize パラメータを設定します。

次のスニペットは、GKE で TGI を使用して bitsandbytes 量子化で Gemma 7B を最適化する方法を示しています。

args:
- --model-id=$(MODEL_ID)
- --num-shard=2
- --quantize=bitsandbytes

この構成を適用するには、次のコマンドを使用します。

kubectl apply -f tgi/tgi-7b-bitsandbytes.yaml

vLLM

GKE は、vLLM での次の量子化オプションをサポートしています。

vLLM でモデル量子化を使用するには、モデルを事前に量子化しておく必要があります。ランタイムを起動するときに、–quantization パラメータを設定します。

次のスニペットは、GKE で vLLM を使用して awq 量子化で Gemma 7B モデルを最適化する方法を示しています。

args:
- --model=$(MODEL_ID)
- --tensor-parallel-size=1
- --quantization=awq
env:
- name: MODEL_ID
  value: google/gemma-7b-AWQ
resources:
  requests:
    nvidia.com/gpu: 1
  limits:
    nvidia.com/gpu: 1

この構成を適用するには、次のコマンドを使用します。

kubectl apply -f vllm/vllm-7b-awq.yaml

KV キャッシュ量子化を使用してレイテンシを改善する

FP8 E5M2 KV キャッシュ量子化を使用すると、特に大きなバッチサイズで、KV キャッシュのメモリ フットプリントを大幅に削減し、レイテンシを改善できます。ただし、これに伴い推論の精度が低下します。

FP8 E5M2 KV キャッシュ量子化を有効にするには、パラメータ --kv-cache-dtype fp8_e5m2 を設定します。

args:
- --model=$(MODEL_ID)
- --tensor-parallel-size=1
- --kv-cache-dtype=fp8_e5m2
- --max-model-len=1200
resources:
  requests:
    cpu: "2"
    memory: "25Gi"
    ephemeral-storage: "25Gi"
    nvidia.com/gpu: 1
  limits:
    cpu: "2"
    memory: "25Gi"
    ephemeral-storage: "25Gi"
    nvidia.com/gpu: 1

この構成を適用するには、次のコマンドを使用します。

kubectl apply -f vllm/vllm-7b-kvcache.yaml

ベスト プラクティス: テンソル並列処理

テンソル並列処理は、計算の負荷を複数の GPU に分散する手法です。これは、単一の GPU メモリ容量を超える大規模なモデルを実行する場合に不可欠です。この手法では、高価な GPU を 1 つ使用する代わりに、手頃なコストの GPU を複数使用できるため、費用対効果が高くなります。また、モデル推論のスループットを向上できる可能性もあります。テンソル並列処理は、テンソル オペレーションを小さなデータチャンクで独立して実行できることを利用します。

この手法の詳細については、Hugging Face のテンソル並列処理ガイドをご覧ください。

次のいずれかのタブを選択して、TGI または vLLM フレームワークを使用してテンソル並列処理を適用する例をご確認ください。

TGI

TGI を使用すると、サービング ランタイムはデフォルトで Pod で使用可能なすべての GPU を使用します。使用する GPU の数を設定するには、--num-shard パラメータの値として GPU の数を指定します。

テンソル並列処理でサポートされているモデルのリストについては、Hugging Face のドキュメントをご覧ください。

次のスニペットは、テンソル並列処理と 2 つの L4 GPU を使用して Gemma 7B 指示チューニング型モデルを最適化する方法を示しています。

args:
- --model-id=$(MODEL_ID)
- --num-shard=2

この構成を適用するには、次のコマンドを使用します。

kubectl apply -f tgi/tgi-7b-it-tensorparallelism.yaml

GKE Autopilot クラスタでこのコマンドを実行すると、21 個の vCPU と 78 GiB のメモリの最小リソース要件を持つ Pod が作成されます。

vLLM

vLLM は分散テンソル並列推論をサポートしています。利用可能な GPU が複数ある場合、vLLM はデフォルトでこの機能を有効にします。

次のスニペットは、テンソル並列処理と 2 つの L4 GPU を使用して Gemma 7B 指示チューニング型モデルを最適化する方法を示しています。

args:
- --model=$(MODEL_ID)
- --tensor-parallel-size=2

この構成を適用するには、次のコマンドを使用します。

kubectl apply -f vllm/vllm-7b-it-tensorparallelism.yaml

GKE Autopilot クラスタでこのコマンドを実行すると、21 個の vCPU と 78 GiB のメモリの最小リソース要件を持つ Pod が作成されます。

ベスト プラクティス: モデルのメモリ最適化

LLM のメモリ使用量を最適化することは、効率的な推論に不可欠です。このセクションでは、Paged Attention や Flash Attention などのアテンション レイヤの最適化戦略について説明します。これらの戦略によりメモリ効率が向上し、入力シーケンスを長くして GPU のアイドル時間を短縮できます。このセクションでは、メモリ制約に合わせてモデルの入力サイズと出力サイズを調整し、特定のサービング フレームワークに合わせて最適化する方法についても説明します。

アテンション レイヤの最適化

単語の意味がコンテキストに応じて変化することがあるため、セルフ アテンション レイヤを使用すると、モデルが言語処理タスクのコンテキストを理解できるようになります。ただしこれらのレイヤは、入力トークンの重み、キー(K)、値(V)を GPU vRAM に格納します。したがって、入力シーケンスが長くなると、サイズと計算時間が二次的に増加します。

長い入力シーケンスを処理する場合、セルフ アテンションのオーバーヘッドが大幅に増加する可能性があるため、KV キャッシュを使用すると特に便利です。この最適化アプローチでは、コンピューティング処理を線形の複雑さにまで低減します。

LLM のアテンション メカニズムを最適化するための具体的な手法には、次のものがあります。

  • Paged Attention: Paged Attentionは、OS 仮想メモリと同様に、ページング手法を使用して大規模なモデルと長い入力シーケンスのメモリ管理を改善します。これにより、KV キャッシュの断片化と重複が効果的に軽減され、GPU メモリを使い果たすことなく長い入力シーケンスを使用できるようになります。
  • Flash Attention: Flash Attention は、トークン生成中に GPU RAM と L1 キャッシュ間のデータ転送を最小限に抑えることで、GPU メモリのボトルネックを軽減します。この手法では、コンピューティング コアのアイドル時間が解消され、GPU の推論とトレーニングのパフォーマンスが大幅に向上します。

モデルの入力と出力のサイズのチューニング

メモリ要件は、入力と出力のサイズによって異なります。出力が長くコンテキストが多い場合は、より多くのリソースが必要になります。一方、出力が短くコンテキストが少ない場合は、小さく低コストの GPU を使用してコストを節約できます。

次のいずれかのタブを選択して、TGI フレームワークと vLLM フレームワークでモデルの入力メモリと出力メモリ要件をチューニングする例をご確認ください。

TGI

TGI サービング ランタイムは起動時にメモリ要件をチェックし、許容可能な最大モデル メモリ フットプリントが使用可能な GPU メモリに収まらない場合は起動しません。このチェックにより、メモリ使用量の多いワークロードでのメモリ不足(OOM)によるクラッシュを回避できます。

GKE は、モデルのメモリ要件を最適化するために次の TGI パラメータをサポートしています。

次のスニペットは、パラメータ設定 --max-total-tokens=3072, --max-batch-prefill-tokens=512, --max-input-length=512 を使用して、1 つの L4 GPU で Gemma 7B 指示チューニング型モデルを提供する方法を示しています。

args:
- --model-id=$(MODEL_ID)
- --num-shard=1
- --max-total-tokens=3072 
- --max-batch-prefill-tokens=512
- --max-input-length=512
env:
- name: MODEL_ID
  value: google/gemma-7b

この構成を適用するには、次のコマンドを使用します。

kubectl apply -f tgi/tgi-7b-token.yaml

vLLM

vLLM ではモデルのコンテキスト長を構成します。これは、KV キャッシュ サイズと GPU RAM 要件に直接影響します。コンテキストの長さを短くすると、より手頃なコストの GPU を使用できます。デフォルト値は、モデルが受け入れるトークンの最大数です。必要に応じて、--max-model-len MAX_MODEL_LEN を使用してコンテキストの最大長を制限します。

たとえば、デフォルトのコンテキスト長が 8,192 の Gemma 7B 指示チューニング型モデルは、1 つの NVIDIA L4 GPU のメモリ容量を超えます。L4 にデプロイするには、プロンプトと出力の合計長を制限するため、--max-model-len を 640 未満の値に設定します。この調整により、デフォルトのコンテキスト長が長いにもかかわらず、1 つの L4 GPU でこのモデルを実行できるようになります。

変更したトークン数の上限でデプロイするには、次のスニペットを使用します。

args:
- --model=$(MODEL_ID)
- --tensor-parallel-size=1
- --max-model-len=600

この構成を適用するには、次のコマンドを使用します。

kubectl apply -f vllm/vllm-7b-token.yaml

チェックリストの概要

最適化の目標 プラクティス
レイテンシ
  • 強力な GPU を優先する: コンピューティング能力とメモリ I/O スループットが高い GPU へのアップグレードを検討してください。
  • 量子化を検討する: AWQ などの手法でレイテンシを改善できますが、精度が低下くなるトレードオフの可能性があるのでご注意ください。
スループット
  • 水平方向にスケーリングする: サービスを提供するレプリカ(Pod)の数を増やしてワークロードを分散します。
  • テンソル並列処理を使用する: 単一の GPU の容量を超える大規模なモデルの場合は、テンソル並列処理を使用して計算を複数の GPU に分散します。小規模なモデルの場合は、オーバーヘッドを回避するために、テンソル並列処理を 1 に設定した複数のレプリカを検討してください。
  • リクエストを一括処理して量子化する: リクエストを組み合わせて、許容できる精度を維持する量子化手法を探します。
費用対効果
  • 小さいモデルを選択する: ファミリー内でリソースの制約と予算に合ったモデルを選択します。
  • 量子化を適用する: 量子化を使用してメモリ要件を削減します(特に大規模なモデルを扱う場合)。
  • コンテキストの長さを制限する: コンテキストの長さを制限してメモリ使用量をさらに削減し、より小さく費用対効果の高い GPU での実行を可能にします。

次のステップ