2026/1/13 23:07:43
网站建设
项目流程
关于建设门户网站,python在线,wordpress ico文件,即时设计是什么软件PyTorch 多卡训练常见坑#xff1a;设置 CUDA_VISIBLE_DEVICES 后仍 OOM 在 GPU 0 的解决之道
问题现象
服务器有两张 GPU#xff08;GPU 0 和 GPU 1#xff09;#xff0c;GPU 0 正在跑一个大模型任务。
你想在 GPU 1 上单独跑另一个 PyTorch 程序#xff0c;于是这样启动…PyTorch 多卡训练常见坑设置 CUDA_VISIBLE_DEVICES 后仍 OOM 在 GPU 0 的解决之道问题现象服务器有两张 GPUGPU 0 和 GPU 1GPU 0 正在跑一个大模型任务。你想在 GPU 1 上单独跑另一个 PyTorch 程序于是这样启动CUDA_VISIBLE_DEVICES1python 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正确限制了可见 GPUPyTorch 在首次初始化 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 torchparserargparse.ArgumentParser()parser.add_argument(--gpu,typestr,default1)argsparser.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 后必须再强制设置当前设备才能完全隔离。希望这个坑能帮到更多遇到同样问题的朋友