GPUメモリ管理とLLM推論最適化: vLLM、PagedAttention、GPTQ、TensorRT-LLMまで

演算強度 (AI) = FLOP数 / メモリアクセス量 (bytes)

性能上限 = min(Peak FLOPS, Peak Memory BW × AI)

def calc_model_memory_gb(
    num_params: int,       # パラメータ数（例: 70e9）
    dtype_bytes: int = 2,  # FP16=2, FP32=4, INT8=1, INT4=0.5
) -> float:
    """モデル重みのメモリを計算"""
    return (num_params * dtype_bytes) / (1024 ** 3)

# 主要モデルのメモリ（FP16基準）
models = {
    "Llama-3.1-8B":   {"params": 8e9,   "bytes": 2},
    "Llama-3.1-70B":  {"params": 70e9,  "bytes": 2},
    "Llama-3.1-405B": {"params": 405e9, "bytes": 2},
    "Mistral-7B":     {"params": 7e9,   "bytes": 2},
    "Qwen2-72B":      {"params": 72e9,  "bytes": 2},
}

for name, cfg in models.items():
    mem_gb = calc_model_memory_gb(cfg["params"], cfg["bytes"])
    print(f"{name}: {mem_gb:.1f} GB")

def calc_kv_cache_memory_gb(
    num_layers: int,
    num_heads: int,
    head_dim: int,
    seq_len: int,
    batch_size: int,
    dtype_bytes: int = 2,  # FP16
) -> float:
    """
    KVキャッシュメモリを計算
    各レイヤー: 2 (K, V) × num_heads × head_dim × seq_len × batch_size
    """
    kv_per_layer = 2 * num_heads * head_dim * seq_len * batch_size
    total_bytes = kv_per_layer * num_layers * dtype_bytes
    return total_bytes / (1024 ** 3)

# Llama-3.1-70B の例
# layers=80, heads=64（GQA: kv_heads=8）, head_dim=128
kv_mem = calc_kv_cache_memory_gb(
    num_layers=80,
    num_heads=8,      # GQAの場合はkv_headsを使用
    head_dim=128,
    seq_len=4096,
    batch_size=1,
    dtype_bytes=2,
)
print(f"KVキャッシュ (seq=4096, bs=1): {kv_mem:.2f} GB")
# 出力: KVキャッシュ (seq=4096, bs=1): 0.50 GB

# バッチサイズ別のKVキャッシュ
for bs in [1, 4, 8, 16, 32]:
    mem = calc_kv_cache_memory_gb(80, 8, 128, 4096, bs, 2)
    print(f"  batch_size={bs:2d}: {mem:.2f} GB")

def calc_activation_memory_gb(
    hidden_size: int,
    seq_len: int,
    batch_size: int,
    num_layers: int,
    dtype_bytes: int = 2,
) -> float:
    """推論時の活性化メモリを近似計算"""
    # 各レイヤー: attention + FFN の活性化
    # 近似値: 2 × hidden_size × seq_len × batch_size per layer
    bytes_per_layer = 2 * hidden_size * seq_len * batch_size * dtype_bytes
    return (bytes_per_layer * num_layers) / (1024 ** 3)

OS仮想メモリ       → PagedAttention
────────────────────────────────────
仮想ページ         → 論理ブロック（logical block）
物理フレーム        → 物理ブロック（physical block）
ページテーブル      → ブロックテーブル（block table）
ページフォルト      → ブロック割り当て

リクエストA: [Block 0] → [Block 1] → [Block 2]
                                           ↕ 物理ブロック共有（共通プロンプト）
リクエストB: [Block 0] → [Block 1] → [Block 3]

# vLLMのインストール
pip install vllm

# サーバー起動（シングルGPU）
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-3.1-8B-Instruct \
    --dtype bfloat16 \
    --max-model-len 8192 \
    --gpu-memory-utilization 0.90

# OpenAI互換APIの呼び出し
curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "meta-llama/Llama-3.1-8B-Instruct",
    "messages": [{"role": "user", "content": "GPUメモリ最適化について説明して"}],
    "max_tokens": 512,
    "temperature": 0.7
  }'

# PythonクライアントからvLLM APIを呼び出す
from openai import OpenAI

client = OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="dummy",
)

response = client.chat.completions.create(
    model="meta-llama/Llama-3.1-8B-Instruct",
    messages=[
        {"role": "system", "content": "あなたはGPU最適化の専門家です。"},
        {"role": "user", "content": "PagedAttentionの動作原理を説明してください。"},
    ],
    max_tokens=1024,
    stream=True,
)

for chunk in response:
    if chunk.choices[0].delta.content:
        print(chunk.choices[0].delta.content, end="", flush=True)

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

# 4ビットNF4量子化の設定
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,  # 計算はBF16で実行
    bnb_4bit_quant_type="nf4",              # NormalFloat4量子化
    bnb_4bit_use_double_quant=True,         # ダブル量子化で追加圧縮
)

model_id = "meta-llama/Llama-3.1-70B-Instruct"

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",      # 自動マルチGPU分散
    trust_remote_code=True,
)

tokenizer = AutoTokenizer.from_pretrained(model_id)

# メモリ使用量の確認
print(f"GPUメモリ: {torch.cuda.memory_allocated() / 1e9:.2f} GB")

from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
from transformers import AutoTokenizer
import torch

model_id = "meta-llama/Llama-3.1-8B-Instruct"

