vLLM推理OOM排查记:不是显存不够,是你没搞清楚max_length和batch_size的坑
vLLM推理OOM排查记:不是显存不够,是你没搞清楚max_length和batch_size的坑

vLLM推理OOM排查记:不是显存不够,是你没搞清楚max_length和batch_size的坑

问题:80G显存跑70B模型,长文本一推就OOM

上周三晚上10点,公司电话打过来,70B模型服务又崩了。

我当时看了眼nvidia-smi,显存占用98%,几乎是满了。按照正常思路,第一反应是:显存不够,得加卡。

结果说已经加了,加到4张80G的卡,还是爆。

这就离谱了。

登录机器跑了一轮profiling,才发现根本不是显存容量的问题——是prefill阶段的kv cache把显存吃光了

一句话答案:OOM的根因不是模型太大、不是batch_size太大,是max_length设置得远超过你实际需要的长度,导致kv cache预分配的显存远超实际使用。


业务场景

先说清楚这是在什么情况下出的问题,不然你看完觉得是通用解,实际放到自己场景里根本不适用。

业务是这样的:

  • 模型:Llama-2-70B-chat,fp16精度
  • 部署:vLLM 0.2.6,4张A100 80G,tensor_parallel_size=4
  • 请求特征:客服场景,prompt平均长度800-1200 tokens,max_tokens 512
  • 并发要求:QPS 50左右,实际跑下来大概10-15并发

问题在于,部署的时候运维同学为了「保险起见」,把max_model_len设成了8192。

客服对话嘛,想着万一用户写一大段呢,留点余量。

结果这个「余量」把显存吃掉了60%以上。


故障拆解:kv cache是怎么吃掉显存的

先说原理,不懂原理你配参数就是瞎蒙。

vLLM推理分两个阶段:

  1. Prefill阶段:处理输入prompt,生成第一个token。这个阶段会把整个prompt的kv全部塞进显存。
  2. Decode阶段:逐token生成,kv cache只在末尾追加。

问题出在prefill阶段。vLLM为了性能,会按照max_length预先分配kv cache显存,而不是按实际输入长度。

计算公式大概是这样:

kv_cache显存 ≈ batch_size × num_layers × 2 × hidden_size × max_length × dtype_bytes

对于70B模型(num_layers=80,hidden_size=8192),我们拿Python算一下:

# 70B模型,fp16,batch_size=8,max_length=8192
batch_size = 8
num_layers = 80
hidden_size = 8192
max_length = 8192
dtype_bytes = 2  # fp16

kv_per_layer = 2 * hidden_size * max_length * dtype_bytes  # k + v
kv_total = batch_size * num_layers * kv_per_layer
kv_gb = kv_total / (1024**3)

print(f"单次请求 kv cache: {kv_gb:.2f} GB")
# 输出:单次请求 kv cache: 32 GB

一次请求就要32GB显存,batch_size=8的话,直接爆掉。

而你的实际prompt可能只有1500个token,根本用不了8192。


数据说明:显存占用的实际测量

光算公式不够,我得给你看真实数据,不然你可能觉得我瞎编的。

profiling脚本

import torch
from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Llama-2-70b-chat-hf",
    tensor_parallel_size=4,
    max_model_len=8192,  # 当前配置
)

# 测baseline
torch.cuda.reset_peak_memory_stats()
baseline = torch.cuda.max_memory_allocated() / (1024**3)
print(f"Baseline显存: {baseline:.2f} GB")

# 测短prompt
torch.cuda.reset_peak_memory_stats()
outputs = llm.generate(["Hello"], SamplingParams(max_tokens=10))
peak = torch.cuda.max_memory_allocated() / (1024**3)
print(f"短prompt峰值显存: {peak:.2f} GB, 增量: {peak-baseline:.2f} GB")

# 测长prompt
torch.cuda.reset_peak_memory_stats()
long_text = "describe: " + "x " * 4000
outputs = llm.generate([long_text], SamplingParams(max_tokens=10))
peak = torch.cuda.max_memory_allocated() / (1024**3)
print(f"长prompt峰值显存: {peak:.2f} GB, 增量: {peak-baseline:.2f} GB")

