エラー解決

【vLLM】PagedAttentionパラメータ最適化でメモリ不足・遅延を解決する実践ガイド

· 14分で読めます · 2 閲覧

問題の概要：vLLM推論時のメモリ不足と遅延

vLLMは、大規模言語モデル(LLM)の推論を高速化するための推論エンジンとして広く利用されています。その核となる技術が、メモリ効率を大幅に向上させる「PagedAttention」です。しかし、PagedAttentionのパラメータ設定を誤ると、以下のような問題が発生し、期待された性能が得られないことがあります。

  • Out of Memory (OOM) エラー: バッチサイズやコンテキスト長を大きくすると、GPUメモリが不足して推論がクラッシュする。
  • 推論速度の低下: トークン生成のスループットが想定より遅く、リソースを十分に活用できていない。
  • 断片化による効率悪化: メモリの断片化が進み、長いシーケンスの処理でパフォーマンスが不安定になる。

具体的なエラーメッセージの例としては、以下のようなものがあります。

RuntimeError: CUDA out of memory. Tried to allocate 2.34 GiB...
# または
vllm.utils.SamplingParamsの使用時に、長いシーケンスで処理が極端に遅くなる。

これらの問題は、PagedAttentionを制御するキーパラメータをアプリケーションの要件（モデルサイズ、バッチサイズ、最大コンテキスト長など）に合わせて最適化することで、多くの場合解決できます。

原因の解説：PagedAttentionの仕組みとパラメータの影響

PagedAttentionは、オペレーティングシステムの仮想メモリとページングの概念を、AttentionのKVキャッシュ管理に応用した画期的なアルゴリズムです。これにより、非連続的なメモリブロックを効率的に利用できるようになりました。

この仕組みを制御する主要なパラメータと、設定ミスが引き起こす問題は以下の通りです。

1. `block_size`（ブロックサイズ）

これはページングの単位となる「ブロック」が保持するトークン数です。デフォルトは16です。

  • 小さすぎる場合: ブロック数が増え、管理オーバーヘッドが大きくなります。メタデータ用のメモリ消費が増え、断片化も起こりやすくなり、結果的に速度が低下する可能性があります。
  • 大きすぎる場合: メモリの内部断片化が発生します。例えば、17トークンのシーケンスにはブロックが2つ必要ですが、2ブロック目は15トークン分が無駄（内部断片化）になります。これによりメモリ使用効率が悪化し、OOMのリスクが高まります。

2. `gpu_memory_utilization`（GPUメモリ使用率）

vLLMエンジンがGPUメモリのうち、KVキャッシュなどのために確保しようとする割合です。デフォルトは0.9（90%）です。

  • 高すぎる場合 (例: 0.95): モデルの重みやアクティベーション用のメモリが不足し、OOMエラーが発生しやすくなります。
  • 低すぎる場合 (例: 0.5): 利用可能なGPUメモリを十分に活用できず、大きなバッチサイズや長いコンテキストを処理する能力が制限され、スループットが低下します。

3. `swap_space`（スワップスペース）

GPUメモリが不足した場合に、KVキャッシュをオフロードするCPUメモリまたはディスクの容量です。デフォルトは4 GiBです。

  • 設定しない/小さすぎる場合: 純粋なGPUメモリのみで処理されるため、大規模な処理ではOOMが発生します。
  • 大きすぎる場合: CPUとGPU間のデータ転送（非常に遅い）が頻繁に発生し、推論レイテンシが劇的に悪化します。あくまで緊急避難的な役割と理解すべきです。

解決方法：PagedAttentionパラメータ最適化のステップバイステップ

以下に、実際のユースケースに合わせてパラメータを調整する手順を示します。

ステップ1: ベースラインの計測

まず、デフォルト設定でのパフォーマンスとリソース使用状況を把握します。`vllm`のエンジン引数を指定せずに実行し、`nvidia-smi`などでGPUメモリ使用量を監視します。

# デフォルト設定での実行例
from vllm import LLM, SamplingParams

llm = LLM(model="meta-llama/Llama-2-7b-chat-hf")
sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=100)
outputs = llm.generate(["こんにちは、AIについて教えてください。"], sampling_params)
print(outputs[0].outputs[0].text)

ステップ2: アプリケーション要件の定義

以下の項目を明確にします。

  • 最大コンテキスト長: 処理する最大のトークン数（入力+生成）。
  • 典型的なバッチサイズ: 同時に処理するリクエスト数。
  • 許容レイテンシ/目標スループット: 性能目標。

ステップ3: パラメータの系統的な調整

要件に基づき、`LLM`クラスの初期化パラメータを調整します。

from vllm import LLM, SamplingParams

# 最適化パラメータを設定したエンジンの初期化例
llm = LLM(
    model="meta-llama/Llama-2-7b-chat-hf",
    max_model_len=8192,          # モデルがサポートする最大長を明示
    gpu_memory_utilization=0.85, # メモリ競合が起きる場合は0.8まで下げてみる
    block_size=32,               # 長いシーケンスを多用する場合は16→32に増加
    swap_space=2,                # 単位はGiB。必要に応じて4や8に。不要なら0。
    enable_prefix_caching=True,  # プロンプトが似通っている場合は有効化で効率UP
)

