Lightning DDP seems to be breaking autograd
🐛 Bug

I am attempting to make a lightning script for the example code in https://github.com/lucidrains/reformer-pytorch
I now have 2 scripts, 1 is a lightning implementation and one is a regular apex DDP implementation,  The regular apex DDP implementation is able to train completly fine, but lightning throws an error about autograd warning that parameters are unused (see error message section)
I enabled find_unused_parameters in the non lightning example and do not get any errors, and disabling find_unused_parameters in the lightning example causes an autograd crash
To Reproduce
Steps to reproduce the behavior:
I personally use the nvidia pytorch container because it has most things installed

git clone https://github.com/lucidrains/reformer-pytorch.git
pip install revtorch
I've attached a zip file containing two training scripts, unzip them to the examples folder of reformer-pytorch
python -m torch.distributed.launch --nproc_per_node=1 example/train_apex_ddp.py -b 4 and note there is training and no crashing
python example/train_lightning.py --distributed to see the crash

-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "/opt/conda/lib/python3.6/site-packages/torch/multiprocessing/spawn.py", line 19, in _wrap
    fn(i, *args)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/distrib_data_parallel.py", line 331, in ddp_train
    self.run_pretrain_routine(model)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 829, in run_pretrain_routine
    self.train()
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/training_loop.py", line 332, in train
    self.run_training_epoch()
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/training_loop.py", line 386, in run_training_epoch
    output = self.run_training_batch(batch, batch_idx)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/training_loop.py", line 506, in run_training_batch
    loss = optimizer_closure()
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/training_loop.py", line 489, in optimizer_closure
    model_ref.backward(self.use_amp, closure_loss, optimizer)
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/core/hooks.py", line 154, in backward
    loss.backward()
  File "/opt/conda/lib/python3.6/site-packages/torch/tensor.py", line 195, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/opt/conda/lib/python3.6/site-packages/torch/autograd/__init__.py", line 99, in backward
    allow_unreachable=True)  # allow_unreachable flag
  File "/opt/conda/lib/python3.6/site-packages/torch/autograd/function.py", line 77, in apply
    return self._forward_cls.backward(self, *args)
  File "/opt/conda/lib/python3.6/site-packages/revtorch/revtorch.py", line 161, in backward
    y, dy = ctx.reversible_blocks[i].backward_pass(y, dy)
  File "/opt/conda/lib/python3.6/site-packages/revtorch/revtorch.py", line 89, in backward_pass
    gy1.backward(dy2)
  File "/opt/conda/lib/python3.6/site-packages/torch/tensor.py", line 195, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/opt/conda/lib/python3.6/site-packages/torch/autograd/__init__.py", line 99, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: Expected to mark a variable ready only once. This error is caused by use of a module parameter outside the `forward` function. The return value of the `forward` function is inspected by the distributed data parallel wrapper to figure out if any of the module's parameters went unused. If this is the case, it knows they won't receive gradients in a backward pass. If any of those parameters are then used outside `forward`, this error condition is triggered. You can disable unused parameter detection by passing the keyword argument `find_unused_parameters=False` to `torch.nn.parallel.DistributedDataParallel`.

Code sample
training_scripts.zip
I'm sorry I was unable to make the code samples any smaller
Environment
Collecting environment information...
PyTorch version: 1.4.0a0+a5b4d78
Is debug build: No
CUDA used to build PyTorch: 10.2

OS: Ubuntu 18.04.3 LTS
GCC version: (Ubuntu 7.4.0-1ubuntu1~18.04.1) 7.4.0
CMake version: version 3.14.0

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: Tesla V100-SXM3-32GB
GPU 1: Tesla V100-SXM3-32GB
GPU 2: Tesla V100-SXM3-32GB
GPU 3: Tesla V100-SXM3-32GB

Nvidia driver version: 418.67
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5

Versions of relevant libraries:
[pip] msgpack-numpy==0.4.3.2
[pip] numpy==1.17.4
[pip] pytorch-lightning==0.6.0
[pip] pytorch-transformers==1.1.0
[pip] revtorch==0.2.3
[pip] torch==1.4.0a0+a5b4d78
[pip] torchtext==0.4.0
[pip] torchvision==0.4.2
[conda] magma-cuda101             2.5.1                         1    local
[conda] mkl                       2019.1                      144  
[conda] mkl-include               2019.1                      144  
[conda] nomkl                     3.0                           0  
[conda] pytorch-lightning         0.6.0                    pypi_0    pypi
[conda] pytorch-transformers      1.1.0                    pypi_0    pypi
[conda] revtorch                  0.2.3                    pypi_0    pypi
[conda] torch                     1.4.0a0+a5b4d78          pypi_0    pypi
[conda] torchtext                 0.4.0                    pypi_0    pypi
[conda] torchvision               0.4.2                    pypi_0    pypi

Additional context
Revtorch has custom autograd, but I have tested revtorch and am able  to train a small revtorch example using lighting, so I don't think the issue is revtorch