Repository: pytorch
Issue: Backward hangs with DDP during training.
Labels: oncall: distributed,triaged,module: ddp

Hello, I am trying to train a network using DDP. The architecture of the network is such that it consists of two sub-networks (a, b) and depending on input either only a or only b or both a and b get executed. Things work fine on a single GPU. But when expanding the network to 2 or more GPUS the backward just hangs. Any help is appreciated. Thanks.

Here is a minimal example,

```
import numpy as np
import os
import torch
import torch.nn as nn
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.optim as optim

class SampleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.a1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.a2 = nn.ReLU(inplace=True)
        self.a3 = nn.Conv2d(64, 2, kernel_size=3, padding=1)


        self.b1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.b2 = nn.ReLU(inplace=True)
        self.b3 = nn.Conv2d(64, 2, kernel_size=3, padding=1)

    def forward(self, data):
        output = {}
        if 'a' in data.keys():
            output['a'] = self.a3(self.a2(self.a1(data['a'])))
        if 'b' in data.keys():
            output['b'] = self.b3(self.b2(self.b1(data['b'])))
        return output

def main():
    mp.set_start_method('spawn')
    rank = int(os.environ['RANK'])
    num_gpus = torch.cuda.device_count()
    torch.cuda.set_device(rank % num_gpus)
    dist.init_process_group(backend='nccl')
    model = SampleNet().cuda()
    ddp_model = torch.nn.parallel.DistributedDataParallel(
                    module=model,
                    device_ids=[torch.cuda.current_device()],
                    broadcast_buffers=False,
                    find_unused_parameters=True)
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
    print('Finetuneing the network... on gpu ', rank)
    
    for i in range(0,20):
        optimizer.zero_grad()
        f = np.random.rand()
        if f < 0.33:
            data = {'a':torch.randn(4,3,128,128), 'b':torch.randn(2,3,32,32)}
            outputs = ddp_model(data)
            labels_a = torch.randn(4,2,128,128).to(rank)
            loss_a = loss_fn(outputs['a'], labels_a)
            labels_b = torch.randn(2,2,32 ,32 ).to(rank)
            loss_b = loss_fn(outputs['b'], labels_b)
            (loss_a + loss_b).backward()
            print(i, loss_a, loss_b)
        elif f < 0.66 and f > 0.33:
            data = {'a':torch.randn(3,3,128,128)}
            outputs = ddp_model(data)
            labels_a = torch.randn(3,2,128,128).to(rank)
            loss_a = loss_fn(outputs['a'], labels_a)
            loss_a.backward()
            print(i, loss_a)
        else:
            data = {'b':torch.randn(5,3,32,32)}
            outputs = ddp_model(data)
            labels_b = torch.randn(5,2,32 ,32 ).to(rank)
            loss_b = loss_fn(outputs['b'], labels_b)
            loss_b.backward()
            print(i, loss_b)
        optimizer.step()
    dist.destroy_process_group()

if __name__ == '__main__':
    main()
```

Any pointers how to fix this? thanks for the help.

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