sampling_params = SamplingParams(max_tokens=200)
prompts = ["日本の首都は？", "フランスの首都は？"] * 5  # バッチサイズ10の例
outputs = llm.generate(prompts, sampling_params)

調整の指針:

  • `block_size`: アプリケーションの平均シーケンス長を考慮します。シーケンス長の変動が大きい場合は16のままが安全です。長いシーケンス（例: 8Kトークン以上）が主で変動が小さい場合は、32や64に増やすことでオーバーヘッドを削減できます。ベンチマークで比較することが重要です。
  • `gpu_memory_utilization`: まずは0.85から始めます。OOMが起きれば0.8, 0.75と下げます。逆にGPUメモリ使用率が低すぎる（`nvidia-smi`で確認）場合は、0.9に上げてより大きなバッチを処理できる可能性があります。
  • `swap_space`: GPUメモリが物理的に不足している場合のみ有効です。例えば、24GB GPUで70Bパラメータモデルを動かす場合などです。設定する場合は、`2`や`4`（GiB）から始め、ディスクI/Oがボトルネックにならないよう注意します。

ステップ4: ベンチマークと検証

変更後に必ずパフォーマンスを計測し、OOMエラーが発生しないか、スループット/レイテンシが改善しているかを確認します。vLLMには組み込みのベンチマークツールが用意されています。

# ベンチマーク実行の例（コマンドライン）
# デフォルト設定
python -m vllm.entrypoints.api_server --model meta-llama/Llama-2-7b-chat-hf &
# 別ターミナルで
python benchmarks/benchmark_throughput.py 
  --backend vllm 
  --model meta-llama/Llama-2-7b-chat-hf 
  --input-len 512 --output-len 128 
  --num-prompts 100

# 最適化パラメータを指定した場合
python -m vllm.entrypoints.api_server 
  --model meta-llama/Llama-2-7b-chat-hf 
  --gpu-memory-utilization 0.85 
  --block-size 32 &
# 同様にベンチマークを実行

ベンチマーク結果から、リクエストのスループット (req/s) と トークンのスループット (tok/s) を比較します。

コード例・コマンド例：具体的なユースケース別設定

ケース1: 長文生成・要約アプリケーション（長いコンテキスト）

入力ドキュメントが長く、生成も長めのユースケース。

llm = LLM(
    model="mistralai/Mixtral-8x7B-Instruct-v0.1",
    max_model_len=32768,         # モデルの最大長を指定
    gpu_memory_utilization=0.82, # 大規模モデルなので少し控えめに
    block_size=32,               # 長いシーケンスに最適化
    enable_prefix_caching=True,
    tensor_parallel_size=2,      # 複数GPUを使用
)
# 長いプロンプトを処理
long_prompt = "以下は長いドキュメントです...（数千トークン）... この要約を生成してください。"
output = llm.generate([long_prompt], SamplingParams(max_tokens=500))

ケース2: チャットボットアプリケーション（短い応答、低レイテンシ）

短い対話が多く、応答速度が重視されるケース。

llm = LLM(
    model="meta-llama/Llama-3-8B-Instruct",
    max_model_len=4096,
    gpu_memory_utilization=0.88, # メモリに余裕を持たせて高速処理
    block_size=16,               # デフォルト。シーケンス長の変動が大きいので安全
    swap_space=0,                # 低レイテンシが命なのでスワップは使わない
)
# 複数の短いチャットリクエストをバッチ処理
chat_prompts = ["おはよう！", "今日の天気は？", "おすすめのレシピは？"]
outputs = llm.generate(chat_prompts, SamplingParams(max_tokens=150, temperature=0.7))

まとめ・補足情報

vLLMのPagedAttentionパラメータ最適化は、単なる数値変更ではなく、ハードウェアリソース（GPUメモリ量）とソフトウェア要件（コンテキスト長、バッチサイズ、レイテンシ）の間の最適なバランス点を見つける作業です。

重要なポイント:

  1. 計測第一: 推測ではなく、ベンチマークツールや`nvidia-smi`での実測データに基づいて判断してください。
  2. `block_size`: シーケンス長の分布に応じて選択。迷ったらデフォルトの16が無難。
  3. `gpu_memory_utilization`: OOMとの戦い。余裕のあるGPUでは高く、ぎりぎりの環境では低く設定。
  4. `swap_space`: パフォーマンス低下の代償を伴う「非常口」。多用しない。
  5. その他の関連パラメータ: `max_model_len`（モデルの真の最大長を正しく設定）、`enable_prefix_caching`（類似プロンプトが多い場合は有効化）、`tensor_parallel_size`/`pipeline_parallel_size`（マルチGPU設定）も全体の効率に影響します。

最適化は反復的なプロセスです。一度設定したら終わりではなく、アプリケーションの利用パターンが変化した際には、再度パラメータを見直すことで、常にvLLMエンジンから最高のパフォーマンスを引き出すことができます。

この記事は役に立ちましたか？