Repository: pytorch
Issue: DistributedDataParallel hangs when not using GPU 0
Labels: oncall: distributed,module: ddp

### 🐛 Describe the bug

Wrapping a model in DDP results in the training hanging when not using GPU 0 on a multi-GPU device. 

The following MVE:

```python
import os
import torch
import torch.multiprocessing as mp
import torch.nn as nn
from torch.optim import Adam
from torch.utils.data import DataLoader, TensorDataset, DistributedSampler
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP



def main():
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '8880'
    gpus = 2
    offset = 2
    mp.spawn(train, nprocs=gpus, args=(gpus, offset))


def train(gpu, world_size, offset):
    _ = torch.manual_seed(0)
    rank = gpu
    dist.init_process_group(
        backend='nccl',
        init_method='env://',
        world_size=world_size,
        rank=rank)
    torch.cuda.set_device(gpu + offset)
    device = torch.device('cuda:' + str(gpu + offset))
    d = 256
    model = nn.Linear(d, 1)
    n_train = 1024
    x_train = torch.randn((n_train, d))
    y_train = torch.randn((n_train, 1))
    ds_train = TensorDataset(x_train, y_train)
    sampler = DistributedSampler(ds_train, num_replicas=world_size, rank=rank)
    dl_train = DataLoader(dataset=ds_train,
                          sampler=sampler,
                          batch_size=32,
                          num_workers=1)


    optimizer = Adam(model.parameters(), lr=1e-3)
    model = model.to(device)
    model = DDP(model, device_ids=[gpu + offset], output_device=gpu + offset)
    loss_func = nn.MSELoss()
    for x, y in dl_train:
        print(rank, device)
        x = x.to(device)
        print(rank, x.shape)
        y = y.to(device)
        p = model(x)
        loss = loss_func(p, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        print(rank, loss)


if __name__ == '__main__':
    main()
```

Produces:

```
1 cuda:3                                                                                                                                                  
1 torch.Size([32, 256])                                                                                                                                   
0 cuda:2                                                                                                                                                  
0 torch.Size([32, 256])                                                                                                                                   
1 tensor(1.0912, device='cuda:3', grad_fn=<MseLossBackward0>)                                                                                             
0 tensor(0.9560, device='cuda:2', grad_fn=<MseLossBackward0>)                                                                                             
1 cuda:3                                                                                                                                                  
1 torch.Size([32, 256])                                                                                                                                   
0 cuda:2                                                                                                                                                  
0 torch.Size([32, 256])                                                                                                                                   
1 tensor(1.3682, device='cuda:3', grad_fn=<MseLossBackward0>)                                                                                             
1 cuda:3                                                                                                                                                  
1 torch.Size([32, 256])  
```

Indicating that after a few batches, the rank 0 process never gets to its backwards pass and everything is stuck. If I set offset to 0 so that I use GPUs 0 and 1 instead of 2 and 3, everything proceeds as expected. If I use GPUs 0 and 2, everything is also fine. 


### Versions

Collecting environment information...
PyTorch version: 1.12.0
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.9.12 (main, Jun  1 2022, 11:38:51)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-122-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 10.1.243
GPU models and configuration: 
GPU 0: NVIDIA RTX A6000
GPU 1: NVIDIA RTX A6000
GPU 2: NVIDIA RTX A6000
GPU 3: NVIDIA RTX A6000

Nvidia driver version: 510.73.08
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.23.1
[pip3] torch==1.12.0
[pip3] torchaudio==0.12.0
[pip3] torchvision==0.13.0
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.6.0              hecad31d_10    conda-forge
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0            py39h7f8727e_0  
[conda] mkl_fft                   1.3.1            py39hd3c417c_0  
[conda] mkl_random                1.2.2            py39h51133e4_0  
[conda] numpy                     1.22.3           py39he7a7128_0  
[conda] numpy-base                1.22.3           py39hf524024_0  
[conda] pytorch                   1.12.0          py3.9_cuda11.6_cudnn8.3.2_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                0.12.0               py39_cu116    pytorch
[conda] torchvision               0.13.0               py39_cu116    pytorch

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang @kwen2501 @ezyang