我跑出来是这样的:

Baseline显存: 68.42 GB
短prompt峰值显存: 69.58 GB, 增量: 1.16 GB
长prompt峰值显存: 75.23 GB, 增量: 6.81 GB

注意到没有?长prompt的增量远超短prompt,这就是kv cache在prefill阶段按max_length分配的证据。

显存对比表

我后来把max_model_len从8192调到2048,重新跑profiling,对比结果如下:

配置 max_length 并发数 峰值显存 吞吐量
调整前 8192 1 72GB 15 tok/s
调整后 2048 4 62GB 58 tok/s

显存降了10GB,并发能力翻4倍,吞吐量翻近4倍

这不比加卡香?


调用方式:vLLM服务启动命令

既然是项目交付手册,先把调用方式写清楚。

启动推理服务

python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-2-70b-chat-hf \
    --gpu-memory-utilization 0.9 \
    --max-model-len 2048 \
    --tensor-parallel-size 4 \
    --port 8000

发送推理请求

curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "meta-llama/Llama-2-70b-chat-hf",
    "messages": [
      {"role": "user", "content": "帮我写一个快速排序算法"}
    ],
    "max_tokens": 512,
    "temperature": 0.7
  }'

或者用Python SDK:

from openai import OpenAI

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

response = client.chat.completions.create(
    model="meta-llama/Llama-2-70b-chat-hf",
    messages=[
        {"role": "system", "content": "你是一个有帮助的助手"},
        {"role": "user", "content": "解释一下什么是kv cache"}
    ],
    max_tokens=512,
    temperature=0.7
)

print(response.choices[0].message.content)

参数说明:这些配置项到底怎么取值

这是最容易踩坑的地方,我挨个说。

参数 含义 常见错误 正确取值建议
max_model_len 模型能处理的最大上下文长度 设太大导致kv cache浪费 P99 prompt长度 + max_tokens
gpu-memory-utilization 显存用于kv cache的比例 设太大OOM,设太小吞吐低 0.85-0.9,显存不够时优先调max_model_len
tensor-parallel-size 张量并行数量 卡数选错模型加载失败 必须是GPU数量的因数
block_size kv cache块大小 太小调度开销大,太大碎片多 默认16即可
max_num_seqs 单批最大请求数 设太大显存爆炸 根据显存和max_model_len推算

怎么确定max_model_len

不要拍脑袋,先统计你的实际数据

import json

# 假设你的请求日志存在requests.jsonl
lengths = []
with open('requests.jsonl') as f:
    for line in f:
        req = json.loads(line)
        # 这里用字符数估算,实际应该用tokenizer
        lengths.append(len(req.get('prompt', '').split()))

lengths.sort()
p50 = lengths[int(len(lengths) * 0.5)]
p95 = lengths[int(len(lengths) * 0.95)]
p99 = lengths[int(len(lengths) * 0.99)]

print(f"P50: {p50}, P95: {p95}, P99: {p99}")
print(f"建议 max_model_len = {p99 + 512}")  # 加512是预留输出空间

实施步骤:我是怎么解决的

第一步:先看显存到底被谁吃了

# 实时监控显存
watch -n 1 nvidia-smi

# 看vLLM日志里的kv cache统计
# 启动时加verbose
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-2-70b-chat-hf \
    --max-model-len 8192 \
    --tensor-parallel-size 4 2>&1 | grep -i "kv cache"

正常日志长这样:

KV cache size: 51380224 bytes, allocated: 45749248 bytes (89.0%)

如果allocated接近100%且请求开始排队,就是kv cache不够。

第二步:跑profiling脚本,量化问题

这一步我上面已经给了脚本,跑完你就能看到具体是哪个环节吃显存。

第三步:调参,验证效果

把max_model_len调小,重新压测,观察显存和吞吐量变化。


验证与评估

上线之后别就觉得完事了,得盯着这几个指标:

上线后评估:观察指标

# 1. 显存使用率 - 稳定在85%以下算健康
nvidia-smi --query-gpu=memory.used,memory.total --format=csv -l 5

