RTX3090(24G)加载小AI模型报CUDA Out Of Memory问题

1. 先清理混乱的运行环境 #

你当前环境同时安装了TensorFlow 2.9.1、CUDA 11.5、适配CUDA 11.6的PyTorch夜间版，版本不匹配叠加两个深度学习框架争抢显存，是问题的核心诱因。

• 直接新建干净的conda虚拟环境，不要在旧环境里凑合用：
  conda create -n dialobot python=3.9 -y
conda activate dialobot
• 安装和CUDA 11.5匹配的正式版PyTorch，不要装nightly版本，该版本的显存管理逻辑本身存在已知bug：
  pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu115
• 安装项目所需依赖，不要安装TensorFlow：gpt2bot本身基于PyTorch开发，安装TensorFlow只会在启动时自动占用显存，没有实际作用。
  pip install transformers==4.21.0 accelerate python-telegram-bot==13.7
• 把gpt2bot源码里所有导入tensorflow的代码全部注释掉，避免启动时自动初始化TF的CUDA上下文锁死显存。

2. 启动前彻底释放显存 #

• Windows系统打开任务管理器，结束所有python、jupyter、tensorboard相关进程；Ubuntu系统用nvidia-smi命令查询残留GPU进程，用kill -9 进程号全部关闭，保证启动前显卡空闲显存不低于23GB。
• Windows系统额外调整设置：进入系统设置→显示→图形设置，关闭「硬件加速GPU计划」后重启电脑，该功能会导致PyTorch显存识别错误，明明存在空闲显存也会提示0字节可用。
• 在运行脚本最开头添加配置，解决显存碎片问题，强制PyTorch使用动态显存分配，不要启动时直接占满全部显存：

import os
import torch
torch.cuda.empty_cache()
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"

3. 修改模型加载逻辑，直接降低一半显存占用 #

你触发报错的节点是加载第二个DialogRPT-updown模型时，默认pipeline使用FP32精度加载模型，显存占用本身偏高，修改两个参数即可解决：

• 找到gpt2bot/utils.py里的load_pipeline函数，所有模型加载时都添加torch_dtype=torch.float16参数，用半精度加载模型，显存占用直接减半，精度损失几乎无法感知，示例：
  pipeline(task, model=model_path, torch_dtype=torch.float16, device=0, **kwargs)
• 不要一次性把所有排序模型全部加载到GPU，修改build_ranker_dict逻辑，用到哪个模型再加载，暂时不用的模型挪到CPU内存，可额外节省4-5GB显存。

4. 效果验证 #

修改完成后先跑最小测试脚本，不要直接启动机器人：

import os
import torch
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
device = "cuda"
# 加载DialoGPT-medium
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
gen_model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium", torch_dtype=torch.float16).to(device)
print(f"生成模型加载完成，显存占用：{torch.cuda.memory_allocated()/1024**3:.1f}GB")
# 加载排序模型
ranker = pipeline("sentiment-analysis", model="microsoft/DialogRPT-updown", torch_dtype=torch.float16, device=0)
print(f"排序模型加载完成，总显存占用：{torch.cuda.memory_allocated()/1024**3:.1f}GB")

正常运行时两个模型全部加载完成显存占用在8-10GB区间，和你之前偶然成功的状态一致，远低于24GB显存上限。
如果依然触发报错，给模型加载参数添加load_in_8bit=True（需要提前安装bitsandbytes库，Windows系统安装对应预编译版本即可），显存占用可再降低一半，24G显卡运行全套模型完全没有压力。

避坑提醒：不要在同一个虚拟环境里同时安装PyTorch和TensorFlow跑GPU任务，两个框架初始化CUDA上下文时会互相锁显存，经常出现模型仅占2G、剩余20多G显存全部锁死无法使用的情况，和你报错日志里“2.07GiB已分配、0字节空闲”的状态完全吻合。

内容的提问来源于stack exchange，提问作者MagiCs ito