GPU服务器显存不足怎么办？模型量化、批处理与显存优化方法

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB (GPU 0; 23.70 GiB total capacity; 21.34 GiB already allocated...)——这段报错几乎是每个做 GPU 部署的人的必经之路。它麻烦的地方在于：报错信息里给出的"已分配"和"总容量"往往对不上，看似还有空间却分配失败；而且同样的模型昨天能跑，今天加了个参数就崩了。

解决显存不足（OOM）不能靠"重启试试"，必须先看清楚显存被谁占了，再针对性地选择量化、改批处理、换并行策略或优化 KV Cache。本文按"定位 → 量化 → 批处理 → 并行与卸载 → 框架层优化"的顺序，给出一套完整的排查和解决路径。

一、先定位：显存到底被谁吃了

在动手调参数之前，先用工具把显存占用拆开看。PyTorch 自带的内存统计可以打印出非常细的分配情况：

# 实时查看进程级显存占用
nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv

# 每秒采样一次显存与利用率，观察峰值
nvidia-smi dmon -s mu -d 1 -c 120

# 查看是否有僵尸进程仍占着显存
fuser -v /dev/nvidia*

在 Python 代码里插入下面几行，可以拿到更精确的分配/保留（allocated / reserved）数据：

import torch

# 打印当前显存快照，单位 GB
def mem(tag):
    alloc = torch.cuda.memory_allocated() / 1024**3
    reserved = torch.cuda.memory_reserved() / 1024**3
    print(f"[{tag}] allocated={alloc:.2f}GB reserved={reserved:.2f}GB")

# 打印按行排序的分配明细，定位是哪个张量占大头
print(torch.cuda.memory_summary())

# 打快照并导出，便于离线分析
torch.cuda.memory._dump_snapshot("./oom_snapshot.pickle")

这里要理解 allocated 和 reserved 的区别：PyTorch 的显存分配器（caching allocator）会缓存已释放的显存块而不还给驱动。所以 nvidia-smi 里看到的占用通常等于 reserved，而实际在用的是 allocated。如果两者差距很大，说明有大量显存被"缓存但未使用"，此时报错信息里的 free 数字会具有误导性。

1.1 显存构成的四块账

把显存拆成四部分来算，问题基本就定位了：

总显存 ≈ 模型权重 + KV Cache（推理）/ 梯度+优化器状态（训练）+ 激活值 + 框架开销

• 权重：参数量 × 每参数字节数，是固定成本。
• KV Cache：随并发数和上下文长度线性增长，推理 OOM 的头号原因。
• 激活值：随 batch size 和序列长度增长，训练时尤其可观。
• 框架开销：CUDA context 约 200～500MB，加上 cuBLAS workspace 和碎片，预留 1～2GB。

二、最有效的一招：模型量化

量化是把权重（有时也包括激活值）从高精度浮点数转成低精度表示。它的收益是直接的线性下降：精度减半，权重显存减半。

2.1 各精度格式对比

2.2 怎么选量化方案

• 有 A100/H100 且追求效果：直接用 BF16，不要量化。数据中心卡显存够，BF16 数值范围与 FP32 一致，不会出现溢出问题。
• 显存差一点（如 24GB 装 13B）：上 INT8，权重从 26GB 降到 13GB，效果损失通常很小。
• 显存差很多（如 24GB 装 70B）：只能 INT4。优先试 AWQ（它保留约 1% 的关键权重不量化，实际效果往往略优于同位数 GPTQ），其次 GPTQ。
• CPU + GPU 混合部署：用 GGUF 格式配合 llama.cpp，可以把部分层卸载到 GPU，其余在内存里跑。
• 不建议无脑降到 Q2_K：低于 4 bit 后困惑度（perplexity）上升明显，生成质量会出现肉眼可见的退化，尤其是代码和数学类任务。

2.3 实操：量化前后显存对比

# 用 vLLM 以 AWQ 量化启动服务，观察显存占用
python -m vllm.entrypoints.api_server \
    --model /data/models/Qwen-14B-AWQ \
    --quantization awq \
    --dtype half \
    --gpu-memory-utilization 0.92 \
    --max-model-len 4096

# 用 llama.cpp 按层卸载到 GPU，-ngl 指定卸载层数
./main -m ./models/llama-2-13b.Q4_K_M.gguf \
       -ngl 40 \
       -c 4096 \
       -n 512

# 加载完成后另开终端确认实际显存
nvidia-smi --query-gpu=memory.used,memory.total --format=csv

提示：--gpu-memory-utilization 0.92 是 vLLM 的重要参数，它决定框架最多使用多少比例显存。设为 0.92 意味着框架会预留 92% 用于权重和 KV Cache，剩下的留给其他进程。如果你在同一张卡上还要跑别的任务，需要相应调低。

三、批处理：把 batch size 和上下文调对

批处理参数直接决定激活值和 KV Cache 的峰值。调错了要么 OOM，要么 GPU 空转。

3.1 推理侧

• max_batch_size / max_num_seqs：控制同时处理多少条请求。这是 KV Cache 的线性乘数，减半它，KV Cache 直接减半。
• max_model_len / 上下文上限：如果业务实际只用 2K，却把上限设成 32K，框架会按上限预留显存，白白浪费。
• chunked prefill：把长 prompt 的 prefill 阶段切成小块处理，避免一次性分配巨大激活值。vLLM 中通过 --max-num-batched-tokens 控制。
• 分页注意力（PagedAttention）：把 KV Cache 按块分配，避免预分配整段连续显存，能显著降低碎片，通常可提升 2～4 倍吞吐。

