Repository: pytorch
Issue: torch.distributed with NCCL can hang if the first operation is a barrier
Labels: oncall: distributed,triaged,module: c10d

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

## To Reproduce

This is a slightly complex use case: running pytorch distributed on 2 nodes with 8 GPUs each, but we use only a single rank/GPU per node (for testing purposes).

Steps to reproduce the behavior:

1. Need a cluster with 2 separate nodes, and each node must have 2 or more GPUs.
2. Create a simple pytorch code that does `torch.distributed.barrier()` followed by `torch.distributed.all_reduce(x)`
```python3
import torch, os, torch.distributed as dist

local_rank = int(os.getenv('LOCAL_RANK', 0))
print(f"local_rank: {local_rank}")
torch.cuda.set_device(local_rank)

dist.init_process_group(backend='nccl', init_method='env://') ;

x = torch.tensor(range(32), dtype=torch.float64, device='cuda')

torch.distributed.barrier()

print('before all_reduce')
torch.distributed.all_reduce(x)
print(x)
```
3. Run the official 1.8.0 pytorch container on those 2 nodes (in our case we use the Slurm resource manager + the pytorch distributed launcher)
```console
$ srun -l --container-image=pytorch/pytorch:1.8.0-cuda11.1-cudnn8-devel bash -c 'set -x ; python -u -m torch.distributed.launch --nproc_per_node=1 --nnodes=2 --master_addr=XXX --node_rank=${SLURM_PROCID} /root/test_pytorch.py'
0: + python -u -m torch.distributed.launch --nproc_per_node=1 --nnodes=2 --master_addr=XXX --node_rank=0 /root/test_pytorch.py
0: local_rank: 0
1: + python -u -m torch.distributed.launch --nproc_per_node=1 --nnodes=2 --master_addr=XXX --node_rank=1 /root/test_pytorch.py
1: local_rank: 0
1: before all_reduce
0: before all_reduce
```
4. As you can see above, the code hangs and never prints the tensor.

I'm providing debug details in the "Additional context" section below,  the root cause seems to be that calling `torch.distributed.barrier()` immediately after `init_process_group` causes global rank 1 to unfortunately choose GPU 1 to perform this barrier:
https://github.com/pytorch/pytorch/blob/9f2aea7b88f69fc74ad90b1418663802f80c1863/torch/lib/c10d/ProcessGroupNCCL.cpp#L1383-L1390
Then, when calling `all_reduce`, rank 1 will try to create a new NCCL communicator, whereas rank 0 will not, so it's a deadlock.

Perhaps the **local** rank should be used instead of the **global** rank in the code above? Or perhaps the `device_ids` argument of the `barrier` function should default to the current device? Today `torch.cuda.set_device` has no impact on this behavior.

There are 2 known workarounds today:
- Set `CUDA_VISIBLE_DEVICES=0`, to hide other GPUs that shouldn't be used.
- Force the barrier to be on the GPU we want: `torch.distributed.barrier(device_ids=[local_rank])`


## Expected behavior
The pytorch code is expected to print the tensor and terminate. Using only 1/8 GPU per node might be unusual, but this hang was tricky to debug, and it was especially confusing since `all_reduce` then `barrier` works fine, but `barrier` then `all_reduce` doesn't:
```python3
import torch, os, torch.distributed as dist

local_rank = int(os.getenv('LOCAL_RANK', 0))
print(f"local_rank: {local_rank}")
torch.cuda.set_device(local_rank)

dist.init_process_group(backend='nccl', init_method='env://') ;

x = torch.tensor(range(32), dtype=torch.float64, device='cuda')

print('before all_reduce')
torch.distributed.all_reduce(x)
print(x)

torch.distributed.barrier()
print('done')
```
```console
$ srun -l --container-image=pytorch/pytorch:1.8.0-cuda11.1-cudnn8-devel bash -c 'set -x ; python -u -m torch.distributed.launch --nproc_per_node=1 --nnodes=2 --master_addr=XXX --node_rank=$SLURM_PROCID /root/test_pytorch2.py'
0: pyxis: importing docker image ...
1: pyxis: importing docker image ...
0: + python -u -m torch.distributed.launch --nproc_per_node=1 --nnodes=2 --master_addr=XXX --node_rank=0 /root/test_pytorch2.py
0: local_rank: 0
1: + python -u -m torch.distributed.launch --nproc_per_node=1 --nnodes=2 --master_addr=XXX --node_rank=1 /root/test_pytorch2.py
1: local_rank: 0
1: before all_reduce
0: before all_reduce
1: tensor([ 0.,  2.,  4.,  6.,  8., 10., 12., 14., 16., 18., 20., 22., 24., 26.,
1:         28., 30., 32., 34., 36., 38., 40., 42., 44., 46., 48., 50., 52., 54.,
1:         56., 58., 60., 62.], device='cuda:0', dtype=torch.float64)
0: tensor([ 0.,  2.,  4.,  6.,  8., 10., 12., 14., 16., 18., 20., 22., 24., 26.,
0:         28., 30., 32., 34., 36., 38., 40., 42., 44., 46., 48., 50., 52., 54.,
0:         56., 58., 60., 62.], device='cuda:0', dtype=torch.float64)
0: done
1: done
```

