修复: CUDA 内存不足
解决训练或推理时的 "CUDA out of memory" 错误
错误信息
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB
或
RuntimeError: CUDA error: out of memory
或
tensorflow.python.framework.errors_impl.ResourceExhaustedError: OOM when allocating tensor
根本原因
GPU 显存耗尽。发生原因：
- 模型太大 - 模型参数超出可用显存
- 批量大小过高 - 一次处理的样本太多
- 内存碎片化 - 低效的内存分配模式
- 内存泄漏 - 张量未正确释放
- 多进程竞争 - 其他进程占用 GPU 内存
模型内存需求
| 模型大小 | FP32 | FP16/BF16 | INT8 | 推荐 GPU |
|---|---|---|---|---|
| 7B 参数 | ~28GB | ~14GB | ~7GB | T4 16GB, A10G 24GB |
| 13B 参数 | ~52GB | ~26GB | ~13GB | A10G 24GB, L4 24GB |
| 70B 参数 | ~280GB | ~140GB | ~70GB | A100 80GB, H100 80GB |
* 仅推理。训练需要 2-3 倍内存用于梯度和优化器状态。
解决方案
1. 减小批量大小
最快的修复方法 - 将批量大小减半直到能运行：
# 不用 batch_size=32，尝试 batch_size=16 或 batch_size=8
2. 使用混合精度 (FP16/BF16)
将内存使用减半，精度损失极小：
# PyTorch
from torch.cuda.amp import autocast
with autocast():
output = model(input)
# 或在 Ampere+ GPU 上使用 BF16
model = model.to(torch.bfloat16)
3. 梯度检查点
用计算换内存：
# PyTorch
model.gradient_checkpointing_enable()
# Transformers
model = AutoModel.from_pretrained(..., gradient_checkpointing=True)
4. 使用量化 (INT8/INT4)
将模型大小减少 2-4 倍：
# bitsandbytes 8位
model = AutoModelForCausalLM.from_pretrained(
"model_name",
load_in_8bit=True,
device_map="auto"
)
5. 定期清理缓存
释放碎片化内存：
import torch
torch.cuda.empty_cache()
import gc
gc.collect()
6. 监控 GPU 内存
检查内存使用情况：
# 终端
watch -n 1 nvidia-smi
# Python
print(torch.cuda.memory_summary())
vLLM 专用解决方案
限制 GPU 内存使用
python -m vllm.entrypoints.openai.api_server --model meta-llama/Llama-3.2-3B --gpu-memory-utilization 0.8
大模型使用张量并行
python -m vllm.entrypoints.openai.api_server --model meta-llama/Llama-3.1-70B --tensor-parallel-size 2
高性能 GPU 与 AI 云服务器
为您的 Docker 容器提供强大的 NVIDIA 算力支持，支持 A100/H100，全球 32 个机房可选。
- 支持 NVIDIA A100/H100 GPU 实例
- 按小时计费，测试成本低至 $0.004/h
- 全球 32+ 数据中心，极低访问延迟
- 一键运行容器化应用与裸金属服务器