Repository: pytorch
Issue: Model parallel with DDP get `Socket Timeout` error when using NCCL, while GLOO works fine
Labels: oncall: distributed,triaged

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->
When training models in multi-machine multi-GPU setting on SLURM cluster, if `dist.init_process_group` with `NCCL` backend, and wrapping my multi-gpu model with `DistributedDataParallel` as [the official tutorial](https://pytorch.org/tutorials/intermediate/ddp_tutorial.html#combine-ddp-with-model-parallelism), a `Socket Timeout` runtime error will raise after the `timeout` (30s in my setting). However, if initialize the group with `GLOO` backend, everything works as expected.

Each of the processes holds a model replica, and I have double checked that GPUs assigned to each process do not overlap. i.e. process 0 get GPU{0, 1}, process 1 get GPU{2, 3}. Parameters in my model are separated to two sets, each set resides on one GPU. 

## To Reproduce

Steps to reproduce the behavior:

1. Launch several processes (8 and 16 in my experiments), call `dist.init_process_group` with `NCCL` backend;
1. Build a multi-gpu model, make sure GPUs assigned to each process do not overlap;
1. Wrap the model with `DistributedDataParallel` and do not assign `device_ids`, `output_device`. After `timeout`, encounter with `Runtime Error: Socket Timeout`.

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->
My codebase basically looks like this:
```python
# fire tasks on SLURM cluster...
os.environ["MASTER_PORT"] = str(port)
os.environ["MASTER_ADDR"] = str(master_ip)
os.environ["WORLD_SIZE"] = str(n_tasks)
os.environ["RANK"] = str(proc_id)
dist.init_process_group(backend=dist.Backend.NCCL, timeout=timedelta(seconds=30))

# ...
class MyModel(nn.Module)
    def __init__(self, ..., device0, device1):
        # ...
        self.part_1.to(device0)
        self.part_2.to(device1)

# task0 get GPU{0, 1}, task1 get GPU(2, 3)...
d0 = torch.device(f"cuda:{rank * 2}")
d1 = torch.device(f"cuda:{rank * 2 + 1}")
model = MyModel(..., d0, d1)
# not all parameters are used in each iteration
ddp_model = DistributedDataParallel(model, find_unused_parameters=True)

# ...
```

Invoking DDP did not raise any error, however after the `timeout` (30s in my setting), I encountered with following error:
### torch-1.1:
```bash
Traceback (most recent call last):
  File "../tools/train_val_classifier.py", line 332, in <module>
    main()
  File "../tools/train_val_classifier.py", line 103, in main
    model, model_without_ddp = get_ddp_model(model, devices=(fp_device, q_device))
  File ".../quant_prob/utils/distributed.py", line 120, in get_ddp_model
    ddp_model = DistributedDataParallel(model, device_ids=devices, find_unused_parameters=True)
  File "/envs/r0.3.0/lib/python3.6/site-packages/torch/nn/parallel/distributed.py", line 286, in __init__
    self.broadcast_bucket_size)
  File "/envs/r0.3.0/lib/python3.6/site-packages/torch/nn/parallel/distributed.py", line 410, in _dist_broadcast_coalesced
    dist._dist_broadcast_coalesced(self.process_group, tensors, buffer_size, False)
RuntimeError: Socket Timeout
```
### torch-1.2
```bash
Traceback (most recent call last):
  File "../tools/train_val_classifier.py", line 332, in <module>
    main()
  File "../tools/train_val_classifier.py", line 103, in main
    model, model_without_ddp = get_ddp_model(model, devices=(fp_device, q_device), debug=CONF.debug)
  File "/mnt/lustre/lirundong/Projects/quant-prob/quant_prob/utils/distributed.py", line 128, in get_ddp_model
    ddp_model = DistributedDataParallel(model, device_ids=devices, find_unused_parameters=True)
  File "/mnt/lustre/lirundong/Program/conda_env/torch-1.2-cuda-9.0/lib/python3.6/site-packages/torch/nn/parallel/distributed.py", line 298, in __init__
    self.broadcast_bucket_size)
  File "/mnt/lustre/lirundong/Program/conda_env/torch-1.2-cuda-9.0/lib/python3.6/site-packages/torch/nn/parallel/distributed.py", line 480, in _distributed_broadcast_coalesced
    dist._broadcast_coalesced(self.process_group, tensors, buffer_size)
RuntimeError: Socket Timeout
```

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->
If initialize with `GLOO` backend, the training and evaluation works properly: for each machine, GPU{0, 2, 4, 6} firstly do some computation, then GPU{1, 3, 5, 7}. Backpropagation works like DDP on single-gpu model.

## 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).

You can get the script and run it with:
```
wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
```

 - PyTorch Version (e.g., 1.0): both 1.1, 1.2
 - OS (e.g., Linux): Linux
 - How you installed PyTorch (`conda`, `pip`, source): 
   - torch-1.1: conda
   - torch-1.2: compiled from source
 - Build command you used (if compiling from source):
   ```python
   TORCH_CUDA_ARCH_LIST="6.0;6.1;7.0" CUDNN_LIB_DIR=$HOME/.local/cudnn-v7.6.2/lib64 CUDNN_INCLUDE_DIR=$HOME/.local/cudnn-v7.6.2/include python setup.py install
   ```
 - Python version: 3.6.9
 - CUDA/cuDNN version: 
   - torch-1.1: CUDA-9.0, cuDNN-7.5.1
   - torch-1.2: CUDA-9.0, cuDNN-7.6.2
 - GPU models and configuration: GTX-1080Ti, TitanXP
 - Any other relevant information: on SLURM cluster

## Additional context

<!-- Add any other context about the problem here. -->


cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera