如何在线程中正确配置DDP模型并分配GPU资源？

问题描述 #

希望在线程内训练DDP模型，当前线程创建代码如下：

t1 = threading.Thread(target=training_func, name='target')
t1.start()

其中training_func是创建DDP模型的函数，由PyTorch Lightning代码处理，但函数在DDP进程初始化阶段卡住，日志如下：

Initializing distributed: GLOBAL_RANK: 0, MEMBER: 1/2
host:3886:3998 [1] NCCL INFO P2P is disabled between connected GPUs 0 and 1. You can repress this message with NCCL_IGNORE_DISABLED_P2P=1.
host:3886:3998 [1] NCCL INFO Could not enable P2P between dev 0(=3000) and dev 1(=88000)
host:3812:3996 [0] NCCL INFO P2P is disabled between connected GPUs 1 and 0. You can repress this message with NCCL_IGNORE_DISABLED_P2P=1.

尝试过PyTorch Lightning Trainer、原生dist.init_process_group、修改后端为gloo，甚至Ray、Huggingface等库均出现相同问题，推测底层依赖PyTorch分布式训练代码，期望线程能正确调用可用GPU完成模型初始化。

解决方案 #

1. 放弃线程，改用多进程启动分布式训练 #

PyTorch DDP（及依赖它的PyTorch Lightning、Huggingface等框架）是多进程架构，线程共享同一进程的GPU上下文，会导致分布式进程间GPU资源冲突、NCCL初始化失败，这是卡住的核心原因。

正确做法是使用torch.multiprocessing启动多进程训练，示例代码：

import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

def training_func(rank, world_size):
    # 绑定当前进程到指定GPU
    os.environ['CUDA_VISIBLE_DEVICES'] = str(rank)
    torch.cuda.set_device(rank)
    # 初始化分布式进程组
    dist.init_process_group(
        backend='nccl',
        init_method='tcp://127.0.0.1:23456',
        world_size=world_size,
        rank=rank
    )
    # 后续模型定义、训练逻辑
    ...

if __name__ == '__main__':
    world_size = 2  # GPU数量
    mp.spawn(training_func, args=(world_size,), nprocs=world_size, join=True)

2. 显式为进程分配GPU资源 #

无论用原生PyTorch还是PyTorch Lightning，都需要显式指定每个进程使用的GPU：

• PyTorch Lightning 示例：

  from pytorch_lightning import Trainer
  
  trainer = Trainer(
      devices=[0, 1],  # 指定要使用的GPU编号
      accelerator='gpu',
      strategy='ddp',
      num_nodes=1
  )
  trainer.fit(model)
  

• 环境变量限制GPU可见性：
  启动训练前设置CUDA_VISIBLE_DEVICES，让每个进程只能看到分配的GPU：

  CUDA_VISIBLE_DEVICES=0 python train.py  # 进程1用GPU0
  CUDA_VISIBLE_DEVICES=1 python train.py  # 进程2用GPU1
  

3. 消除NCCL P2P警告 #

日志中的P2P警告虽不是卡住的直接原因，但可通过环境变量消除：

import os
os.environ['NCCL_IGNORE_DISABLED_P2P'] = '1'

或在终端启动时设置：

export NCCL_IGNORE_DISABLED_P2P=1

4. 不推荐：线程内强制绑定GPU（仅用于调试） #

如果必须在线程内尝试，需确保线程绑定单个GPU且主线程释放CUDA上下文，但仍可能出现问题：

def training_func():
    # 线程内仅可见指定GPU
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    torch.cuda.set_device(0)
    # 清空主线程可能残留的CUDA缓存
    torch.cuda.empty_cache()
    # 后续DDP初始化逻辑
    ...

此方法不适合生产环境，仅作临时调试用。

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