使用torch.distributed异步all-reduce进程仍阻塞，求排查

问题描述 #

我尝试使用torch.distributed的异步all-reduce，设置async_op=True后进程还是出现阻塞，想知道问题出在哪。

我复制了文档示例代码，添加了sleep和print来检测阻塞：

import torch
import torch.distributed as dist
import os
import time

rank = int(os.getenv('RANK', '0'))
dist.init_process_group(
        backend='nccl',
        world_size=2,
        rank=rank,
        )

output = torch.tensor([rank]).cuda(rank)
if rank == 1:
       time.sleep(5)

s = torch.cuda.Stream()
print(f"Process {rank}: begin aysnc all-reduce", flush=True)
handle = dist.all_reduce(output, async_op=True)
# Wait ensures the operation is enqueued, but not necessarily complete.
handle.wait()
# Using result on non-default stream.
print(f"Process {rank}: async check")
with torch.cuda.stream(s):
    s.wait_stream(torch.cuda.default_stream())
    output.add_(100)
if rank == 0:
    # if the explicit call to wait_stream was omitted, the output below will be
    # non-deterministically 1 or 101, depending on whether the allreduce overwrote
    # the value after the add completed.
    print(output)

实际输出 #

Process 0: begin aysnc all-reduce
Process 1: begin aysnc all-reduce
Process 1: async check
Process 0: async check
tensor([101], device='cuda:0')

预期结果 #

我预期Process 0: async check应该在Process 1: begin aysnc all-reduce之前打印。

问题原因与解决方法 #

问题核心在于handle.wait()的实际行为和你理解的不一样——这个方法会同步阻塞当前进程，直到异步all-reduce操作完全完成，而非仅仅确保操作入队。

虽然代码里的注释提到"Wait ensures the operation is enqueued, but not necessarily complete"，但这一表述不适用于NCCL后端的集合操作：对于NCCL来说，handle.wait()会强制等待所有参与进程的all-reduce操作全部结束，才会继续执行后续代码。

因为rank1进程先sleep了5秒，rank0的all-reduce必须等rank1参与才能完成，所以rank0会卡在handle.wait()处，直到rank1完成sleep并执行all-reduce后，两者的操作都结束，才会打印async check。

要实现你预期的非阻塞效果，直接去掉handle.wait()调用即可，让异步操作在后台执行，当前进程可以继续往下走：

print(f"Process {rank}: begin aysnc all-reduce", flush=True)
handle = dist.all_reduce(output, async_op=True)
# 移除handle.wait()调用
print(f"Process {rank}: async check")

这样rank0发起all-reduce后会立即打印async check，无需等待rank1的操作，符合预期。

另外注意：如果后续代码需要使用output的最终结果，必须在使用前调用handle.wait()或通过CUDA流同步来确保异步操作已完成，避免读取未更新的数据。

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