Repository: pytorch
Issue: retro inductor OOM
Labels: triaged,oncall: pt2,module: aotdispatch,module: pt2-dispatcher

## Repro
Another lucidrains model

` pip install retro-pytorch`

```python
import torch
from retro_pytorch import RETRO
import torchdynamo

retro = RETRO(
    chunk_size = 64,                         # the chunk size that is indexed and retrieved (needed for proper relative positions as well as causal chunked cross attention)
    max_seq_len = 2048,                      # max sequence length
    enc_dim = 896,                           # encoder model dim
    enc_depth = 2,                           # encoder depth
    dec_dim = 796,                           # decoder model dim
    dec_depth = 12,                          # decoder depth
    dec_cross_attn_layers = (3, 6, 9, 12),   # decoder cross attention layers (with causal chunk cross attention)
    heads = 8,                               # attention heads
    dim_head = 64,                           # dimension per head
    dec_attn_dropout = 0.25,                 # decoder attention dropout
    dec_ff_dropout = 0.25,                   # decoder feedforward dropout
    use_deepnet = True                       # turn on post-normalization with DeepNet residual scaling and initialization, for scaling to 1000 layers
).cuda()


seq = torch.randint(0, 20000, (2, 2048 + 1)).cuda()      # plus one since it is split into input and labels for training
retrieved = torch.randint(0, 20000, (2, 32, 2, 128)).cuda() # retrieved tokens - (batch, num chunks, num retrieved neighbors, retrieved chunk with continuation)

@torchdynamo.optimize("inductor")
def train():
    loss = retro(seq, retrieved, return_loss = True)
    loss.backward()

train()
```

## Logs

