PyTorch 多卡训练常见坑：设置 CUDA_VISIBLE_DEVICES 后仍 OOM 在 GPU 0 的解决之道

cvyoutian2025-12-15 8:48

PyTorch 多卡训练常见坑：设置 CUDA_VISIBLE_DEVICES 后仍 OOM 在 GPU 0 的解决之道

问题现象

服务器有两张 GPU（GPU 0 和 GPU 1），GPU 0 正在跑一个大模型任务。

你想在 GPU 1 上单独跑另一个 PyTorch 程序，于是这样启动：

bash 复制代码

CUDA_VISIBLE_DEVICES=1 python train.py --gpu 1

程序日志显示：

复制代码

using GPU : 1
可见 GPU 数量: 1
当前 GPU ID: 0
GPU 名称: NVIDIA GeForce RTX 3090  # 正确是 GPU 1

一切看似正常，但运行几步后报错：

复制代码

RuntimeError: CUDA out of memory ... (GPU 0; 23.69 GiB total capacity; ...)

明明只看到一张卡，为什么还在物理 GPU 0 上爆显存？

根本原因

这是 PyTorch（尤其是 1.x 版本）的历史遗留 bug：

即使通过 CUDA_VISIBLE_DEVICES 正确限制了可见 GPU，PyTorch 在首次初始化 CUDA 上下文时，仍会在原始的物理 GPU 0 上分配少量内存（用于内部通信、缓存等）。

后续模型运行虽然在指定的 GPU 上，但只要触发某些操作（如大 tensor 分配、某些模块 forward），就会唤醒 GPU 0 的旧上下文，导致显存被占用，最终 OOM。

解决方案

在所有模型创建和数据移动之前，强制设置当前设备：

python 复制代码

import torch

# 在创建 net、前，紧跟 import torch 之后
torch.cuda.set_device(0)   # 这里的 0 是逻辑 ID，即你通过 CUDA_VISIBLE_DEVICES 指定的那张卡

加上这行后，PyTorch 会彻底绑定到新上下文，旧的 GPU 0 上下文不再被触碰，问题瞬间解决。

最佳实践代码结构

python 复制代码

import argparse
import os
import torch   # 先 import torch

parser = argparse.ArgumentParser()
parser.add_argument("--gpu", type=str, default='1')
args = parser.parse_known_args()[0]
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

# 关键：强制设置设备
torch.cuda.set_device(0)
torch.cuda.empty_cache()  # 可选，保险起见再清一次

# 现在才导入模型等
from net import Net
# ...

额外建议

单卡训练时不要使用 torch.nn.DataParallel，它会引入额外开销并可能触发旧 bug。
先用小 batchSize（如 2~4）验证跑通，再逐步增大。
监控命令：watch -n 1 nvidia-smi

总结

一句 torch.cuda.set_device(0) 就能彻底解决这个让人抓狂的" GPU 0 占用"问题。记住：限制可见 GPU 后，必须再强制设置当前设备，才能完全隔离。

希望这个坑能帮到更多遇到同样问题的朋友！