本文介绍 PYTORCH_CUDA_ALLOC_CONF 环境变量的配置方法，帮助你在显存紧张时最大化利用 GPU 资源，减少 OOM（Out of Memory）错误。

📌 问题背景

你是否遇到过这样的情况？

RuntimeError: CUDA out of memory. Tried to allocate 256.00 MiB 
(GPU 0; 24.00 GiB total capacity; 20.50 GiB already allocated; 
180.00 MiB free; 21.80 GiB reserved in total by PyTorch)

明明 GPU 还有 180MB 空闲，为什么分配 256MB 就失败了？
答案是：显存碎片化（Memory Fragmentation）

PyTorch 默认的显存分配策略会导致内存碎片，即使总空闲显存足够，也可能因为没有连续的大块内存而分配失败。

🛠️ 解决方案：PYTORCH_CUDA_ALLOC_CONF

PYTORCH_CUDA_ALLOC_CONF 是 PyTorch 提供的环境变量，用于配置 CUDA 内存分配器的行为。

基本用法

Linux / macOS

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

Windows PowerShell

$env:PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"

Docker Compose

environment:
  - PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

设置后，PyTorch 在任何代码中导入时都会自动读取这个配置

import torch  # 自动读取 PYTORCH_CUDA_ALLOC_CONF
model.to("cuda")  # 内存分配器配置生效

📊 核心配置选项详解

1. expandable_segments:True（推荐）

作用：启用可扩展内存段，减少内存碎片

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

原理：

• 默认情况下，PyTorch 分配固定大小的内存块
• 启用后，内存段可以动态扩展
• 显著减少碎片化问题

适用场景：

• 加载大模型（如 LLM、VLM）
• 多个模型共享 GPU
• 显存紧张的环境

2. max_split_size_mb:128

作用：限制内存块的最大分割大小

export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128

原理：

• 防止大块内存被过度分割
• 减少小碎片的产生

适用场景：

• 频繁分配/释放不同大小的张量
• 训练过程中显存波动大

3. garbage_collection_threshold:0.6

作用：设置触发垃圾回收的阈值

export PYTORCH_CUDA_ALLOC_CONF=garbage_collection_threshold:0.6

原理：

• 当已分配显存达到总显存的 60% 时触发回收
• 默认值是 0.8（80%）

适用场景：

• 需要更积极的内存回收
• 与其他进程共享 GPU

4. 组合使用（推荐配置）

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:256,garbage_collection_threshold:0.7

🔧 实际应用示例

示例 1：本地运行大模型

运行前设置

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

然后运行你的脚本

python run_model.py

示例 2：Docker 部署

docker-compose.yml

services:
  model-server:
    image: pytorch/pytorch:2.4.0-cuda12.4-cudnn9-runtime
    environment:
      - PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:256
      - CUDA_VISIBLE_DEVICES=0
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]

示例 3：Jupyter Notebook

在 Notebook 的第一个 Cell 中：

import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

必须在 import torch 之前设置！

import torch

⚠️ 注意：必须在 import torch 之前设置环境变量！

示例 4：Python 脚本启动

#!/usr/bin/env python
import os

在导入 torch 之前设置

os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True,max_split_size_mb:256'

import torch
from transformers import AutoModel

# 你的代码...

📈 效果对比

以加载 4B 参数的视觉语言模型为例（24GB 显卡）：

配置	显存占用	OOM 发生率	推理速度
默认配置	~12GB	偶发	基准
expandable_segments:True	~10GB	极少	基准
组合配置	~9.5GB	几乎无	略慢 5%

🎯 最佳实践

显存紧张时的推荐配置

# 24GB 及以下显卡推荐
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:128

多模型共享 GPU

# 多个模型/进程共享时
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,garbage_collection_threshold:0.5

训练场景

# 训练时显存波动大
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:256,garbage_collection_threshold:0.6

🔍 其他有用的显存优化技巧

配合 PYTORCH_CUDA_ALLOC_CONF 使用：

1. 查看显存使用情况

import torch

print(f"已分配: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
print(f"已缓存: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")
print(f"最大分配: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

2. 手动清理缓存

import torch
import gc
# 清理 Python 垃圾
gc.collect()
# 清理 CUDA 缓存
torch.cuda.empty_cache()

3. 限制显存使用比例

import torch
# 限制只使用 80% 的显存
torch.cuda.set_per_process_memory_fraction(0.8, device=0)

4. 使用混合精度

# 使用 FP16 可以减少一半显存
model = AutoModel.from_pretrained("model_name", torch_dtype=torch.float16)

📚 总结

问题	解决方案
OOM 但显存有空闲	expandable_segments:True
显存碎片化严重	max_split_size_mb:128
多进程共享 GPU	garbage_collection_threshold:0.5
通用推荐配置	expandable_segments:True,max_split_size_mb:256

一行命令，显著减少 OOM 错误：

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

参考资料

• PyTorch CUDA Memory Management 官方文档
• PyTorch CUDA Allocator 源码