Repository: pytorch
Issue: [rfc] Trigger callback when backwards begins for DDP with custom autograd function
Labels: oncall: distributed,triaged,module: ddp

# [rfc] Trigger callback when backwards begins for DDP with custom autograd function

with @zhaojuanmao 

*Context*

A variety of improvements to `DistributedDataParallel`, namely around ensuring local training parity and enhancing performance  where possible for certain cases are planned or already in the works. In particular, with respect to local training parity, we are aiming to support the following use cases:

* Not all outputs of DDP `forward()` used in loss computation. Today, this results in an error during backward as these parameters are not marked as unused.
* Calling `backward()` twice. This is useful for a variety of use cases (i.e. computing different losses, loss with respect to different targets, etc) but DDP does not support this due to the way we account for autoggrad hooks.
* Activation checkpointing with multiple checkpoints, as well as activation checkpointing with unused parameters

For the above cases, it would be very useful to have a callback that fires when DDP user issues a call to backward, which helps DDP know that backwards is being called - currently, we only have access to the backwards pass via the autograd hooks installed on the reducer. Such a callback will allow the following features:

* Traverse graph beginning from outputs to find complete set of unused parameters (unused parameters means parameters that don’t get gradient in this backward pass). This set will include the parameters that don’t get gradient because a particular model output was unused when computing loss.
* `Calling backward()` twice, by resetting the appropriate flags each time our callback fires. This mainly involves moving `Reducer::prepare_for_backward` to be called by this function.
* Miscellaneous fixes such as moving backward start log time to this function.
* Enqueuing a delay_allreduce after backwards pass is over, which @zhaojuanmao is doing in https://github.com/pytorch/pytorch/pull/55248/. 

_Proposal_

In order to achieve this, the proposal is to implement a custom `torch.autograd.Function` as follows:

```
class PassThrough(Function):
    @staticmethod
    def forward(ctx, *inputs):
        # Set any needed attributes from DDP onto ctx so it can be used during backwards pass.
        ctx.save_for_backward(inputs)
        return inputs

    @staticmethod
    def backward(ctx, *grad_outputs):
        # Implement various functionality discussed above
        return grad_outputs
```

Calling this function on single-tensor DDP outputs is quite straightforward:

```
ret: torch.Tensor = DDP_forward(inputs)
ret = PassThrough.apply(ret)
return ret
```

However, the more interesting use case is models with non-tensor outputs. DDP supports `dict`, `tuple` and `list` data structures that contain tensors, in  nested fashion. Adapting the above, to, for example, a `tuple`:

```
ret: Tuple[torch.Tensor, torch.Tensor]= DDP_forward(inputs)
ret = PassThrough.apply(ret)
return ret
```

would *not* be correct, as the custom function only considers tensors in the input/output and not general python data structures, as per https://github.com/pytorch/pytorch/issues/55509. The result of this would be errenous since we would feed a python structure as input, and get back outputs that don't require grad.

Another non-solution is calling this function once for each tensor. Ideally, we would want this function called only once at the start of backwards,  and grad_outputs would contain gradients for each tensor that was output by the user’s forward pass (so we can do things like use zero gradients to determine unused parameters).

The following, for tuples, provides the expected behavior:

```
ret: torch.Tensor = DDP_forward(inputs)
a, b = ret
new_a, new_b = PassThrough.apply(a, b)
# We need to reassign the return value here, because a and b were not modified in place
ret = (new_a, new_b)
```
The above implies the main things we need to do are:

1. Get output data structure from user forward
2. Parse all tensors from output (we already have _find_tensors that does this)
3. Call PassThrough with all outputs as argument: tensor_outputs = PassThrough.apply(*tensor_outputs)
4. *Reconstruct* output return type, by iterating through the output data structure recursively and filling in the appropriate tensor_outputs. 

The following is a full prototype of this proposal for `tuple` data structures, which can be generalized to the output data structures that we support:

```
import torch
import torch.nn as nn
from torch.autograd import Function


class PassThrough(Function):
    @staticmethod
    def forward(ctx, *inputs):
        return inputs

    @staticmethod
    def backward(ctx, *grad_outputs):
        print(f"grad_outputs {grad_outputs}")
        # Use gradients to search through the graph as in
        # https://gist.github.com/rohan-varma/7c8dab3635193c04c607e67c4951f519
        return grad_outputs


class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.a = nn.Linear(1, 1, bias=False)
        self.b = nn.Linear(1, 1, bias=False)

    def forward(self, x):
        a, b = self.a(x), self.b(x)
        # Get tensors from tuple. This would be a more general call to
        # _find_tensors.
        ret = a, b
        new_a, new_b = PassThrough.apply(a, b)
        # Reconstruct tuple from output tensors. This would require a more general
        # function that repacks the tensor(s) into the data structure.
        ret = new_a, new_b
        return ret


model = MyModel()
inp = torch.ones(1)
out = model(inp)
loss = out[0].sum()
print("Calling backward...")
loss.backward()
print("Done with bwd")
```

Note that the custom function reports `0` for b's grad_output, which is what we expected since it was unused in loss computation.


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