Consider the following MWE, where I attempt to simply sum random tensors that are generated in different GPUs. If I generate tensors of size e.g., 5.000, it still works, but for size 10.000 it timeouts.
I’m using CUDA 11.2 w/ 4 RTX A6000. Tried both torch-1.9.1+cu111 and the nightly one compiled directly from the repo. Note that if I use gloo as the backend, then it works.
Is this a bug, or maybe there is something wrong with my environment? Any idea of what could I try? Thanks.
import os

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


def fn(rank, world_size):
    # Set up distributed job
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '11235'
    dist.init_process_group('nccl', rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

    # Generate random tensor and do all reduce
    x = torch.randn(10_000, device='cuda')
    dist.all_reduce(tensor=x)
    print(x[0])

    dist.destroy_process_group()

if __name__ == '__main__':
    if not torch.cuda.is_available():
        print('No GPUs available, cannot run test')
    else:
        n_gpus = torch.cuda.device_count()
        mp.spawn(
            fn,
            args=(n_gpus,),
            nprocs=n_gpus,
            join=True,
        )