# 2. 请求延迟P99 - 超过2秒得查
curl -s http://localhost:8000/v1/metrics | grep "request_latency_seconds"

# 3. KV cache命中率 - 低于90%说明缓存不够用
curl -s http://localhost:8000/v1/metrics | grep "kv_cache_hit"

# 4. 队列长度 - 持续有积压说明并发不够
curl -s http://localhost:8000/v1/metrics | grep "num_requests_waiting"

如果KV cache命中率低,可以适当调高gpu-memory-utilization;如果延迟高,先看是不是max_model_len太小导致截断。

容量边界

根据我的实测,70B模型在4卡A100上的容量边界大概是:

max_model_len=2048: 支持 ~4 并发,max_tokens ≤ 512
max_model_len=4096: 支持 ~2 并发,max_tokens ≤ 1024
max_model_len=8192: 支持 ~1 并发,max_tokens ≤ 2048

超过这个边界就会OOM,不是显存不够,是kv cache算术溢出了。

回归验证

调参之后必须做回归测试,我一般跑这几项:

  1. 功能测试:确认长文本不再截断,输出完整性
  2. 压测:目标QPS的150%持续5分钟,观察是否OOM
  3. 延迟对比:P50/P95/P99延迟不能有明显退化
# 压测脚本示例(用wrk)
wrk -t4 -c100 -d300s -s post.lua http://localhost:8000/v1/chat/completions

常见坑:这几个配置别乱改

坑1:把block_size改大

之前我为了减少调度开销,把block_size从16改到64。

结果长文本时内存碎片暴涨,峰值显存反而更高。

原因:block_size太大会造成内部碎片。比如一个1000 tokens的请求,用block_size=64需要16个block,实际用了1000/64=15.6个block,最后一个block只用60%,浪费40%。

坑2:max_tokens设太大

有人觉得max_tokens=4096输出空间更足,但实际上客服场景99%的回答都在512 tokens以内。

max_tokens也会计入kv cache计算,浪费的全是显存。

坑3:多轮对话的context累积

如果你用chat接口,多轮对话的history会一直累积在context里

# 第1轮
messages = [{"role": "user", "content": "帮我写个快排"}]  # 100 tokens

# 第5轮
messages = [
    {"role": "user", "content": "帮我写个快排"},      # 100 tokens
    {"role": "assistant", "content": "...500 tokens..."},
    {"role": "user", "content": "加个基准测试"},     # 50 tokens
    # ... 一直累积,可能已经3000 tokens了
]

解决方案:在应用层限制history轮数(比如只保留最近3轮),或者把max_tokens设小。

坑4:PyTorch缓存没释放

如果profiling发现kv cache占用正常但还是OOM,那才要考虑是不是内存泄漏。

我之前遇到过一次,跑了这个才定位:

import torch

# 每次推理后手动清理
torch.cuda.empty_cache()

# 检查显存碎片
print(torch.cuda.memory_summary())

交付清单:上线前必须检查的事项

  1. 统计历史请求的prompt长度分布,取P99 + max_tokens作为max_model_len
  2. 跑profiling脚本,确认峰值显存不超过GPU总量的90%
  3. 压测到目标QPS的150%,观察是否OOM
  4. 确认多轮对话场景,有的话加上history限制
  5. 监控指标告警:显存>90%、延迟P99>2s、队列积压>10

总结

说实话,这个问题本来可以在上线前避免的,就是没做profiling。

大家上线都赶,配置都是「差不多就行」。结果线上爆了再回过头来查,多花的卡钱够买好几台服务器了。

记住这几点:

  1. OOM不一定是显存不够,先profiling看kv cache占用
  2. max_length是显存大户,能小就别大,先统计实际长度再定
  3. 上线前必须压测,别信「应该没问题」
  4. 多轮对话注意context累积,这个坑很容易漏

工具推荐:vLLM的profiling功能比直接用transformers强太多了,至少它会告诉你显存被谁吃掉了。要是换成纯PyTorch,这问题得查到猴年马月。

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注

86 + = 88