## Environment
 - PyTorch Version (e.g., 1.0): 1.8.0
 - OS (e.g., Linux): Linux
 - How you installed PyTorch (`conda`, `pip`, source): official Docker Hub container image
 - GPU models and configuration: V100

## Additional context

Rank 0 has a context active on GPU 0 only, as expected, and is blocked on `cuMemcpyDtoHAsync_v2`:
```
#0  0x00007f64213cc5cd in ?? () from /usr/lib/x86_64-linux-gnu/libcuda.so.1
[...]
#7  0x00007f64212b4852 in ?? () from /usr/lib/x86_64-linux-gnu/libcuda.so.1
#8  0x00007f6421362e4e in cuMemcpyDtoHAsync_v2 () from /usr/lib/x86_64-linux-gnu/libcuda.so.1
#9  0x00007f6527404d99 in ?? () from /opt/conda/lib/python3.8/site-packages/torch/lib/../../../../libcudart.so.11.0
#10 0x00007f65273e0d6d in ?? () from /opt/conda/lib/python3.8/site-packages/torch/lib/../../../../libcudart.so.11.0
#11 0x00007f652741ff0e in cudaMemcpyAsync () from /opt/conda/lib/python3.8/site-packages/torch/lib/../../../../libcudart.so.11.0
#12 0x00007f64cbe51045 in void at::native::nonzero_cuda_out_impl<bool>(at::Tensor const&, at::Tensor&) () from /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_cuda_cu.so
#13 0x00007f64cbddc5eb in at::native::nonzero_out_cuda(at::Tensor&, at::Tensor const&) () from /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_cuda_cu.so
```
It is likely waiting for the NCCL collective to finish, but it never happens since rank 1 is not participating in this collective.

