企业官网建设流程全解析

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

问题现象

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

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

CUDA_VISIBLE_DEVICES=1python train.py --gpu1

程序日志显示：

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。

解决方案

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

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

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

最佳实践代码结构

importargparseimportosimporttorch# 先 import torchparser=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()# 可选，保险起见再清一次# 现在才导入模型等fromnetimportNet# ...

额外建议

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

总结

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

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