Repository: pytorch
Issue: DDP hangs on forward pass of transformer
Labels: oncall: distributed

### 🐛 Describe the bug

I'm facing an issue where DDP is hanging sometimes. A relevant code snippet is below. The setup is fairly simple: I have an encoder, and I want to do a forward pass on one GPU but not the other. Based on which GPU I let it do the forward pass on, I see varying results as to what processes complete.

I'm running the file as follows:
```
export CUDA_VISIBLE_DEVICES=5,6
python -m torch.distributed.launch --nproc_per_node=2 ddp.py
```

Here is the file ddp.py:
```python
import torch
import argparse
from transformers import RobertaModel
from torch.nn.parallel import DistributedDataParallel as DDP

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--local_rank", type=int, default=-1,
                        help="For distributed training: local_rank")   

    args = parser.parse_args()
    torch.cuda.set_device(args.local_rank)
    device = torch.device("cuda", args.local_rank)
    torch.distributed.init_process_group(backend='nccl')
    args.device = device
        
    encoder = RobertaModel.from_pretrained("microsoft/codebert-base")
    encoder.to(device)
    
    encoder = DDP(encoder, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=False)
    
    source_ids = torch.zeros([8, 256]).long().to(device)
    source_mask = torch.ones([8, 256]).long().to(device)

    # devices = [-1, 1] # 0 started, 1 started, 1 done, unexpected
    devices = [-1, 0] # 0 started, 1 started, unexpected
    # devices = [-1, 0, 1] # 0 started, 0 started, 0 done, 0 done, expected
    if args.local_rank not in devices:
        print("process 1 started")
        output = encoder(source_ids, attention_mask=source_mask)
        torch.distributed.barrier()
        print("process 1 done")
    
    if args.local_rank in devices:
        print("process 0 started")
        torch.distributed.barrier()
        print("process 0 done")
                            
if __name__ == "__main__":
    main()
```

### Versions


Collecting environment information...
PyTorch version: 1.12.1+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 16.04.7 LTS (x86_64)
GCC version: (Ubuntu 5.5.0-12ubuntu1~16.04) 5.5.0 20171010
Clang version: Could not collect
CMake version: version 3.24.2
Libc version: glibc-2.23

Python version: 3.8.15 (default, Nov 24 2022, 15:19:38)  [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-4.4.0-210-generic-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: 
GPU models and configuration: 
GPU 0: GeForce GTX 1080 Ti
GPU 1: GeForce GTX 1080 Ti
GPU 2: GeForce GTX 1080 Ti
GPU 3: GeForce GTX 1080 Ti
GPU 4: GeForce GTX 1080 Ti
GPU 5: GeForce GTX 1080 Ti
GPU 6: GeForce GTX 1080 Ti
GPU 7: GeForce GTX 1080 Ti

Nvidia driver version: 460.27.04
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.6.0.21
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.22.4
[pip3] torch==1.12.1
[pip3] torchaudio==0.12.1
[pip3] torchvision==0.13.1
[pip3] torchviz==0.0.2
[conda] mkl                       2022.2.1                 pypi_0    pypi
[conda] mkl-fft                   1.3.1                    pypi_0    pypi
[conda] mkl-random                1.2.2                    pypi_0    pypi
[conda] mkl-service               2.4.0                    pypi_0    pypi
[conda] numpy                     1.22.4                   pypi_0    pypi
[conda] torch                     1.12.1                   pypi_0    pypi
[conda] torchaudio                0.12.1                   pypi_0    pypi
[conda] torchvision               0.13.1                   pypi_0    pypi
[conda] torchviz                  0.0.2                    pypi_0    pypi

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu