PyTorch DDP中all_gather_object阻塞或NCCL报错问题求助

背景 #

在独立GPU上通过PyTorch DDP训练模型，使用all_gather_object函数收集本地对象时，遇到两类异常。

问题1：程序卡在dist.all_gather_object处 #

代码版本1 #

import os
import logging
import torch
import torch.nn as nn
import torch.multiprocessing as mp
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import Dataset, DataLoader, TensorDataset

from util import *

logger = logging.getLogger("demo")
logger.addHandler(logging.FileHandler("demo.out"))
logger.setLevel(logging.INFO)

class ToyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(3, 3)
    def forward(self, x):
        return self.linear(x)

if __name__ == "__main__":
    '''torch setup'''
    reserve_gpu([2,3])
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    device_id = rank % torch.cuda.device_count()

    '''train'''
    x = []
    for inst in range(12):
        x.append(torch.randn(3))
    x=torch.stack(x)
    dataset = TensorDataset(x)
    sampler = DistributedSampler(dataset=dataset)
    loader = DataLoader(
        dataset=dataset,
        num_workers=8,
        pin_memory=True,
        sampler=sampler
    )
    model = DDP(
        ToyModel().to(device_id),
        device_ids=[device_id],
    )
    metrics = {"rank":rank}
    # for epoch in range(1):
    #     for idx, x in enumerate(loader):
    #         logger.info(f'<{os.getpid()}>: {idx}, {x}')
    output = [None for _ in range(dist.get_world_size())]
    dist.all_gather_object(output, metrics)
    print(output)

预期输出 #

[{"rank":0},{"rank":1}]
[{"rank":0},{"rank":1}]

实际异常 #

程序卡在dist.all_gather_object(output, metrics)调用处，无响应。

问题2：触发NCCL invalid usage错误 #

代码版本2 #

import os
import logging
import torch
import torch.nn as nn
import torch.multiprocessing as mp
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import Dataset, DataLoader, TensorDataset

from util import *

logger = logging.getLogger("demo")
logger.addHandler(logging.FileHandler("demo.out"))
logger.setLevel(logging.INFO)

class ToyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(3, 3)
    def forward(self, x):
        return self.linear(x)

if __name__ == "__main__":
    '''torch setup'''
    # reserve_gpu([2,3])
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    device_id = rank % torch.cuda.device_count()
    gather_objects = ["foo", 12, {1: 2}] # any picklable object
    output = [None for _ in gather_objects]
    dist.all_gather_object(output, gather_objects[dist.get_rank()])
    print(output)

报错回溯 #

Traceback (most recent call last):
Traceback (most recent call last):
  File "demo.py", line 32, in <module>
  File "demo.py", line 32, in <module>
        dist.all_gather_object(output, gather_objects[dist.get_rank()])dist.all_gather_object(output, gather_objects[dist.get_rank()])

  File "/data_HDD/zhuxingyu/anaconda3/envs/p11/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py", line 1649, in all_gather_object
  File "/data_HDD/zhuxingyu/anaconda3/envs/p11/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py", line 1649, in all_gather_object
    all_gather(object_size_list, local_size, group=group)
  File "/data_HDD/zhuxingyu/anaconda3/envs/p11/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py", line 2060, in all_gather
    all_gather(object_size_list, local_size, group=group)
  File "/data_HDD/zhuxingyu/anaconda3/envs/p11/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py", line 2060, in all_gather
    work = default_pg.allgather([tensor_list], [tensor])
RuntimeError: NCCL error in: /opt/conda/conda-bld/pytorch_1646755903507/work/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:1169, invalid usage, NCCL version 21.0.3
ncclInvalidUsage: This usually reflects invalid usage of NCCL library (such as too many async ops, too many collectives at once, mixing streams in a group, etc).
    work = default_pg.allgather([tensor_list], [tensor])
