Repository: pytorch
Issue: Alternatives to compiling a for loop
Labels: feature,triaged,months,oncall: pt2,module: dynamo,module: higher order operators,module: pt2-dispatcher

### 🐛 Describe the bug

So I've been trying to compile some auto-regressive transformer models and they consist of for loops within the model architecture. I'm posting a simplified version of the architecture here

```
import torch
import torch.nn as nn
import math

def init_biased_mask(n_head, max_seq_len, period):
    def get_slopes(n):
        def get_slopes_power_of_2(n):
            start = (2**(-2**-(math.log2(n)-3)))
            ratio = start
            return [start*ratio**i for i in range(n)]
        if math.log2(n).is_integer():
            return get_slopes_power_of_2(n)                   
        else:                                                 
            closest_power_of_2 = 2**math.floor(math.log2(n)) 
            return get_slopes_power_of_2(closest_power_of_2) + get_slopes(2*closest_power_of_2)[0::2][:n-closest_power_of_2]
    
    slopes = torch.Tensor(get_slopes(n_head))
    bias = torch.arange(start=0, end=max_seq_len, step=period).unsqueeze(1).repeat(1,period).view(-1)//(period)
    bias = - torch.flip(bias,dims=[0])
    alibi = torch.zeros(max_seq_len, max_seq_len)
    for i in range(max_seq_len):
        alibi[i, :i+1] = bias[-(i+1):]
    alibi = slopes.unsqueeze(1).unsqueeze(1) * alibi.unsqueeze(0)
    mask = (torch.triu(torch.ones(max_seq_len, max_seq_len)) == 1).transpose(0, 1)
    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
    mask = mask.unsqueeze(0) + alibi
    return mask

def enc_dec_mask(device, T, S):
    mask = torch.ones(T, S)
    for i in range(T):
        mask[i, i] = 0
    return (mask==1).to(device=device)

class DummyTransformer(nn.Module):
    def __init__(self, device):
        super(DummyTransformer, self).__init__()
        self.device = device
        self.feature_dim = 64
        self.nhead = 4
        self.max_len = 600
        self.biased_mask = init_biased_mask(n_head = 4, max_seq_len = 600, period=30)
        self.style = torch.randn((1, 64)).to(torch.float32).to(self.device)
        decoder_layer = nn.TransformerDecoderLayer(d_model=self.feature_dim, nhead=self.nhead, dim_feedforward=2*self.feature_dim, batch_first=True)        
        self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=1)
        self.decoder = nn.Linear(64, 5023*3, bias=False)
        self.loss = nn.MSELoss()

    def forward(self, input, hidden_states):
        frame_num = input.shape[1]
        for i in range(frame_num):
            if i==0:
                vertice_emb = torch.zeros((1,64)).unsqueeze(1) # (1,1,feature_dim)
                vertice_input = vertice_emb
            else:
                vertice_input = vertice_emb
            tgt_mask = self.biased_mask[:, :vertice_input.shape[1], :vertice_input.shape[1]].clone().detach().to(device=self.device)
            memory_mask = enc_dec_mask(self.device, vertice_input.shape[1], hidden_states.shape[1])
            vertice_out = self.transformer_decoder(vertice_input, hidden_states, tgt_mask=tgt_mask, memory_mask=memory_mask)
            new_output = vertice_out[:,-1,:].unsqueeze(1)
            vertice_emb = torch.cat((vertice_emb, new_output), 1)
        vertice_out = self.decoder(vertice_out)
        loss = self.loss(vertice_out, input)
        return vertice_out, loss

# frameNums = [500, 200, 500, 300]
frameNums = [5, 5, 5, 5, 2, 2, 2, 2, 5, 5, 5, 5, 3, 3, 3, 3]

model = DummyTransformer('cpu')
model.to('cpu')
for i in frameNums:
    hidden_states = torch.randn((1, i, 64)).to(torch.float32).to('cpu')
    input = torch.randn((1, i, 5023*3))
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    output, loss = model(input, hidden_states)
    end.record()
    torch.cuda.synchronize()
    print("eager", start.elapsed_time(end) / 1000)
print("--------------------------")
opt_model = torch.compile(model)
for i in frameNums:
    hidden_states = torch.randn((1, i, 64)).to(torch.float32).to('cpu')
    input = torch.randn((1, i, 5023*3))
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    output, loss = opt_model(input, hidden_states)
    end.record()
    torch.cuda.synchronize()
    print("compiled", start.elapsed_time(end) / 1000)
```

