Repository: pytorch
Issue: Activation Checkpointing composability with split backward computation
Labels: oncall: distributed,triaged,activation-checkpointing

Activation checkpointing avoids saving intermediate tensors in order to save memory. It does so by recomputing the forward pass on demand to obtain the intermediate values required for gradient computation during backward.

For pipelining, we are splitting up the backward computation into `stage_backward_input` and `stage_backward_weight` (https://github.com/pytorch/pytorch/blob/main/torch/distributed/pipelining/_backward.py) which represent the derivative with respect to input and the derivative with respect to weights, respectively. This is mentioned in more detail in the [ZeroBubble paper](https://arxiv.org/pdf/2401.10241). We are computing backwards for each parameter group which leads to repeated forward computation that is impacting performance.


## Minimal example
```python
import torch
from torch.utils.checkpoint import checkpoint

a = torch.tensor(1., requires_grad=True)
b = torch.tensor(1., requires_grad=True)

def fn(x, z):
    print("computing forward")
    y = x.sin().exp() * z
    return y, y.sin().exp()

# compute forward under AC
y, out = checkpoint(fn, a, b, use_reentrant=False)

# 1. compute wrt y and x
da, dy = torch.autograd.grad(out, inputs=(a, y), retain_graph=True)

# 2. compute wrt z using y's grad
db = torch.autograd.grad((y,), inputs=(b,), grad_outputs=(dy,))
```

Output:
```
computing forward
computing forward
computing forward
```

## torch.distributed.pipelining minimal example

```python
import torch
from torch.distributed.pipelining._backward import stage_backward_input, stage_backward_weight
import torch.nn as nn
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper


def example(checkpoint):
    torch.manual_seed(0)
    x = torch.tensor((1.,), requires_grad=True)
    class SimpleModel(nn.Module):
        def __init__(self, num_layers=5):
            super().__init__()
            self.layers = nn.Sequential(*[nn.Linear(1, 1) for _ in range(num_layers)])

        def forward(self, x):
            print("computing forward")
            return self.layers(x)

    model = SimpleModel()
    if checkpoint:
        model = checkpoint_wrapper(model)

    y = model(x)
    loss = y.mean()

    dx, param_groups = stage_backward_input([loss], output_grads=None, input_values=[x], weights=model.parameters())
    print(f"{len(param_groups)=}")
    stage_backward_weight(model.parameters(), param_groups=param_groups)

print("without AC:", "=" * 25)
example(checkpoint=False)
print("with AC:", "=" * 25)
example(checkpoint=True)
```

Output:
```
without AC: =========================
computing forward
len(param_groups)=5
with AC: =========================
computing forward
computing forward
len(param_groups)=5
computing forward
computing forward
computing forward
computing forward
computing forward
```

## Potential solution

@soulitzer mentioned potentially adding a flag to AC to cache the recomputed intermediates. Still need discussion on how the API would look like for that and how we could explictly free the intermediates when we decide we are done. 

For pipelining, saving the intermediates would increase the lifespan in memory for the time between `backward_input` and `backward_weight`, but this should still help peak memory. Ideally we could have a situation where forward only needs to be called once for `backward_input` and once for `backward_weight`.

cc @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o