```
(dynamo) ubuntu@ip-172-31-31-152:~/tests$ python retro.py 
torchdynamo.symbolic_convert: [WARNING] Graph break: missing: BUILD_MAP_UNPACK_WITH_CALL from user code at   File "retro.py", line 26, in train
    loss = retro(seq, retrieved, return_loss = True)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/retro_pytorch/retro_pytorch.py", line 618, in forward
    pos_emb = rearrange(pos_emb, 'n d -> 1 n d')
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/einops/einops.py", line 487, in rearrange
    return reduce(tensor, pattern, reduction='rearrange', **axes_lengths)

torchdynamo.symbolic_convert: [WARNING] Graph break: call_method ConstDictVariable() __contains__ [GetAttrVariable(UserDefinedClassVariable(), framework_name)] {} from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/einops/einops.py", line 231, in _apply_recipe
    backend = get_backend(tensor)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/einops/_backends.py", line 42, in get_backend
    if BackendSubclass.framework_name not in _backends:

torchdynamo.symbolic_convert: [WARNING] Graph break: call_function UserDefinedObjectVariable(_lru_cache_wrapper) [UserDefinedObjectVariable(TransformRecipe), ShapeVariable()] {} from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/einops/einops.py", line 233, in <graph break in _apply_recipe>
    _reconstruct_from_shape(recipe, backend.shape(tensor))

torchinductor.compile_fx: [WARNING] skipping cudagraphs due to complex input striding
torchinductor.compile_fx: [WARNING] skipping cudagraphs due to complex input striding
torchinductor.lowering: [WARNING] using triton random, expect difference from eager
torchinductor.lowering: [WARNING] using triton random, expect difference from eager
torchinductor.compile_fx: [WARNING] skipping cudagraphs due to complex input striding
torchinductor.compile_fx: [WARNING] skipping cudagraphs due to complex input striding
torchinductor.compile_fx: [WARNING] skipping cudagraphs due to complex input striding
torchinductor.compile_fx: [WARNING] skipping cudagraphs due to complex input striding
torchdynamo.symbolic_convert: [WARNING] Graph break: missing: BUILD_MAP_UNPACK_WITH_CALL from user code at   File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/retro_pytorch/retro_pytorch.py", line 87, in forward
    out = self.fn(x, *args, **kwargs) + residual
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/retro_pytorch/retro_pytorch.py", line 276, in forward
    k_pos_emb = repeat(k_pos_emb, 'b h n d -> b h (r n) d', r = num_retrieved)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/einops/einops.py", line 537, in repeat
    return reduce(tensor, pattern, reduction='repeat', **axes_lengths)

torchinductor.compile_fx: [WARNING] skipping cudagraphs due to complex input striding
torchinductor.compile_fx: [WARNING] skipping cudagraphs due to complex input striding
torchinductor.compile_fx: [WARNING] skipping cudagraphs due to complex input striding
torchinductor.compile_fx: [WARNING] skipping cudagraphs due to complex input striding
torchinductor.lowering: [WARNING] using triton random, expect difference from eager
torchinductor.lowering: [WARNING] using triton random, expect difference from eager
torchinductor.lowering: [WARNING] using triton random, expect difference from eager
torchinductor.lowering: [WARNING] using triton random, expect difference from eager
torchinductor.lowering: [WARNING] using triton random, expect difference from eager
torchinductor.lowering: [WARNING] using triton random, expect difference from eager
torchinductor.lowering: [WARNING] using triton random, expect difference from eager
Traceback (most recent call last):
  File "retro.py", line 29, in <module>
    train()
  File "/home/ubuntu/torchdynamo/torchdynamo/eval_frame.py", line 166, in _fn
    return fn(*args, **kwargs)
  File "retro.py", line 24, in train
    @torchdynamo.optimize("inductor")
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/retro_pytorch/retro_pytorch.py", line 564, in forward
    def forward(
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/retro_pytorch/retro_pytorch.py", line 618, in <graph break in forward>
    pos_emb = rearrange(pos_emb, 'n d -> 1 n d')
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/retro_pytorch/retro_pytorch.py", line 627, in <graph break in forward>
    encoder_retrieved_mask = rearrange(mask, 'b k r n -> (b k r) n')
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/retro_pytorch/retro_pytorch.py", line 404, in forward
    def forward(self, x, *, encoder = None, encoder_retrieved_mask = None, context_mask = None, retrieved = None):
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/retro_pytorch/retro_pytorch.py", line 406, in <graph break in forward>
    self_attn_pos_emb = self.rotary_pos_emb(seq_len, device = device)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/retro_pytorch/retro_pytorch.py", line 418, in <graph break in forward>
    cross_attn_q_pos_emb = self.rotary_pos_emb(self.chunk_size, device = device, offset = self.chunk_size - 1)  # need to add extra chunk size, since it will be shifted
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/retro_pytorch/retro_pytorch.py", line 430, in <graph break in forward>
    x = attn(x, pos_emb = self_attn_pos_emb)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/retro_pytorch/retro_pytorch.py", line 85, in forward
    def forward(self, x, *args, **kwargs):
  File "/home/ubuntu/torchdynamo/torchdynamo/eval_frame.py", line 166, in _fn
    return fn(*args, **kwargs)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/functorch/_src/aot_autograd.py", line 904, in forward
    return compiled_f(
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/functorch/_src/aot_autograd.py", line 895, in new_func
    return compiled_fn(args)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/functorch/_src/aot_autograd.py", line 228, in g
    return f(*args)
  File "/home/ubuntu/torchdynamo/torchdynamo/eval_frame.py", line 166, in _fn
    return fn(*args, **kwargs)
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/functorch/_src/aot_autograd.py", line 354, in forward
    fw_outs = call_func_with_args(
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/functorch/_src/aot_autograd.py", line 250, in call_func_with_args
    out = normalize_as_list(f(args))
  File "/opt/conda/envs/dynamo/lib/python3.8/site-packages/functorch/_src/aot_autograd.py", line 228, in g
    return f(*args)
  File "/home/ubuntu/torchdynamo/torchinductor/compile_fx.py", line 125, in run
    compiled_fn = cudagraphify_impl(model, new_inputs, static_input_idxs)
  File "/home/ubuntu/torchdynamo/torchinductor/compile_fx.py", line 150, in cudagraphify_impl
    static_inputs = [
  File "/home/ubuntu/torchdynamo/torchinductor/compile_fx.py", line 151, in <listcomp>
    static_input(x) if idx not in static_input_idxs else inputs[idx]
  File "/home/ubuntu/torchdynamo/torchinductor/compile_fx.py", line 143, in static_input
    buffer = torch.zeros(needed_size, dtype=x.dtype, device=x.device)
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 14.00 MiB (GPU 0; 15.78 GiB total capacity; 13.71 GiB already allocated; 11.94 MiB free; 14.77 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
```

cc @ezyang @wconstab @bdhirsh @anijain2305 @zou3519 @soumith @ngimel