Repository: pytorch
Issue: Python program locks up when using some innocuous code inside mp.spawn()
Labels: oncall: distributed,triaged

## 🐛 Bug

PyTorch/Python locks up when using mp.spawn() and DistributedDataParallel with some seemingly innocuous code.

## To Reproduce

Steps to reproduce the behavior:

1. Run the script below
2. Note that this locks up somewhere inside mp.spawn and never prints "DONE {BACKEND} ....". Note also that the memories for the GPU's are still loaded and being utilized.
3. If you then kill the locked python processes, remove the lines `if rank != 0: assert torch.equal(loss_res, loss_res)`, and rerun this script, everything works fine.

```
import os
import sys 
from time import time
import tempfile
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp

from torch.nn.parallel import DistributedDataParallel as DDP 


BACKEND = "nccl"
NUM_OBS = 10000
NUM_DIM = 300 
NUM_OUT = 5 


def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    dist.init_process_group(BACKEND, rank=rank, world_size=world_size)


def cleanup():
    dist.destroy_process_group()


class ToyModel(nn.Module):
    def __init__(self, hidden_dim=6000):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(NUM_DIM, hidden_dim)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(hidden_dim, NUM_OUT)

    def forward(self, x): 
        return self.net2(self.relu(self.net1(x)))


def demo_basic(rank, world_size, epochs=10):
    print(f"Running basic DDP example on rank {rank} with {world_size=}.")
    setup(rank, world_size)

    # create model and move it to GPU with id rank
    model = ToyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    for epoch in range(epochs):
        optimizer.zero_grad()
        outputs = ddp_model(torch.randn(NUM_OBS, NUM_DIM))
        labels = torch.randn(NUM_OBS, NUM_OUT).to(rank)
        loss_res = loss_fn(outputs, labels)
        print(f"{rank=} {epoch=} {loss_res.item()=}")
        loss_res.backward()
        optimizer.step()
        if rank != 0:
            assert torch.equal(loss_res, loss_res)

    cleanup()


def run_demo(demo_fn, world_size):
    mp.spawn(demo_fn,
             args=(world_size,),
             nprocs=world_size,
             join=True)


if __name__ == "__main__":
    n_gpus = torch.cuda.device_count()
    print(f"Running {n_gpus=}.")
    t0 = time()
    run_demo(demo_basic, n_gpus)
    d = time() - t0
    print(f"DONE {BACKEND=} {round(d, 2)}")
```

## Expected behavior

I expect this program to run through to completion without locking up.

## 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: 1.8.1+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.2 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.8.10 (default, Jun  2 2021, 10:49:15)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.8.0-1038-aws-x86_64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: Tesla V100-SXM2-16GB
GPU 1: Tesla V100-SXM2-16GB
GPU 2: Tesla V100-SXM2-16GB
GPU 3: Tesla V100-SXM2-16GB

Nvidia driver version: 460.80
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.21.0
[pip3] torch==1.8.1+cu111
[pip3] torchaudio==0.8.1
[pip3] torchvision==0.9.1+cu111
[conda] Could not collect

## Additional context

<!-- Add any other context about the problem here. -->


cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @agolynski @SciPioneer @H-Huang @mrzzd @cbalioglu @gcramer23