Repository: pytorch
Issue: DDP fails if you have multiple forward passes and a single backwards pass with `find_unused_parameters=True`
Labels: oncall: distributed,module: ddp

## 🐛 Bug

We have a multitask model with different output heads and a shared encoder. Each output head operates on a different dataset, producing a different loss (e.g. one head could be classification on one dataset and another could be ranking on a different dataset), and the model may contain output heads which are not used in a particular training run. As we may have entirely unused output heads we need to supply `find_unused_parameters=True` when constructing DDP.

Our training procedure receives a batch of tensors, one per output head we're training in this run, and we call forward on each batch separately. Each output head produces its own loss, which we combine using `loss = torch.mean(torch.stack(losses))` outside the DDP forward call, and then call `loss.backwards()` on. This triggers an error message which is pasted below and the computation fails.

While triaging this issue we determined it's due to the combination of unused output heads and multiple forward passes. If we remove the unused output heads so each remaining output head is used once before the backward call, and set `find_unused_parameters=False` then DDP correctly aggregates the gradients. If we have `find_unused_parameters=True`, even when all output heads are used in a backward call we still get the same error. In some of our other multitask model experiments where we always have a single forward call followed by a single backward call then DDP works, even when there are multiple output heads and not all of them are used (assuming `find_unused_parameters=True`).

The error message returned is a bit confusing as we aren't using the model parameters outside of the forward method, nor do we have multiple re-entrant backwards passes.

## To Reproduce

Steps to reproduce the behavior:

1. Have some unused parameters in a model.
2. Wrap the model in DDP with `find_unused_parameters=True`.
3. Call forward at least twice with two different batches which use two different output losses in the model.
4. Aggregate the losses via sum or mean.
5. Call `loss.backwards()`. 

I snipped out the top of the traceback stack as it refers to internal code.

```   Traceback (most recent call last):
     File "/path/to/code/trainers/combine.py", line 183, in _training_step_base
       loss.backward()
     File "/path/to/venv/lib64/python3.6/site-packages/torch/tensor.py", line 245, in backward
       torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
     File "/path/to/venv/lib64/python3.6/site-packages/torch/autograd/__init__.py", line 147, in backward
       allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag
   RuntimeError: Expected to mark a variable ready only once. This error is caused by one of the following reasons:
1) Use of a module parameter outside the `forward` function. Please make sure model parameters are not shared across multiple concurrent forward-backward passes
2) Reused parameters in multiple reentrant backward passes. For example, if you use multiple `checkpoint` functions to wrap the same part of your model, it would result in the same set of parameters been used by different reentrant backward passes multiple times, and hence marking a variable ready multiple times. DDP does not support such use cases yet.
```

## Expected behavior

Gradients to be aggregated across processes for the encoder and set of output heads that we are training.

## Environment

 - PyTorch Version (e.g., 1.0): 1.8.0+cu111
 - OS (e.g., Linux): Oracle Linux 7
 - How you installed PyTorch (`conda`, `pip`, source): pip
 - Build command you used (if compiling from source): n/a
 - Python version: Python 3.6
 - CUDA/cuDNN version: 11.1
 - GPU models and configuration: V100

I have tested this on pytorch 1.9 as well, where it fails giving slightly different error messages, or hangs the process.

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