PyTorch多进程调用init_process_group出现CUDA初始化错误如何解决

该CUDA初始化错误是PyTorch多进程与CUDA适配的经典问题：CUDA运行时不支持在fork模式启动的子进程中初始化上下文。Python多进程默认使用fork启动方式，若主进程（pytest调度进程）提前触发过CUDA上下文初始化（哪怕仅调用过torch.cuda.is_available()），fork出的子进程会继承损坏的CUDA上下文，在子进程中初始化NCCL分布式进程组时就会触发该报错。

1. 强制开启多进程spawn启动模式 #

你代码中set_start_method('spawn')处于注释状态，打开并增加重复设置容错即可解决核心问题，修改dist_launcher函数开头逻辑：

def dist_launcher(num_procs, *func_args, **func_kwargs):
    """Launch processes and gracefully handle failures. """
    # 新增spawn启动模式设置
    from torch.multiprocessing import set_start_method
    try:
        set_start_method('spawn')
    except RuntimeError:
        # 避免重复设置启动模式时报错
        pass
    # 剩下的原有逻辑不变

spawn模式会完全新建独立子进程，不会继承父进程的CUDA上下文，从根源避免CUDA上下文损坏问题。

2. 动态分配可用端口 #

你当前写死了MASTER_PORT为29503，多次运行测试时端口可能未及时释放，也会导致进程组初始化失败，新增动态端口获取逻辑：

def get_free_port() -> str:
    import socket
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(('127.0.0.1', 0))
        return str(s.getsockname()[1])

在dist_init中设置环境变量时替换固定端口：

os.environ['MASTER_PORT'] = get_free_port()

3. 补充进程组销毁逻辑 #

你代码中destroy_process_group处于注释状态，测试跑完后销毁进程组可以避免残留状态影响后续测试，在dist_init函数末尾run_func执行完成后增加：

run_func(*func_args, **func_kwargs)
# 新增销毁逻辑
dist.barrier()
dist.destroy_process_group()

4. 避免主进程触碰CUDA操作 #

pytest主进程、装饰器外层逻辑中不要调用任何CUDA相关接口（包括torch.cuda.is_available()、创建CUDA张量等），所有CUDA操作都放到子进程中执行。

如果不想手动维护多进程逻辑，可以直接使用PyTorch自带的torch.multiprocessing.spawn接口，该接口已经内置了多进程启动适配、错误传递等能力，稳定性比手动创建Process更高。

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