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 后，必须再强制设置当前设备，才能完全隔离。
希望这个坑能帮到更多遇到同样问题的朋友！
