Repository: pytorch
Issue: Using  `torch.utils.checkpoint.checkpoint_sequential` and `torch.autograd.grad` breaks when used in combination with `DistributedDataParallel`
Labels: oncall: distributed,module: checkpoint,feature,triaged

## 🐛 Bug

Using  `torch.utils.checkpoint.checkpoint_sequential` and `torch.autograd.grad` breaks when used in combination with `DistributedDataParallel` resulting in the following stacktrace

```
Traceback (most recent call last):                                                                    
  File "minimal_buggy_2.py", line 198, in <module>                                                    
    train(hps)                                                                                        
  File "minimal_buggy_2.py", line 179, in train                                                       
    loss.backward()                                                                                   
  File "/opt/conda/lib/python3.7/site-packages/torch/tensor.py", line 107, in backward                
    torch.autograd.backward(self, gradient, retain_graph, create_graph)                               
  File "/opt/conda/lib/python3.7/site-packages/torch/autograd/__init__.py", line 93, in backward      
    allow_unreachable=True)  # allow_unreachable flag                                                 
  File "/opt/conda/lib/python3.7/site-packages/torch/autograd/function.py", line 77, in apply         
    return self._forward_cls.backward(self, *args)                                                    
  File "/opt/conda/lib/python3.7/site-packages/torch/utils/checkpoint.py", line 99, in backward       
    torch.autograd.backward(outputs, args)                                                            
  File "/opt/conda/lib/python3.7/site-packages/torch/autograd/__init__.py", line 93, in backward      
    allow_unreachable=True)  # allow_unreachable flag                                                 
RuntimeError: has_marked_unused_parameters_ ASSERT FAILED at /pytorch/torch/csrc/distributed/c10d/reducer.cpp:181, please report a bug to PyTorch.
```

## To Reproduce

Steps to reproduce the behavior:

1. Running `python -m torch.distributed.launch --nproc_per_node=2 minimal_buggy.py` works
1. Running `python -m torch.distributed.launch --nproc_per_node=2 minimal_buggy.py --parallel` breaks