RuntimeError: NCCL error in: /opt/conda/conda-bld/pytorch_1646755903507/work/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:1169, invalid usage, NCCL version 21.0.3
ncclInvalidUsage: This usually reflects invalid usage of NCCL library (such as too many async ops, too many collectives at once, mixing streams in a group, etc).
ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 0 (pid: 3101) of binary: /data_HDD/zhuxingyu/anaconda3/envs/p11/bin/python
Traceback (most recent call last):
  File "/data_HDD/zhuxingyu/anaconda3/envs/p11/bin/torchrun", line 33, in <module>
    sys.exit(load_entry_point('torch==1.11.0', 'console_scripts', 'torchrun')())
  File "/data_HDD/zhuxingyu/anaconda3/envs/p11/lib/python3.8/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 345, in wrapper
    return f(*args, **kwargs)
  File "/data_HDD/zhuxingyu/anaconda3/envs/p11/lib/python3.8/site-packages/torch/distributed/run.py", line 724, in main
    run(args)
  File "/data_HDD/zhuxingyu/anaconda3/envs/p11/lib/python3.8/site-packages/torch/distributed/run.py", line 715, in run
    elastic_launch(
  File "/data_HDD/zhuxingyu/anaconda3/envs/p11/lib/python3.8/site-packages/torch/distributed/launcher/api.py", line 131, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "/data_HDD/zhuxingyu/anaconda3/envs/p11/lib/python3.8/site-packages/torch/distributed/launcher/api.py", line 245, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
============================================================
demo.py FAILED
------------------------------------------------------------
Failures:
[1]:
  time      : 2022-11-02_12:45:22
  host      : ubuntu
  rank      : 1 (local_rank: 1)
  exitcode  : 1 (pid: 3102)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2022-11-02_12:45:22
  host      : ubuntu
  rank      : 0 (local_rank: 0)
  exitcode  : 1 (pid: 3101)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
============================================================

问题分析与解决 #

问题1解决思路 #

1. 核心原因：NCCL后端下，all_gather_object依赖的CUDA张量通信需要同步，未显式同步会导致进程阻塞；另外reserve_gpu若未正确隔离GPU资源，也会引发通信异常。
2. 修复步骤：
   • 在调用all_gather_object前添加torch.cuda.synchronize(device_id)，确保CUDA操作完成后再执行集合通信。
   • 检查reserve_gpu函数是否通过os.environ["CUDA_VISIBLE_DEVICES"]正确限制每个进程的可见GPU，避免资源冲突。
   • 确保所有进程都能执行到all_gather_object，无提前退出情况。

问题2解决思路 #

1. 核心原因：all_gather_object要求output列表长度必须等于进程数（dist.get_world_size()），代码中output长度与gather_objects绑定为3，但实际进程数为2，导致通信逻辑不匹配触发NCCL错误。
2. 修复步骤：
   • 将output长度改为进程数：output = [None for _ in range(dist.get_world_size())]。
   • 确保待收集对象可被pickle序列化（代码中对象满足要求，避免传入CUDA张量等不可序列化对象）。

修复后示例代码（问题1） #

import os
import logging
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import TensorDataset, DataLoader

from util import *

logger = logging.getLogger("demo")
logger.addHandler(logging.FileHandler("demo.out"))
logger.setLevel(logging.INFO)

class ToyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(3, 3)
    def forward(self, x):
        return self.linear(x)

if __name__ == "__main__":
    '''torch setup'''
    reserve_gpu([2,3])
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    device_id = rank % torch.cuda.device_count()
    torch.cuda.set_device(device_id)

    '''train'''
    x = torch.randn(12, 3)
    dataset = TensorDataset(x)
    sampler = DistributedSampler(dataset=dataset)
    loader = DataLoader(
        dataset=dataset,
        num_workers=8,
        pin_memory=True,
        sampler=sampler
    )
    model = DDP(
        ToyModel().to(device_id),
        device_ids=[device_id],
    )
    metrics = {"rank":rank}
    
    # 同步CUDA操作
    torch.cuda.synchronize(device_id)
    
    output = [None for _ in range(dist.get_world_size())]
    dist.all_gather_object(output, metrics)
    print(output)
    
    # 清理进程组
    dist.destroy_process_group()

修复后示例代码（问题2） #

import os
import torch
import torch.distributed as dist

from util import *

if __name__ == "__main__":
    '''torch setup'''
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    device_id = rank % torch.cuda.device_count()
    torch.cuda.set_device(device_id)
    
    gather_objects = ["foo", 12]  # 长度与进程数一致
    output = [None for _ in range(dist.get_world_size())]
    dist.all_gather_object(output, gather_objects[rank])
    
    print(output)
    dist.destroy_process_group()

内容的提问来源于stack exchange，提问作者K. Zhu