【vLLM】バッチ推論のスループット最大化チューニング：メモリ不足エラーと遅延を解決

問題の概要：vLLMバッチ推論におけるスループット低下とメモリ不足

vLLMは大規模言語モデル(LLM)の高速推論を実現するライブラリですが、バッチ推論（複数の入力テキストを同時に処理）を行う際、以下のようなパフォーマンス課題に直面することがあります。

• スループットの低下: 期待するほど1秒あたりの処理トークン数(TPS)が上がらない。
• Out of Memory (OOM) エラー: GPUメモリが不足し、実行が中断される。
• リクエストレイテンシの増加: 個々のリクエストの応答時間が長くなる。

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

RuntimeError: CUDA out of memory. Tried to allocate 2.34 GiB...
# または
The server is experiencing heavy traffic. Please try again later. (実際にはトラフィックが少ない場合)

これらの問題は、vLLMの強力な機能であるPagedAttentionやメモリ管理を適切に設定せずに使用した場合に発生します。特に、バッチサイズ(`max_num_batched_tokens`, `max_num_seqs`)やKVキャッシュの設定が鍵となります。

原因の解説：スループット低下の根本的な要因

vLLMのバッチ推論におけるボトルネックは、主に以下の3点に集約されます。

1. 不適切なバッチサイズ設定

vLLMは、一度に処理するシーケンス数(`max_num_seqs`)とトークン数(`max_num_batched_tokens`)を動的に管理します。これらの値が小さすぎるとGPUの計算リソースを十分に活用できず、スループットが低下します。逆に大きすぎると、メモリ不足(OOM)を引き起こします。

2. KVキャッシュのメモリ割り当て

LLM推論の主要なメモリ消費者は、KeyとValueのキャッシュ(KVキャッシュ)です。vLLMはPagedAttentionによりこれを効率化しますが、ブロックサイズ(`block_size`)やGPU割り当て比率(`gpu_memory_utilization`)の設定が不適切だと、メモリフラグメンテーションが発生したり、利用可能なメモリを最大限に活用できなかったりします。

3. 推論エンジンのパラメータチューニング不足

`AsyncLLMEngine`や`LLMEngine`を使用する際のパラメータ（`max_model_len`, `enforce_eager`など）は、デフォルト値のままでは特定のハードウェアやモデルサイズに最適化されていません。

解決方法：スループット最大化のためのステップバイステップチューニング

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

まず、現状のパフォーマンスを把握します。vLLMサーバーを起動し、簡単なベンチマークツール（例：`benchmark_throughput.py`）で現状のスループットを計測します。

# vLLM付属のベンチマークスクリプトの例
python -m vllm.entrypoints.benchmark_throughput 
    --model meta-llama/Llama-2-7b-chat-hf 
    --dataset ShareGPT_V3_unfiltered 
    --num-prompts 1000 
    --request-rate 10 
    --output output.json

ステップ2: メモリ設定の最適化

OOMエラーを防ぎつつ最大メモリを活用するため、`–gpu-memory-utilization`と`–swap-space`を調整します。一般的に、`gpu_memory_utilization`は0.9（90%）程度から始め、OOMが発生すれば少し下げます。

# vLLMサーバー起動コマンド例（チューニング後）
python -m vllm.entrypoints.api_server 
    --model meta-llama/Llama-2-13b-chat-hf 
    --tensor-parallel-size 1 
    --gpu-memory-utilization 0.88 
    --swap-space 4 
    --max-model-len 4096 
    --enforce-eager  # メモリフラグメンテーション対策（速度は少し低下）
    --block-size 16   # PagedAttentionのブロックサイズ（小さいほど細かい管理）

ステップ3: バッチパラメータの調整

スループット向上の核心は、`–max-num-batched-tokens`と`–max-num-seqs`です。これらの最適値はハードウェアと入力長に依存します。

