Repository: pytorch
Issue: torch.compile has significantly higher peak memory for a simple checkpointed layer of matmul + layernorm + gelu
Labels: triaged,oncall: pt2,module: inductor

### 🐛 Describe the bug

```
import torch
from torch import nn
import torch.nn.functional as F
import torch.autograd as autograd
torch.set_default_device('cuda')

def bench(f, name=None, iters=1000, warmup=5, display=True, profile=False):
    import time
    from triton.testing import do_bench

    for _ in range(warmup):
        f()
    if profile:
        with torch.profiler.profile() as prof:
            f()
        prof.export_chrome_trace(f"{name if name is not None else 'trace'}.json")

    us_per_iter = do_bench(lambda: f())*1000

    if name is None:
        res = us_per_iter
    else:
        res= f"{name}: {us_per_iter:.3f}us"

    if display:
        print(res)
    return res

class Layer(nn.Module):
    def __init__(self):
        super().__init__()
        dim = 2048
        self.mods = nn.Sequential(nn.Linear(dim, dim), nn.LayerNorm([dim]), nn.GELU())

    def forward(self, x):
        return self.mods(x)

class Layered(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(*[Layer() for _ in range(5)])

    def forward(self, x):
        return self.layers(x)

mod = Layered().to(dtype=torch.bfloat16).cuda()
for idx, layer in enumerate(mod.layers):
    old_forward = mod.layers[idx].forward
    def foo(old_forward):
        def wrapped_fun(x):
            return torch.utils.checkpoint.checkpoint(old_forward, x, use_reentrant=False)
        return wrapped_fun
    mod.layers[idx].forward = foo(old_forward)

comp_mod = torch.compile(mod, fullgraph=True)

inp = torch.randn(2048*128, 2048, dtype=torch.bfloat16).cuda()
def clear_grads_and_run(mod, f):
    def foo():
        for param in mod.parameters():
            param.grad = None
        f()
    return foo

comp_mod(inp).pow(2).sum().backward()
for param in mod.parameters():
    param.grad = None
for param in comp_mod.parameters():
    param.grad = None

torch.cuda.reset_peak_memory_stats()

bench(clear_grads_and_run(mod, lambda: mod(inp).sum().backward()), name="eager")
print(torch.cuda.max_memory_allocated()/1e9)
torch.cuda.reset_peak_memory_stats()

bench(clear_grads_and_run(comp_mod, lambda: comp_mod(inp).sum().backward()), name="compiled")
print(torch.cuda.max_memory_allocated()/1e9)
```

Results in
```
eager: 242199.387us
8.915477504
compiled: 244289.825us
13.200999424
```

So, not only do we have significantly higher peak memory - we also have a slower runtime :'( Not great.

The peak memory usage repros at one layer. Looking at the memory trace, we see that it's due to significantly higher peak memory being materialized in the backwards

<img width="1671" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/5c4c30f5-7e44-4543-88da-88d981cf5818">

cc: @anijain2305 @eellison @shunting314 @jansel 

### Versions

N/A

cc @ezyang @msaroufim @bdhirsh @anijain2305 @zou3519 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler @amjames @desertfire