# vLLM 中限制并发序列数与单次批处理 token 数
python -m vllm.entrypoints.api_server \
    --model /data/models/llama-2-7b \
    --max-num-seqs 8 \
    --max-num-batched-tokens 2048 \
    --max-model-len 2048 \
    --gpu-memory-utilization 0.90

# Triton Inference Server 中调整动态批处理配置片段
# dynamic_batching { max_queue_delay_microseconds: 100 }
# instance_group [ { count: 2, kind: KIND_GPU } ]

3.2 训练侧

训练时 OOM 的根源通常是激活值。有几个成熟手段：

# PyTorch 开启梯度检查点
from torch.utils.checkpoint import checkpoint
model.gradient_checkpointing_enable()

# 或在训练脚本中统一设置
model.config.use_cache = False          # 训练时必须关掉 KV 缓存
training_args.gradient_checkpointing = True
training_args.gradient_accumulation_steps = 8
training_args.per_device_train_batch_size = 2

# 显存碎片优化：改用可扩展的分配策略
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python train.py

gradient_checkpointing 与 use_cache=False 是训练脚本里必须同时设置的两项——后者尤其容易被忽略，开着 KV 缓存做训练会白白多占一大块显存。

四、单卡装不下：并行与卸载

当量化做到极限、批处理也压不下去时，就要考虑横向扩展了。

4.1 三种并行方式

• 张量并行（Tensor Parallel, TP）：把一层内部的矩阵切开分到多卡。通信量大，要求卡间有 NVLink。PCIe 互联下 TP 效率损失严重。适合单机内多卡。
• 流水线并行（Pipeline Parallel, PP）：按层切分，不同卡负责不同层。通信量小，但会有流水线气泡（bubble），需要合理的 micro-batch 划分。
• 序列并行 / 上下文并行：把长序列切开，专门应对超长上下文场景。

4.2 卸载到 CPU 内存与磁盘

这是"实在没办法"的兜底方案：把部分层放在系统内存，甚至 NVMe 上，用时再搬到 GPU。llama.cpp 的 -ngl 参数就是干这个的；DeepSpeed 的 ZeRO-Offload / ZeRO-Infinity 也支持把优化器状态和参数卸载到内存和 NVMe。

代价非常明确：PCIe 带宽（约 63 GB/s @ Gen4 x16）比 HBM 带宽（2039 GB/s ~ 3.35 TB/s）低一到两个数量级，卸载会严重拖慢速度。所以卸载只适合能跑起来比跑得快更重要的场景，不适合在线服务。

五、排查与优化步骤清单

1. 确认 OOM 发生在哪个阶段：加载权重时崩是权重不够；第一个 batch 崩是激活值；跑一段时间后崩是 KV Cache 或内存泄漏。
2. 打印 allocated 与 reserved：判断是真实占用还是分配器碎片，碎片问题用 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True 缓解。
3. 检查是否有残留进程：用 nvidia-smi --query-compute-apps 和 fuser -v /dev/nvidia* 找出占着显存不干活的老进程。
4. 按公式重算显存：权重 + KV + 激活 + 2GB 开销，看理论值是否就超过物理显存。
5. 先降上下文上限：把 max_model_len 调到业务真实需要的值，这一步往往能立刻释放出可观的 KV Cache。
6. 再降并发数：max_num_seqs 减半，观察吞吐是否可接受。
7. 上量化：FP16 → INT8 → INT4，每降一档权重减半，用业务测试集验证效果是否可接受。
8. 训练场景加梯度检查点和梯度累积：并确认 use_cache=False。
9. 启用 PagedAttention 类分页技术：降低碎片，提升可用显存比例。
10. 最后才考虑并行与卸载：先确认卡间有 NVLink，再决定并行策略。

六、常见误区

• 只看 nvidia-smi 的空闲显存就下结论：那是 reserved 口径，实际可分配的可能更少，也可能更多（分配器缓存）。
• 一上来就量化到 4bit：先确认是不是并发和上下文参数设错了，很多时候调参数比量化更有效且无精度损失。
• 忽略显存碎片：长时间运行的服务反复分配释放不同大小的块，会产生大量碎片。定期重启或使用分页机制。
• 多进程共享一张卡却不做显存限制：每个进程都以为自己能用满整卡。用 torch.cuda.set_per_process_memory_fraction(0.45) 显式限制。
• 训练时不关 use_cache：等于白扔一大块显存给用不上的 KV 缓存。

# 限制单进程可用显存比例（多进程共享一张卡时必备）
import torch
torch.cuda.set_per_process_memory_fraction(0.45, device=0)

# 主动释放分配器缓存（注意：不会把显存还给其他进程，只在本进程内复用）
torch.cuda.empty_cache()

# 检查是否有进程残留占用，必要时清理
nvidia-smi --query-compute-apps=pid,used_memory --format=csv
kill -9 

小结

解决 GPU 显存不足，正确的顺序是：先定位再动手。用 memory_summary() 和 nvidia-smi 确认显存构成，再把账算清楚，最后按"调参数 → 量化 → 并行 → 卸载"的顺序逐层施策。

优先级上，调批处理参数零成本且无损，永远应该先试；量化是最有效的减重手段，代价是精度，需要业务侧验证；并行和卸载是最后的手段，会显著增加系统复杂度和通信开销。另外务必记住，推理 OOM 的第一嫌疑人是 KV Cache 而不是权重——尤其在长上下文、高并发场景下，把上下文上限和并发数管住，比换更大的卡划算得多。

本文涉及的框架参数以 vLLM、PyTorch、llama.cpp 的常见版本为准，不同版本参数名可能略有差异，请以对应版本官方文档为准。