如何解决多GPU训练时出现的RuntimeError: CUDA error: out of memory

问题背景 #

单GPU训练正常，但切换到4卡训练时，即便batch size设为1仍报CUDA内存不足，错误触发点为torch.distributed.barrier()调用阶段。

已做配置 #

• Docker启动命令：

  docker exec -e NVIDIA_VISIBLE_DEVICES=0,1,2,3 -it jy /bin/bash
  

• 代码内GPU可见性设置：

  os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3'
  

当前GPU状态 #

|===============================+======================+======================|
|   0  GeForce RTX 208...  On   | 00000000:3D:00.0 Off |                  N/A |
| 24%   26C    P8    21W / 250W |      8MiB / 11019MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  GeForce RTX 208...  On   | 00000000:3E:00.0 Off |                  N/A |
| 25%   27C    P8     2W / 250W |      8MiB / 11019MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   2  GeForce RTX 208...  On   | 00000000:40:00.0 Off |                  N/A |
| 25%   25C    P8    20W / 250W |      8MiB / 11019MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   3  GeForce RTX 208...  On   | 00000000:41:00.0 Off |                  N/A |
| 26%   25C    P8    15W / 250W |      8MiB / 11019MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

错误日志 #

/root/anaconda3/envs/pytorch171/lib/python3.7/site-packages/torch/distributed/launch.py:186: FutureWarning: The module torch.distributed.launch is deprecated
and will be removed in future. Use torchrun.
Note that --use_env is set by default in torchrun.
If your script expects `--local_rank` argument to be set, please
change it to read from `os.environ['LOCAL_RANK']` instead. See 
https://pytorch.org/docs/stable/distributed.html#launch-utility for 
further instructions

  FutureWarning,
WARNING:torch.distributed.run:
*****************************************
Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
*****************************************
| distributed init (rank 2): env://
| distributed init (rank 1): env://
| distributed init (rank 3): env://
| distributed init (rank 0): env://
Traceback (most recent call last):
  File "main_track.py", line 398, in <module>
    main(args)
  File "main_track.py", line 159, in main
    utils.init_distributed_mode(args)
  File "/jy/TransTrack/util/misc.py", line 459, in init_distributed_mode
    torch.distributed.barrier()
  File "/root/anaconda3/envs/pytorch171/lib/python3.7/site-packages/torch/distributed/distributed_c10d.py", line 2709, in barrier
    work = default_pg.barrier(opts=opts)
RuntimeError: CUDA error: out of memory
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
WARNING:torch.distributed.elastic.multiprocessing.api:Sending process 355 closing signal SIGTERM
WARNING:torch.distributed.elastic.multiprocessing.api:Sending process 357 closing signal SIGTERM
WARNING:torch.distributed.elastic.multiprocessing.api:Sending process 358 closing signal SIGTERM
ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 1 (pid: 356) of binary: /root/anaconda3/envs/pytorch171/bin/python3
Traceback (most recent call last):
  File "/root/anaconda3/envs/pytorch171/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/root/anaconda3/envs/pytorch171/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/root/anaconda3/envs/pytorch171/lib/python3.7/site-packages/torch/distributed/launch.py", line 193, in <module>
    main()
  File "/root/anaconda3/envs/pytorch171/lib/python3.7/site-packages/torch/distributed/launch.py", line 189, in main
    launch(args)
  File "/root/anaconda3/envs/pytorch171/lib/python3.7/site-packages/torch/distributed/launch.py", line 174, in launch
    run(args)
  File "/root/anaconda3/envs/pytorch171/lib/python3.7/site-packages/torch/distributed/run.py", line 713, in run
    )(*cmd_args)
  File "/root/anaconda3/envs/pytorch171/lib/python3.7/site-packages/torch/distributed/launcher/api.py", line 131, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "/root/anaconda3/envs/pytorch171/lib/python3.7/site-packages/torch/distributed/launcher/api.py", line 261, in launch_agent
    failures=result.failures,
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
============================================================
main_track.py FAILED
------------------------------------------------------------
Failures:
  <NO_OTHER_FAILURES>
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2022-10-13_08:54:25
  host      : 2f923a848f88
  rank      : 1 (local_rank: 1)
  exitcode  : 1 (pid: 356)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
============================================================

排查与解决步骤 #

1. 修正分布式初始化逻辑
   错误发生在barrier()阶段，大概率是进程未正确绑定GPU导致多进程抢占同一设备内存。检查init_distributed_mode函数，确保调用barrier()前完成：

   • 通过local_rank绑定对应GPU：torch.cuda.set_device(args.local_rank)
   • 正确初始化进程组：torch.distributed.init_process_group(backend='nccl')

2. 清理重复GPU配置
   Docker启动命令和代码内同时设置CUDA_VISIBLE_DEVICES可能导致冲突，建议只保留其中一种配置，比如移除代码内的os.environ["CUDA_VISIBLE_DEVICES"]设置，依赖Docker环境变量生效。

3. 验证进程-GPU绑定关系
   在初始化流程中加入打印语句，确认每个进程绑定的GPU是否唯一：

   print(f"Rank {args.rank} using GPU {torch.cuda.current_device()}")
   

   若多个进程打印同一GPU，说明local_rank传递或解析存在问题，需检查启动参数或环境变量读取逻辑。

4. 调整模型加载时机
   避免在进程组初始化前加载模型，改为初始化完成后再加载，防止多进程同时在同一GPU预加载模型导致内存溢出。

5. 启用同步调试
   设置环境变量CUDA_LAUNCH_BLOCKING=1启动训练，获取更精准的错误栈，定位具体内存溢出环节。

6. 更换官方推荐的启动方式
   弃用已被官方标记为过时的torch.distributed.launch，改用torchrun启动，它默认启用--use_env，能更可靠传递LOCAL_RANK等环境变量，减少配置错误：

   torchrun --nproc_per_node=4 main_track.py [你的训练参数]
   

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