The code for `minimal_buggy.py` is here:
```python
from argparse import ArgumentParser, Namespace
import numpy as np

import torch
import torch.nn as nn
import torch.distributed as dist
from torch.utils.checkpoint import checkpoint_sequential

import torchvision.datasets as ds
from torchvision import transforms


class GradientStep(nn.Module):

    def __init__(self, energy):
        super(GradientStep, self).__init__()
        self._energy = energy
        self.step_size = 0.1

    def forward(self, x: torch.Tensor):
        with torch.enable_grad():
            x.requires_grad_()
            omega = self._energy(x).sum()
            grad_out = torch.ones_like(omega).to(x.device)
            dx = torch.autograd.grad(outputs=omega, inputs=x, grad_outputs=grad_out,
                                     create_graph=True, retain_graph=True,
                                     allow_unused=True)[0]
            dx.requires_grad_()
            x = (x - self.step_size ** 2 * dx)
        return x


class FFN(nn.Module):

    def __init__(self):
        super(FFN, self).__init__()
        self.n_hidden = 1024
        self._energy = nn.Sequential(
            nn.Linear(28**2, self.n_hidden),
            nn.LeakyReLU(),
            nn.Linear(self.n_hidden, self.n_hidden),
            nn.LeakyReLU(),
            nn.Linear(self.n_hidden, self.n_hidden),
            nn.LeakyReLU(),
            nn.Linear(self.n_hidden, 1))
        self.L = 10

    def forward(self, x: torch.Tensor):
        y = x.clone().to(x.device)
        y.requires_grad_()
        fwd = nn.Sequential(*[GradientStep(self._energy) for _ in range(self.L)])
        y = checkpoint_sequential(fwd, self.L, y)
        return y


def get_distributed_mnist_iterators(batch_size, **kwargs):
    def _worker_init_fn(worker_id):
        np.random.seed(np.random.get_state()[1][0] + worker_id)

    base_transforms = [
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Lambda(lambda x: x.reshape(-1, ))]

    train_dataset = ds.MNIST(
        '/tmp/mnist/train',
        train=True,
        download=True,
        transform=transforms.Compose(base_transforms))

    test_dataset = ds.MNIST(
        '/tmp/mnist/test',
        train=False,
        download=True,
        transform=transforms.Compose(base_transforms))

    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
    test_sampler = torch.utils.data.distributed.DistributedSampler(test_dataset)

    train_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=kwargs.get('workers', 4),
        pin_memory=True,
        sampler=train_sampler,
        worker_init_fn=_worker_init_fn)

    test_loader = torch.utils.data.DataLoader(
        test_dataset,
        batch_size=kwargs.get('test_batch_size', 100),
        shuffle=False,
        num_workers=kwargs.get('workers', 4),
        pin_memory=True,
        sampler=test_sampler,
        worker_init_fn=_worker_init_fn)

    return train_loader, test_loader, train_sampler, test_sampler


def get_mnist_iterators(batch_size, **kwargs):

    base_transforms = [
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Lambda(lambda x: x.reshape(-1, ))]

    train_dataset = ds.MNIST(
        '/tmp/mnist/train',
        train=True,
        download=True,
        transform=transforms.Compose(base_transforms))

    test_dataset = ds.MNIST(
        '/tmp/mnist/test',
        train=False,
        download=True,
        transform=transforms.Compose(base_transforms))

    train_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=kwargs.get('workers', 4),
        pin_memory=True,
        sampler=None)

    test_loader = torch.utils.data.DataLoader(
        test_dataset,
        batch_size=kwargs.get('test_batch_size', 100),
        shuffle=True,
        num_workers=kwargs.get('workers', 4),
        pin_memory=True,
        sampler=None)

    return train_loader, test_loader


def parse_args() -> Namespace:
    parser = ArgumentParser()
    parser.add_argument('--local_rank',
                            type=int,
                            default=0,
                            help="Is being set by the pytorch distributed launcher")

    parser.add_argument('--parallel',
                        default=False,
                        action='store_true')

    hps = parser.parse_args()
    return hps


def train(hps: Namespace):

    if hps.parallel:
        train_loader, test_loader, _, _ = get_distributed_mnist_iterators(
            batch_size=32)
    else:
        train_loader, test_loader = get_mnist_iterators(
            batch_size=32)

    model = FFN()
    model.cuda()
    if hps.parallel:
        model = nn.parallel.DistributedDataParallel(model)
    crit = nn.MSELoss()

    opt = torch.optim.Adam(model.parameters(), lr=1e-4)
    for epoch in range(100):
        for step, (b, lbl) in enumerate(train_loader):

            model.train()
            opt.zero_grad()
            model.zero_grad()
            corrupt_b = (b + 0.3 * torch.randn_like(b)).cuda()
            recons = model(corrupt_b)
            loss = crit(recons, b.cuda())
            loss.backward()
            opt.step()

            if dist.get_rank() == 0:
                if step % 10 == 0:
                    print(f'epoch {epoch}, batch: {step}, loss: {float(loss.cpu())}')


def setup():
    hps = parse_args()
    dist.init_process_group(backend='nccl', init_method=f'env://', rank=hps.local_rank)

    size = dist.get_world_size()
    group = torch.distributed.new_group(ranks=list(range(size)))
    return group, hps


if __name__ == "__main__":
    group, hps = setup()
    train(hps)
```

## Expected behavior

Expect either command to work independent of using `DistributedDataParallel` or not.

## Environment

```
PyTorch version: 1.1.0
Is debug build: No
CUDA used to build PyTorch: 9.0.176

OS: Ubuntu 16.04.5 LTS
GCC version: (Ubuntu 5.4.0-6ubuntu1~16.04.11) 5.4.0 20160609
CMake version: version 3.11.1

Python version: 3.7
Is CUDA available: Yes
CUDA runtime version: 10.0.130
GPU models and configuration:
GPU 0: GeForce GTX 1080 Ti
GPU 1: GeForce GTX 1080 Ti

Nvidia driver version: 410.79
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.4.2

Versions of relevant libraries:
[pip3] numpy==1.16.3
[pip3] pytorch-memlab==0.0.3
[pip3] torch==1.1.0
[pip3] torchvision==0.3.0
[conda] pytorch-memlab            0.0.3                    pypi_0    pypi
[conda] torch                     1.1.0                    pypi_0    pypi
[conda] torchvision               0.3.0                    pypi_0    pypi
```

## Additional context

<!-- Add any other context about the problem here. -->


cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @xush6528 @osalpekar @jiayisuse @agolynski