Repository: pytorch
Issue: SAC doesn't support nesting with different recompute plans
Labels: module: checkpoint,module: autograd,triaged,needs design

Using selective activation checkpoint (SAC) in a nested fashion came up because it might be useful when using float8 in conjunction with torchtitan:

(1) The torchtitan repo uses SAC to recompute all ops [except matmuls](https://github.com/pytorch/torchtitan/blob/main/torchtitan/parallelisms/parallelize_llama.py#L205)

(2) When torchtitan uses `Float8Linear` layers, however, this will cause any `torch.abs(torch.max())` calls that float8 uses during quantization to be recomputed in the backward.

One option is to tweak the outer SAC region linked above to also mark `abs/max` ops as always saved for backward (at least for max, if you are reducing to a single tensor then saving it for backward is cheap). But in general, it might be cleaner to express these two cases separately: user code that calls `aten.abs` (or `aten.max` in a situation where might not actually be better to save) may want to be treated differently from an inner `Float8Linear` layer that always knows its max() should be saved and not recomputed.

Another alternative would just be to make sure that the partitioner figures out that it should not recompute the `aten.max()` call from `Float8Linear`, since (a) it was tagged with `PREFER_RECOMPUTE` (giving the partitioner flexibility to ignore the user intent), and (b) it should hopefully be clear that saving it is better.

A third option (the reason for this issue) is to nest SAC: e.g., have the outer region express that matmuls must be saved for backward, and have a smaller region, local to `Float8Linear`, that just says "always recompute abs/max, leave all other ops alone". In particular, the outer SAC will mark `aten.max` with `PREFER_RECOMPUTE`, while the inner SAC will mark `aten.max` with `MUST_SAVE`, which should override the PREFER annotation.

I'm not sure that this works today though. Small repro:

```
import functools
import torch
from torch.utils.checkpoint import (
    CheckpointPolicy,
    create_selective_checkpoint_contexts,
)


def _save_sin(ctx, func, *args, **kwargs):
    return CheckpointPolicy.MUST_SAVE if func in [torch.ops.aten.sin.default] else CheckpointPolicy.PREFER_RECOMPUTE

def _save_cos(ctx, func, *args, **kwargs):
    return CheckpointPolicy.MUST_SAVE if func in [torch.ops.aten.cos.default] else CheckpointPolicy.PREFER_RECOMPUTE

def save_sin():
    return create_selective_checkpoint_contexts(_save_sin)
def save_cos():
    return create_selective_checkpoint_contexts(_save_cos)

def g(tmp):
    return tmp.sin().cos().sin()

def f(x):
    out1 = x.sin().cos().sin()
    out2 = torch.utils.checkpoint.checkpoint(g, x+1, context_fn=save_cos, use_reentrant=False)
    return out1, out2

@torch.compile(backend="aot_eager")
def f_checkpointed(x):
    return torch.utils.checkpoint.checkpoint(f, x, context_fn=save_sin, use_reentrant=False)

#torch.cuda.memory._record_memory_history(max_entries=100)
x = torch.randn(16, 16, device='cuda', requires_grad=True)
out = f_checkpointed(x)
#torch.cuda.memory._dump_snapshot(f"mem_prof.pickle")
```

The outer `f()` is supposed to save its first `sin()` output, and the inner `g()` should save its first `cos()` output. When I run with compile, I get this fw graph:
```
===== Forward graph 0 =====
 /home/hirsheybar/local/a/pytorch/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
    def forward(self, primals_1: "f32[16, 16][16, 1]cuda:0"):
         # File: /home/hirsheybar/local/a/pytorch/tmp3.py:27 in f, code: out1 = x.sin().cos().sin()
        sin: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.sin.default(primals_1)
        cos: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.cos.default(sin)
        sin_1: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.sin.default(cos);  cos = None

         # File: /home/hirsheybar/local/a/pytorch/tmp3.py:28 in f, code: out2 = torch.utils.checkpoint.checkpoint(g, x+1, context_fn=save_cos, use_reentrant=False)
        add: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.add.Tensor(primals_1, 1)

         # File: /home/hirsheybar/local/a/pytorch/tmp3.py:24 in g, code: return tmp.sin().cos().sin()
        sin_2: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.sin.default(add);  add = None
        cos_1: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.cos.default(sin_2)
        sin_3: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.sin.default(cos_1);  cos_1 = None
        return (sin_1, sin_3, primals_1, sin, sin_2)
        
 ===== Backward graph 0 =====
 <eval_with_key>.3 class GraphModule(torch.nn.Module):
    def forward(self, primals_1: "f32[16, 16][16, 1]cuda:0", sin: "f32[16, 16][16, 1]cuda:0", sin_2: "f32[16, 16][16, 1]cuda:0", tangents_1: "f32[16, 16][16, 1]cuda:0", tangents_2: "f32[16, 16][16, 1]cuda:0"):
         # File: /home/hirsheybar/local/a/pytorch/tmp3.py:24 in g, code: return tmp.sin().cos().sin()
        cos_1: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.cos.default(sin_2)
        cos_2: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.cos.default(cos_1);  cos_1 = None
        mul: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.mul.Tensor(tangents_2, cos_2);  tangents_2 = cos_2 = None
        sin_4: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.sin.default(sin_2);  sin_2 = None
        neg: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.neg.default(sin_4);  sin_4 = None
        mul_1: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.mul.Tensor(mul, neg);  mul = neg = None

         # File: /home/hirsheybar/local/a/pytorch/tmp3.py:28 in f, code: out2 = torch.utils.checkpoint.checkpoint(g, x+1, context_fn=save_cos, use_reentrant=False)
        add: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.add.Tensor(primals_1, 1)

         # File: /home/hirsheybar/local/a/pytorch/tmp3.py:24 in g, code: return tmp.sin().cos().sin()
        cos_3: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.cos.default(add);  add = None
        mul_2: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.mul.Tensor(mul_1, cos_3);  mul_1 = cos_3 = None

         # File: /home/hirsheybar/local/a/pytorch/tmp3.py:27 in f, code: out1 = x.sin().cos().sin()
        cos: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.cos.default(sin)
        cos_4: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.cos.default(cos);  cos = None
        mul_3: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.mul.Tensor(tangents_1, cos_4);  tangents_1 = cos_4 = None
        sin_5: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.sin.default(sin);  sin = None
        neg_1: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.neg.default(sin_5);  sin_5 = None
        mul_4: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.mul.Tensor(mul_3, neg_1);  mul_3 = neg_1 = None
        cos_5: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.cos.default(primals_1);  primals_1 = None
        mul_5: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.mul.Tensor(mul_4, cos_5);  mul_4 = cos_5 = None

         # File: /home/hirsheybar/local/a/pytorch/tmp3.py:27 in f, code: out1 = x.sin().cos().sin()
        add_1: "f32[16, 16][16, 1]cuda:0" = torch.ops.aten.add.Tensor(mul_2, mul_5);  mul_2 = mul_5 = None
        return (add_1,)
``` 

Where it looks like we properly saved the `sin()` outputs for bw, but we are not saving the `cos()` output from `g()` (and are recomputing it instead).

Today, SAC is implemented with a set of `TorchDispatchModes` ([link](https://github.com/pytorch/pytorch/blob/main/torch/utils/checkpoint.py#L1270)). When running with nested regions of SAC, we will end up with multiple instances of these modes on the mode stack, each with their own separate recompute logic. If we want to support this, we would need some way of:

(1) allowing all of these modes to run, and get a chance to specify their recompute preferences (we seem to already [do this](https://github.com/pytorch/pytorch/blob/0e19522122b0d1aa36ac4eceb53d1d5d2cf1caf9/torch/utils/checkpoint.py#L1292))

(2) aggregating the results in the outer-most mode to make a final call about what should be recomputed or saved. One annoyance is that we will need to figure out a way to share this state properly across the modes, e.g. in some side-car state.

One meta question: is this something we want to support? (I imagine this might be useful to the float8/torchtitan case, cc @soulitzer @ezyang @albanD @gqchen @pearu @nikitaved @Varal7 @xmfan @vkuzo, although they might have some other workarounds in mind).