• `max-num-batched-tokens`: 一度に処理する総トークン数。GPUメモリ容量とモデルサイズから逆算。例：13BモデルでGPUメモリ16GBなら、2048程度から開始。
• `max-num-seqs“: 同時処理可能なリクエスト数の上限。スループットとレイテンシのトレードオフ。通常は`max-num-batched-tokens`を平均入力トークン数で割った値より少し大きめに設定。

# バッチパラメータを指定したサーバー起動
python -m vllm.entrypoints.api_server 
    --model meta-llama/Llama-2-7b-chat-hf 
    --max-num-batched-tokens 2048 
    --max-num-seqs 32 
    --max-model-len 4096

ステップ4: 推論エンジンの選択と詳細設定

大量の非同期リクエストを処理する場合は`AsyncLLMEngine`が適しています。また、`–enforce-eager`モードはメモリフラグメンテーションを防ぎ安定性を高めますが、若干の速度低下を伴います。速度優先の場合はこのフラグを外します。

# AsyncLLMEngineを使用したPythonコード例
from vllm import AsyncLLMEngine, SamplingParams
from vllm.utils import random_uuid
import asyncio

async def main():
    engine = AsyncLLMEngine.from_engine_args(engine_args) # engine_argsに上記パラメータを設定
    prompts = ["日本の首都は？", "AIとは何ですか？"] * 10 # バッチ推論用のプロンプトリスト
    sampling_params = SamplingParams(temperature=0, max_tokens=50)
    
    request_ids = [random_uuid() for _ in prompts]
    outputs = await engine.generate(prompts, sampling_params, request_ids=request_ids)
    # 出力処理...

# 実行
asyncio.run(main())

コード例・コマンド例：総合的なチューニングスクリプト

以下は、チューニングパラメータをまとめて検証するためのPythonスクリプトの例です。

import subprocess
import time
import json

def run_vllm_server_with_params(params):
    """パラメータを変えてvLLMサーバーを起動・テストする"""
    cmd = [
        "python", "-m", "vllm.entrypoints.api_server",
        "--model", "meta-llama/Llama-2-7b-chat-hf",
        "--port", "8000",
        "--max-num-batched-tokens", str(params["max_num_batched_tokens"]),
        "--max-num-seqs", str(params["max_num_seqs"]),
        "--gpu-memory-utilization", str(params["gpu_memory_utilization"]),
        "--disable-log-stats", # ログを簡素化
    ]
    if params.get("enforce_eager"):
        cmd.append("--enforce-eager")
    
    # サーバー起動（実際の運用では別プロセス管理が必要）
    print(f"起動コマンド: {' '.join(cmd)}")
    # ここでベンチマークツールを実行し、スループットを計測
    # ...

# チューニングするパラメータの組み合わせ
param_sets = [
    {"max_num_batched_tokens": 1024, "max_num_seqs": 16, "gpu_memory_utilization": 0.85},
    {"max_num_batched_tokens": 2048, "max_num_seqs": 32, "gpu_memory_utilization": 0.88},
    {"max_num_batched_tokens": 4096, "max_num_seqs": 64, "gpu_memory_utilization": 0.90, "enforce_eager": True},
]

for params in param_sets:
    print(f"n=== パラメータセットでテスト: {params} ===")
    run_vllm_server_with_params(params)
    time.sleep(2) # クールダウン

まとめ・補足情報

vLLMのバッチ推論スループットを最大化するには、メモリ管理とバッチ処理のパラメータをハードウェアリソースとワークロードに合わせて継続的にチューニングすることが不可欠です。

主要なチューニングパラメータまとめ

• `–max-num-batched-tokens`: スループットに直結。OOMが発生する手前まで大きくする。
• `–max-num-seqs`: 同時リクエスト数。レイテンシとスループットのバランスを取る。
• `–gpu-memory-utilization`: 0.85〜0.95の範囲で調整。OOM時は下げる。
• `–block-size`: メモリ効率と管理オーバーヘッドのトレードオフ。通常は8, 16, 32が候補。
• `–enforce-eager`: 安定性重視なら有効、最大スループット重視なら無効。

トラブルシューティングのアドバイス

1. 「CUDA out of memory」が発生した場合: まず`–gpu-memory-utilization`を下げ（例: 0.9→0.85）、次に`–max-num-batched-tokens`を減らします。
2. スループットが頭打ちの場合: `–max-num-seqs`を増やし、より多くのリクエストを並行処理できるようにします。また、入力テキストの長さがばらつく場合は、長さでソートしてバッチ化する「パディング」の検討も有効です。
3. レイテンシが長すぎる場合: `–max-num-seqs`を小さくし、個々のリクエストの待ち時間を短縮します。

最終的には、実際の本番環境に近いクエリ分布での負荷テストが最も重要です。vLLMは動的なバッチ処理を得意としていますが、適切なパラメータ設定によりその性能を最大限に引き出すことができます。