Rank 1 has a context active on both GPU 0 *and* GPU 1, because the global rank was used when the barrier was called.
```
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A     33368      C   /opt/conda/bin/python            1125MiB |
|    1   N/A  N/A     33368      C   /opt/conda/bin/python            1135MiB |
+-----------------------------------------------------------------------------+
```
And it is blocked on `broadcastUniqueNCCLID`:
```
#0  0x00007fa3ff79e91e in __libc_recv (fd=32, buf=0x7ffcb08b21a6, len=1, flags=0) at ../sysdeps/unix/sysv/linux/recv.c:28
#1  0x00007fa3ce2ce169 in c10d::TCPStore::waitHelper_(std::vector<std::string, std::allocator<std::string> > const&, std::chrono::duration<long, std::ratio<1l, 1000l> > const&) () from /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_python.so
#2  0x00007fa3ce2ceb52 in c10d::TCPStore::getHelper_(std::string const&) () from /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_python.so
#3  0x00007fa3ce2ced2d in c10d::TCPStore::get(std::string const&) () from /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_python.so
#4  0x00007fa3ce2cbd42 in c10d::PrefixStore::get(std::string const&) () from /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_python.so
#5  0x00007fa3ce2d81d5 in c10d::ProcessGroupNCCL::broadcastUniqueNCCLID(ncclUniqueId*, c10d::OpType, std::string const&, int) () from /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_python.so
#6  0x00007fa3ce2dedfd in c10d::ProcessGroupNCCL::getNCCLComm(std::string const&, std::vector<c10::Device, std::allocator<c10::Device> > const&, c10d::OpType, int, bool) () from /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_python.so
#7  0x00007fa3ce2e8d67 in c10::detail::intrusive_target_default_null_type<c10d::ProcessGroup> c10d::ProcessGroupNCCL::collective<c10d::ProcessGroupNCCL::allreduce(std::vector<at::Tensor, std::allocator<at::Tensor> >&, c10d::AllreduceOptions const&)::{lambda(at::Tensor&, at::Tensor&, ncclComm*, c10::cuda::CUDAStream&)#1}, c10::intrusive_ptr<c10d::ProcessGroup::Work, c10::detail::intrusive_target_default_null_type<c10d::ProcessGroup> > c10d::ProcessGroupNCCL::collective<{lambda(at::Tensor&, at::Tensor&, ncclComm*, c10::cuda::CUDAStream&)#1}>(std::vector<at::Tensor, std::allocator<at::Tensor> >&, std::vector<at::Tensor, std::allocator<at::Tensor> >&, {lambda(at::Tensor&, at::Tensor&, ncclComm*, c10::cuda::CUDAStream&)#1}, c10d::OpType, char const*)::{lambda(std::vector<c10::cuda::CUDAStream, std::allocator<c10::cuda::CUDAStream> >&)#1}, c10::detail::intrusive_target_default_null_type<c10d::ProcessGroup> c10d::ProcessGroupNCCL::allreduce(std::vector<at::Tensor, std::allocator<at::Tensor> >&, c10d::AllreduceOptions const&)::{lambda(at::Tensor&, at::Tensor&, ncclComm*, c10::cuda::CUDAStream&)#1}<{lambda(at::Tensor&, at::Tensor&, ncclComm*, c10::cuda::CUDAStream&)#1}>(std::vector<at::Tensor, std::allocator<at::Tensor> >&, std::vector<at::Tensor, std::allocator<at::Tensor> >&, c10::intrusive_ptr<c10d::ProcessGroup::Work, c10::detail::intrusive_target_default_null_type<c10d::ProcessGroup> >, {lambda(at::Tensor&, at::Tensor&, ncclComm*, c10::cuda::CUDAStream&)#1}, char const)::{lambda(std::vector<c10::cuda::CUDAStream, std::allocator<c10::cuda::CUDAStream> >)#2}>(std::vector<at::Tensor, std::allocator<at::Tensor> >&, std::vector<at::Tensor, std::allocator<at::Tensor> >&, c10::intrusive_ptr<c10d::ProcessGroup::Work, c10::detail::intrusive_target_default_null_type<c10d::ProcessGroup> >, c10::intrusive_ptr<c10d::ProcessGroup::Work, c10::detail::intrusive_target_default_null_type<c10d::ProcessGroup> > c10d::ProcessGroupNCCL::collective<{lambda(at::Tensor&, at::Tensor&, ncclComm*, c10::cuda::CUDAStream&)#1}>(std::vector<at::Tensor, std::allocator<at::Tensor> >&, std::vector<at::Tensor, std::allocator<at::Tensor> >&, {lambda(at::Tensor&, at::Tensor&, ncclComm*, c10::cuda::CUDAStream&)#1}, c10d::OpType, char const*)::{lambda(std::vector<c10::cuda::CUDAStream, std::allocator<c10::cuda::CUDAStream> >&)#1}, c10::detail::intrusive_target_default_null_type<c10d::ProcessGroup> c10d::ProcessGroupNCCL::allreduce(std::vector<at::Tensor, std::allocator<at::Tensor> >&, c10d::AllreduceOptions const&)::{lambda(at::Tensor&, at::Tensor&, ncclComm*, c10::cuda::CUDAStream&)#1}<{lambda(at::Tensor&, at::Tensor&, ncclComm*, c10::cuda::CUDAStream&)#1}>(std::vector<at::Tensor, std::allocator<at::Tensor> >&, std::vector<at::Tensor, std::allocator<at::Tensor> >&, c10::intrusive_ptr<c10d::ProcessGroup::Work, c10::detail::intrusive_target_default_null_type<c10d::ProcessGroup> >, {lambda(at::Tensor&, at::Tensor&, ncclComm*, c10::cuda::CUDAStream&)#1}, char const)::{lambda(std::vector<c10::cuda::CUDAStream, std::allocator<c10::cuda::CUDAStream> >)#2}, c10d::ProcessGroupNCCL::allreduce(std::vector<at::Tensor, std::allocator<at::Tensor> >&, c10d::AllreduceOptions const&)::{lambda(at::Tensor&, at::Tensor&, ncclComm*, c10::cuda::CUDAStream&)#1}, char const) [clone .isra.821] [clone .constprop.868] ()
   from /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_python.so
#8  0x00007fa3ce2e9b52 in c10d::ProcessGroupNCCL::allreduce(std::vector<at::Tensor, std::allocator<at::Tensor> >&, c10d::AllreduceOptions const&) () from /opt/conda/lib/python3.8/site-packages/torch/lib/libtorch_python.so
```
I believe this rank is trying to create a new NCCL communicator because it has no cached communicator for GPU 0 (as the barrier was done on GPU 1).

Thanks a lot to my colleagues @kwen2501 @mcarilli @addyladdy for the pytorch/nccl debug help.

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @agolynski @SciPioneer @H-Huang @mrzzd @cbalioglu