Repository: pytorch
Issue: NCCL backend fails when calling broadcast from different threads
Labels: high priority,oncall: distributed,triaged,module: nccl

## 🐛 Bug

When I use the NCCL backend, `torch.distributed.broadcast()` results in a `NCCL error` when called from multiple threads (each thread uses a separate process group). `broadcast()` works fine from multiple threads when using the Gloo backend; NCCL backend works fine in the same environment when using a single thread.

## To Reproduce

Code to reproduce:

```python
import argparse
import os
import threading
import time
import torch
import torch.distributed as dist

NUM_TRIALS = 20

def receive_tensor_helper(tensor, src_rank, group, tag, num_iterations,
                          intra_server_broadcast):
    for i in range(num_iterations):
        if intra_server_broadcast:
            dist.broadcast(tensor=tensor, group=group, src=src_rank)
        else:
            dist.recv(tensor=tensor, src=src_rank, tag=tag)
    print("Done with tensor size %s" % tensor.size())

def send_tensor_helper(tensor, dst_rank, group, tag, num_iterations,
                       intra_server_broadcast):
    for i in range(num_iterations):
        if intra_server_broadcast:
            dist.broadcast(tensor=tensor, group=group, src=1-dst_rank)
        else:
            dist.send(tensor=tensor, dst=dst_rank, tag=tag)
    print("Done with tensor size %s" % tensor.size())

def start_helper_thread(func, args):
    helper_thread = threading.Thread(target=func,
                                     args=tuple(args))
    helper_thread.start()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Test lightweight communication library')
    parser.add_argument("--backend", type=str, default='gloo',
                        help="Backend")
    parser.add_argument('-s', "--send", action='store_true',
                        help="Send tensor (if not specified, will receive tensor)")
    parser.add_argument("--master_addr", required=True, type=str,
                        help="IP address of master")
    parser.add_argument("--use_helper_threads", action='store_true',
                        help="Use multiple threads")
    parser.add_argument("--rank", required=True, type=int,
                        help="Rank of current worker")
    parser.add_argument('-p', "--master_port", default=12345,
                        help="Port used to communicate tensors")
    parser.add_argument("--intra_server_broadcast", action='store_true',
                        help="Broadcast within a server")

    args = parser.parse_args()

    num_ranks_in_server = 1
    if args.intra_server_broadcast:
        num_ranks_in_server = 2
    local_rank = args.rank % num_ranks_in_server
    torch.cuda.set_device(local_rank)

    os.environ['MASTER_ADDR'] = args.master_addr
    os.environ['MASTER_PORT'] = str(args.master_port)
    world_size = 2
    dist.init_process_group(args.backend, rank=args.rank, world_size=world_size)

    tensor_sizes = [10, 100, 1000, 10000, 100000, 1000000, 10000000]

    groups = []
    for tag in range(len(tensor_sizes)):
        if args.intra_server_broadcast:
            group = dist.new_group([0, 1])
            groups.append(group)

    for tag, tensor_size in enumerate(tensor_sizes):
        if args.intra_server_broadcast:
            group = groups[tag]
        else:
            group = None

        if args.send:
            if args.intra_server_broadcast:
                tensor = torch.tensor(range(tensor_size), dtype=torch.float32).cuda(
                    local_rank)
            else:
                tensor = torch.tensor(range(tensor_size), dtype=torch.float32).cpu()
            if args.use_helper_threads:
                start_helper_thread(send_tensor_helper,
                                    [tensor, 1-args.rank,
                                     group, tag, NUM_TRIALS,
                                     args.intra_server_broadcast])
            else:
                send_tensor_helper(tensor, 1-args.rank, group, tag,
                                   NUM_TRIALS, args.intra_server_broadcast)
        else:
            if args.intra_server_broadcast:
                tensor = torch.zeros((tensor_size,), dtype=torch.float32).cuda(
                    local_rank)
            else:
                tensor = torch.zeros((tensor_size,), dtype=torch.float32).cpu()
            if args.use_helper_threads:
                start_helper_thread(receive_tensor_helper,
                                    [tensor, 1-args.rank,
                                     group, tag, NUM_TRIALS,
                                     args.intra_server_broadcast])
            else:
                receive_tensor_helper(tensor, 1-args.rank, group, tag,
                                      NUM_TRIALS, args.intra_server_broadcast)
```

Commands to reproduce the error (on the same multi-GPU server):
```bash
python harness.py --master_addr localhost --rank 0 --intra_server_broadcast --backend nccl --use_helper_threads --send
python harness.py --master_addr localhost --rank 1 --intra_server_broadcast --backend nccl --use_helper_threads
```

Error trace:
```bash
Traceback (most recent call last):
  File "/opt/conda/lib/python3.6/threading.py", line 916, in _bootstrap_inner
    self.run()
  File "/opt/conda/lib/python3.6/threading.py", line 864, in run
    self._target(*self._args, **self._kwargs)
  File "harness.py", line 15, in receive_tensor_helper
    dist.broadcast(tensor=tensor, group=group, src=src_rank)
  File "/opt/conda/lib/python3.6/site-packages/torch/distributed/distributed_c10d.py", line 741, in broadcast
    work = group.broadcast([tensor], opts)
RuntimeError: NCCL error in: ../torch/lib/c10d/ProcessGroupNCCL.cpp:260, unhandled system error
```

These commands work as expected:
```bash
python harness.py --master_addr localhost --rank 0 --intra_server_broadcast --backend gloo --use_helper_threads --send
python harness.py --master_addr localhost --rank 1 --intra_server_broadcast --backend gloo --use_helper_threads
```
```bash
python harness.py --master_addr localhost --rank 0 --intra_server_broadcast --backend nccl --send
python harness.py --master_addr localhost --rank 1 --intra_server_broadcast --backend nccl
```

## Expected behavior

No runtime error / tensors successfully broadcasted among participating ranks.

## Environment

Please copy and paste the output from our [environment collection script](https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py)
(or fill out the checklist below manually).

```
Collecting environment information...
PyTorch version: 1.0.0a0+056cfaf
Is debug build: No
CUDA used to build PyTorch: 10.0.130

OS: Ubuntu 16.04.5 LTS
GCC version: (Ubuntu 5.4.0-6ubuntu1~16.04.11) 5.4.0 20160609
CMake version: version 3.5.1

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: 10.0.130
GPU models and configuration:
GPU 0: Tesla V100-PCIE-16GB
GPU 1: Tesla V100-PCIE-16GB
GPU 2: Tesla V100-PCIE-16GB
GPU 3: Tesla V100-PCIE-16GB

Nvidia driver version: 384.130
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.4.2

Versions of relevant libraries:
[pip] msgpack-numpy==0.4.3.2
[pip] numpy==1.15.4
[pip] torch==1.0.0a0+056cfaf
[pip] torchtext==0.4.0
[pip] torchvision==0.2.1
[conda] blas                      1.0                         mkl
[conda] magma-cuda100             2.4.0                         1    pytorch
[conda] mkl                       2018.0.3                      1
[conda] mkl-include               2019.1                      144
[conda] mkl_fft                   1.0.6            py36h7dd41cf_0
[conda] mkl_random                1.0.1            py36h4414c95_1
[conda] torch                     1.0.0a0+056cfaf           <pip>
[conda] torchtext                 0.4.0                     <pip>
[conda] torchvision               0.2.1                     <pip>
```