An interesting observation I had while running this snip was the timings I got for compiled models
```
eager 0.01397555160522461
eager 0.01042841625213623
eager 0.013084671974182128
eager 0.009634816169738769
eager 0.0038338561058044434
eager 0.0038000640869140625
eager 0.0023736319541931153
eager 0.003575808048248291
eager 0.00902143955230713
eager 0.009564160346984863
eager 0.005406720161437988
eager 0.009522175788879395
eager 0.007098368167877197
eager 0.005195775985717774
eager 0.005480447769165039
eager 0.00552345609664917
--------------------------
compiled 28.39280078125
compiled 0.005262335777282715
compiled 0.005636096000671387
compiled 0.0052111358642578124
compiled 10.9308525390625
compiled 7.168000098317862e-06
compiled 0.0016711679697036742
compiled 0.0015872000455856323
compiled 0.003831808090209961
compiled 0.0035819520950317383
compiled 0.005676032066345215
compiled 0.002215935945510864
compiled 25.4304765625
compiled 0.0032000000476837156
compiled 0.0028743679523468017
compiled 0.0027555840015411376
```
`torch.compile()` is caching the compiles for previously encountered `frameNums` and using them whenever it encounters them again

From my limited understanding of the issues https://github.com/pytorch/pytorch/issues/111441 and https://github.com/pytorch/pytorch/issues/105279#issuecomment-1636949119 `torch.compile()` unrolls the for loop. 
My issue here is that the dataloader serves different input sizes (temporally) so the number of iterations change for every sample. 
Every time a new input size is encountered it compiled the model again (which is infeasible). 
One solution was to cap the data being served (during training) to the smallest sequence length but that doesn't solve the issue during inference since I don't know what would be the input size. 

Is there a workaround I can use to get rid of loops in such models (or any model which have loops dependent on input size) ? 

### Versions

```
Collecting environment information...
PyTorch version: 2.1.0
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux 11 (bullseye) (x86_64)
GCC version: (Debian 10.2.1-6) 10.2.1 20210110
Clang version: Could not collect
CMake version: version 3.18.4
Libc version: glibc-2.31

Python version: 3.8.18 | packaged by conda-forge | (default, Oct 10 2023, 15:44:36)  [GCC 12.3.0] (64-bit runtime)
Python platform: Linux-5.10.0-26-cloud-amd64-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA L4
Nvidia driver version: 525.105.17
cuDNN version: Probably one of the following:
/usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudnn.so.8.9.0
/usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.9.0
/usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.9.0
/usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.9.0
/usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.9.0
/usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.9.0
/usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.9.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Byte Order:                         Little Endian
Address sizes:                      46 bits physical, 48 bits virtual
CPU(s):                             16
On-line CPU(s) list:                0-15
Thread(s) per core:                 2
Core(s) per socket:                 8
Socket(s):                          1
NUMA node(s):                       1
Vendor ID:                          GenuineIntel
CPU family:                         6
Model:                              85
Model name:                         Intel(R) Xeon(R) CPU @ 2.20GHz
Stepping:                           7
CPU MHz:                            2200.244
BogoMIPS:                           4400.48
Hypervisor vendor:                  KVM
Virtualization type:                full
L1d cache:                          256 KiB
L1i cache:                          256 KiB
L2 cache:                           8 MiB
L3 cache:                           38.5 MiB
NUMA node0 CPU(s):                  0-15
Vulnerability Gather data sampling: Unknown: Dependent on hypervisor status
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Mitigation; Clear CPU buffers; SMT Host state unknown
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Retbleed:             Mitigation; Enhanced IBRS
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Mitigation; Clear CPU buffers; SMT Host state unknown
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves arat avx512_vnni md_clear arch_capabilities

Versions of relevant libraries:
[pip3] numpy==1.21.6
[pip3] pytorch-lightning==1.6.5
[pip3] pytorch-triton==2.1.0+6e4932cda8
[pip3] pytorch3d==0.7.5
[pip3] torch==2.1.0
[pip3] torchaudio==2.1.0
[pip3] torchmetrics==0.11.4
[pip3] torchvision==0.13.1a0
[pip3] triton==2.1.0
[conda] blas                      1.0                         mkl    conda-forge
[conda] cudatoolkit               10.2.89             hdec6ad0_12    conda-forge
[conda] libjpeg-turbo             2.0.0                h9bf148f_0    pytorch
[conda] mkl                       2023.1.0         h213fc3f_46344  
[conda] numpy                     1.21.6                   pypi_0    pypi
[conda] pytorch                   2.1.0           py3.8_cuda11.8_cudnn8.7.0_0    pytorch
[conda] pytorch-cuda              11.8                 h7e8668a_5    pytorch-nightly
[conda] pytorch-lightning         1.6.5                    pypi_0    pypi
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] pytorch-triton            2.1.0+6e4932cda8          pypi_0    pypi
[conda] pytorch3d                 0.7.5           py38_cu118_pyt210    pytorch3d
[conda] torchaudio                2.1.0                py38_cu118    pytorch
[conda] torchmetrics              0.11.4                   pypi_0    pypi
[conda] torchtriton               2.1.0                      py38    pytorch
[conda] torchvision               0.13.1          cpu_py38h164cc8f_0  
```

cc @ezyang @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames @rec @zou3519 @ydwu4 @bdhirsh @msaroufim @wconstab @anijain2305 @aakhundov