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


Google Kubernetes Engine(GKE)は、最適なパフォーマンスと費用で大規模言語モデル(LLM)推論をきめ細かく制御します。このガイドでは、vLLMText 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 を使用してコンテキストの長さに上限を設定します。

たとえば、Gemma 7B 指示チューニング型モデルは、デフォルトのコンテキストの長さが 8,192 であり、NVIDIA L4 GPU 1 つのメモリ容量を超えます。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 での実行を可能にします。

次のステップ