Repository: pytorch
Issue: DDP Gradient reduction is not triggered after calling autograd.backward()
Labels: oncall: distributed

## 🐛 Bug
Hi everyone,  I found that autograd.backward() doesn’t trigger reduction under certain circumstance when using DDP. 
@mrshenli @rohan-varma @pritamdamania 

## To Reproduce
Here is a minimal code to reproduce. I initialized two linear networks with constant so that everything is deterministic here.
Remark: run this code in Pytorch version <= 1.8.1. 

```python
import torch
import torch.multiprocessing as mp
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.autograd as autograd
import torch.distributed as dist
import os
from argparse import ArgumentParser


def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '7777'
    # initialize the process group
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

def zero_grad(params):
    '''
    Clean the gradient of each parameter
    '''
    for p in params:
        if p.grad is not None:
            p.grad.detach()
            p.grad.zero_()

def collect_grad(params):
    '''
    Collect grads of parameters and concatenate them into a vector.
    If grad is None, it will be filled with zeros
    :param params: list of parameters
    :return: vector
    '''
    grad_list = []
    for p in params:
        if p.grad is not None:
            grad_list.append(p.grad.contiguous().view(-1))
            del p.grad
        else:
            # replace None with zeros
            grad_list.append(torch.zeros_like(p).view(-1))
    return torch.cat(grad_list)

def subprocess_fn(rank, args):
    setup(rank, args.num_gpus)
    print(f'Running on rank {rank}')
    D = nn.Linear(2, 1, bias=True).to(rank)
    G = nn.Linear(1, 2, bias=True).to(rank)
    # initialize weights
    nn.init.constant_(D.weight, 2.0)
    nn.init.constant_(D.bias, -1.0)
    nn.init.constant_(G.weight, 4.0)
    nn.init.constant_(G.bias, 1.0)

    if args.distributed:
        G = DDP(G, device_ids=[rank], broadcast_buffers=False)
        D = DDP(D, device_ids=[rank], broadcast_buffers=False)

    d_params = list(D.parameters())
    g_params = list(G.parameters())

    if not args.distributed:
        z = torch.tensor([[2.0], [1.0]]).to(rank)
    elif rank == 0:
        z = torch.tensor([[1.0]]).to(rank)
    elif rank == 1:
        z = torch.tensor([[2.0]]).to(rank)

    loss = D(G(z)).mean()
    grad_g = autograd.grad(loss, g_params, create_graph=True)
    gradvec_g = torch.cat([g.contiguous().view(-1) for g in grad_g])
    zero_grad(d_params)
    autograd.backward(gradvec_g,
                      grad_tensors=torch.ones_like(gradvec_g),
                      inputs=d_params)
    hvp_d = collect_grad(d_params)
    print(f'Hessian vector product of d param: {hvp_d}')
    cleanup()


if __name__ == '__main__':
    torch.backends.cudnn.benchmark = True
    parser = ArgumentParser()
    parser.add_argument('--num_gpus', type=int, help='Number of GPUs', default=1)
    args = parser.parse_args()
    args.distributed = args.num_gpus > 1

    if args.distributed:
        mp.spawn(subprocess_fn, args=(args, ), nprocs=args.num_gpus)
    else:
        subprocess_fn(0, args)

    print('Done!')
```
Below is the output, which shows that gradient reduction is not triggered in this situation. 

```
Running on rank 0
Running on rank 1
Hessian vector product of d param: tensor([3., 3., 0.], device='cuda:1')
Hessian vector product of d param: tensor([2., 2., 0.], device='cuda:0')
Done!
```

## Expected behavior
DDP is supposed to do gradient reduction after autograd.backward() is called. The expected output should be 
```
Running on rank 0
Running on rank 1
Hessian vector product of d param: tensor([2.5., 2.5, 0.], device='cuda:1')
Hessian vector product of d param: tensor([2.5, 2.5, 0.], device='cuda:0')
Done!
```

## Environment

```
Collecting environment information...
PyTorch version: 1.8.0
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.10.2
Libc version: glibc-2.27

Python version: 3.8.8 (default, Apr 13 2021, 19:58:26)  [GCC 7.3.0] (64-bit runtime)
Python platform: Linux-4.15.0-154-generic-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: 11.1.74
GPU models and configuration: 
GPU 0: Tesla V100-DGXS-32GB
GPU 1: Tesla V100-DGXS-32GB
GPU 2: Tesla V100-DGXS-32GB
GPU 3: Tesla V100-DGXS-32GB

Nvidia driver version: 455.23.05
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.20.1
[pip3] numpydoc==1.1.0
[pip3] torch==1.8.0
[pip3] torchaudio==0.8.0a0+a751e1d
[pip3] torchvision==0.9.0
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               10.2.89              hfd86e86_1  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.2.0           h06a4308_296  
[conda] mkl-service               2.3.0            py38h27cfd23_1  
[conda] mkl_fft                   1.3.0            py38h42c9631_2  
[conda] mkl_random                1.2.1            py38ha9443f7_2  
[conda] mypy_extensions           0.4.3                    py38_0  
[conda] numpy                     1.20.1           py38h93e21f0_0  
[conda] numpy-base                1.20.1           py38h7d8b39e_0  
[conda] numpydoc                  1.1.0              pyhd3eb1b0_1  
[conda] pytorch                   1.8.0           py3.8_cuda10.2_cudnn7.6.5_0    pytorch
[conda] torchaudio                0.8.0                      py38    pytorch
[conda] torchvision               0.9.0                py38_cu102    pytorch

```

## Additional context
More context will be provided below


cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @agolynski @SciPioneer @H-Huang @mrzzd @cbalioglu @gcramer23