quantize_config = BaseQuantizeConfig(
    bits=4,              # 4ビット量子化
    group_size=128,      # グループサイズ（小さいほど精度高いがメモリ増加）
    desc_act=False,      # 活性化順序の記述
    damp_percent=0.01,   # Hessianダンピング係数
)

tokenizer = AutoTokenizer.from_pretrained(model_id)

# キャリブレーションデータの準備（代表的なテキストサンプル）
calibration_data = [
    tokenizer("The GPU accelerates machine learning by...", return_tensors="pt").input_ids,
    tokenizer("Quantization reduces model size while...", return_tensors="pt").input_ids,
    # 実際には1024サンプル以上を推奨
]

# モデルのロードと量子化
model = AutoGPTQForCausalLM.from_pretrained(
    model_id,
    quantize_config=quantize_config,
    torch_dtype=torch.float16,
)

model.quantize(calibration_data)
model.save_quantized("llama-3.1-8b-gptq-4bit")
print("GPTQ量子化完了!")

# 量子化モデルのロード
quantized_model = AutoGPTQForCausalLM.from_quantized(
    "llama-3.1-8b-gptq-4bit",
    use_safetensors=True,
    device="cuda:0",
)

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_id = "meta-llama/Llama-3.1-8B-Instruct"
quant_path = "llama-3.1-8b-awq"

# AWQ量子化の設定
quant_config = {
    "zero_point": True,   # ゼロポイント量子化
    "q_group_size": 128,  # グループサイズ
    "w_bit": 4,           # 4ビット
    "version": "GEMM",    # 行列乗算カーネル
}

model = AutoAWQForCausalLM.from_pretrained(
    model_id,
    low_cpu_mem_usage=True,
    use_cache=False,
)
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)

model.quantize(tokenizer, quant_config=quant_config)
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)
print("AWQ量子化完了!")

Static Batching（batch_size=3）:

時間 →
[リクエストA: ████████████░░░░░░░░]  (12トークン生成)
[リクエストB: ████░░░░░░░░░░░░░░░░]  (4トークン生成)
[リクエストC: ████████░░░░░░░░░░░░]  (8トークン生成)
                └─ B、CはAの完了を待つ必要あり（GPU浪費）

Continuous Batching:

Step 1: [A1][B1][C1]  ← 3件を同時処理
Step 2: [A2][B2][C2]
Step 3: [A3][B3][C3]  ← B完了、新規リクエストDを追加
Step 4: [A4][C4][D1]  ← 空きスロットを即座に埋める
Step 5: [A5][C5][D2]
Step 6: [A6][C6][D3]  ← C完了、新規リクエストEを追加
...

from vllm import LLM, SamplingParams

# テンソル並列で4つのGPUに分散
llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    tensor_parallel_size=4,        # 4 GPUにテンソル並列分散
    dtype="bfloat16",
    max_model_len=8192,
    gpu_memory_utilization=0.90,
    enforce_eager=False,           # CUDAグラフ最適化を使用
)

sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    max_tokens=512,
    stop=["</s>", "[INST]"],
)

prompts = [
    "GPUメモリ階層を説明して",
    "PagedAttentionの利点は?",
    "量子化手法を比較して",
]

outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated = output.outputs[0].text
    print(f"プロンプト: {prompt[:50]}...")
    print(f"生成: {generated[:100]}...")
    print()

Attentionヘッドの分散（4-wayテンソル並列）:

GPU 0: Head 0〜15
GPU 1: Head 16〜31
GPU 2: Head 32〜47
GPU 3: Head 48〜63

各GPUが独立して計算し、AllReduceで結果を集約

Llama-3.1-70B（80レイヤー）→ 4-wayパイプライン:

GPU 0: Layer 0〜19
GPU 1: Layer 20〜39
GPU 2: Layer 40〜59
GPU 3: Layer 60〜79

レイヤー順に処理し、GPU間でactivationを転送

import torch

def profile_gpu_memory(func, *args, **kwargs):
    """GPUメモリ使用量をプロファイリング"""
    torch.cuda.reset_peak_memory_stats()
    torch.cuda.synchronize()

    before = torch.cuda.memory_allocated()
    result = func(*args, **kwargs)
    torch.cuda.synchronize()

    after = torch.cuda.memory_allocated()
    peak = torch.cuda.max_memory_allocated()

    print(f"メモリ増加: {(after - before) / 1e9:.3f} GB")
    print(f"ピークメモリ: {peak / 1e9:.3f} GB")
    print()
    print(torch.cuda.memory_summary())
    return result

# メモリ統計出力の例
def load_and_infer():
    from transformers import pipeline
    pipe = pipeline(
        "text-generation",
        model="microsoft/phi-2",
        torch_dtype=torch.float16,
        device_map="auto",
    )
    return pipe("GPU memory management is", max_new_tokens=50)

profile_gpu_memory(load_and_infer)

# 最適化されたvLLMサーバー設定例
vllm_config = {
    "model": "meta-llama/Llama-3.1-8B-Instruct",
    "dtype": "bfloat16",
    "tensor_parallel_size": 1,
    "gpu_memory_utilization": 0.90,    # GPUメモリの90%を使用
    "max_model_len": 8192,
    "max_num_batched_tokens": 8192,    # バッチあたりの最大トークン数
    "max_num_seqs": 256,               # 同時処理シーケンス数
    "enable_chunked_prefill": True,    # Chunked prefillを有効化
    "block_size": 16,                  # KVキャッシュブロックサイズ（PagedAttention）
    "swap_space": 4,                   # CPUスワップ領域（GB）
    "enforce_eager": False,            # CUDAグラフを使用
    "disable_log_stats": False,
}