[Dynamo][Compile]Torch compile with dynamic shapes not working

[bohnstingl]

🐛 Describe the bug

My networks rely on varying shapes during training as well as during inference. Thus, I tried to use torch.compile(... dynamic=True) as well as the torch._dynamo.optimize(..., dynamic=True) feature. However, unfortunately I can't get it to work properly and the function gets recompiled always as soon as the input shape changes. I tried it with the latest nightly build verison 2.1.0.dev20230712 as well as with the current main branch from Friday Jul 14th 2023. I printed the guards that trigger the recompilation and the culprit guard is precisely what the dynamic feature targets:
GuardFail(reason="tensor 'L['input']' size mismatch at index 0. expected 146, actual 149", orig_code=<code object RNNScript at 0x2758840 ...)

I also tried to add a warm-up phase, where the input length is varied, but whenever the shape of the input changes, the function gets recompiled.

@ezyang could you maybe please advice here?

Error logs

GuardFail(reason="tensor 'L['input']' size mismatch at index 0. expected 146, actual 149", orig_code=<code object RNNScript at 0x2758840 ...)

Minified repro

from typing import List, Tuple, Optional, overload, Union, cast
import torch
import numpy as np
import time
import torch.optim as optim
from torch.nn.parameter import Parameter

def RNNScript(
    input,
    param1,
    param2,
    ):

    state1 = torch.zeros(32, 340, dtype=input.dtype, device=input.device)
    
    outs = []

    Wx = input @ param1
    Wx_inp, Wx_rec = torch.tensor_split(Wx, 2, 2)
    for wt_inp, wt_rec in zip(Wx_inp, Wx_rec):
        rec_mul_inp, rec_mul_rec = torch.tensor_split(state1 @ param2, 2, 1)
        input_prev = (wt_inp + rec_mul_inp)
        output_gate = (wt_rec + rec_mul_rec)

        state1 = input_prev * torch.sigmoid(output_gate)
        outs.append(state1)
    
    outs = torch.stack(outs)

    return outs, (outs)

if __name__ == "__main__":

    input_size = 140
    hidden_size = 340
    num_layers = 1
    num_timesteps = 111
    batch_size = 32

    bi_dir = True 
    rnnt_input = False
    num_threads = -1
    use_gpu = True
    load_weights = False

    forward_times = []
    backward_times = []

    if use_gpu:
        device = torch.device('cuda:0')
    else:
        device = None

    parameters = []

    w_ih = torch.empty((input_size, hidden_size), device=device)
    w_io = torch.empty((input_size, hidden_size), device=device)
    w_i_comb = Parameter(torch.cat([w_ih,w_io],1))
    parameters.append(w_i_comb)

    w_hh = torch.empty((hidden_size, hidden_size), device=device)
    w_ho = torch.empty((hidden_size, hidden_size), device=device)
    w_h_comb = Parameter(torch.cat([w_hh,w_ho],1))
    parameters.append(w_h_comb)
    
    def count_kernels(guard):
        print("[pt2_compile] guard failed: ", guard)

    rnnscript = torch.compile(RNNScript, mode='reduce-overhead', dynamic=True, fullgraph=True)
    #backend = torch._TorchCompileInductorWrapper('reduce-overhead', None, True)
    #rnnscript = torch._dynamo.optimize(backend=backend, nopython=True, dynamic=True, guard_fail_fn=count_kernels)(RNNScript)
    #rnnscript = RNNScript
    snu = lambda x: rnnscript(x, w_i_comb, w_h_comb)

    optimizer = optim.SGD(parameters, 0.1)

    inp = torch.rand((num_timesteps, batch_size, input_size))

    if use_gpu:
        inp = inp.cuda()

    optimizer.zero_grad()
    for execution in range(5):
        start_forward = time.time_ns()
        t_rnd = np.random.randint(0, 200)
        inp = torch.rand((t_rnd, batch_size, input_size))
        if use_gpu:
            inp = inp.cuda()
        out, state = snu(inp)

        if use_gpu:
            torch.cuda.synchronize()
        stop_forward = time.time_ns()
        forward_times.append((stop_forward - start_forward) / (10 ** 9))

        loss = 1. - torch.sum(out)

        start_time_backward = time.time_ns()
        #loss.backward()
        if use_gpu:
            torch.cuda.synchronize()
        stop_time_backward = time.time_ns()
        backward_times.append((stop_time_backward - start_time_backward) / (10 ** 9))

    print('================================================================')
    print('Model with sSNU-os:')
    print('# Layers: ' + str(num_layers))
    print('# Units per layer: ' + str(hidden_size))
    print('Bidirectional: ' + str(bi_dir))
    print('Load weights: '  + str(load_weights))
    print('RNN-T input: ' + str(rnnt_input))
    print('# CPU threads: ' + str(num_threads))
    print('GPU support: ' + str(use_gpu))
    print('----------------------------------------------------------------')
    print('Timing summary')
    print('Time of forward computation: {:.4f} +- {:.4f} s'.format(np.mean(np.array(forward_times)), np.std(np.array(forward_times))))
    print('Time of backward computation: {:.4f} +- {:.4f} s'.format(np.mean(np.array(backward_times)), np.std(np.array(backward_times))))

Versions

Collecting environment information...
PyTorch version: 2.1.0a0+gitfb376f8
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Red Hat Enterprise Linux release 8.7 (Ootpa) (x86_64)
GCC version: (GCC) 8.5.0 20210514 (Red Hat 8.5.0-16)
Clang version: Could not collect
CMake version: version 3.26.4
Libc version: glibc-2.28

Python version: 3.11.3 (main, May 15 2023, 15:45:52) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-4.18.0-425.10.1.el8_7.x86_64-x86_64-with-glibc2.28
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA A100-SXM4-40GB
Nvidia driver version: 525.60.13
cuDNN version: Could not collect
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
CPU(s): 128
On-line CPU(s) list: 0-127
Thread(s) per core: 1
Core(s) per socket: 64
Socket(s): 2
NUMA node(s): 2
Vendor ID: AuthenticAMD
CPU family: 23
Model: 49
Model name: AMD EPYC 7742 64-Core Processor
Stepping: 0
CPU MHz: 3361.974
CPU max MHz: 2250.0000
CPU min MHz: 1500.0000
BogoMIPS: 4499.91
Virtualization: AMD-V
L1d cache: 32K
L1i cache: 32K
L2 cache: 512K
L3 cache: 16384K
NUMA node0 CPU(s): 0-63
NUMA node1 CPU(s): 64-127
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr wbnoinvd amd_ppin arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif v_spec_ctrl umip rdpid overflow_recov succor smca sme sev sev_es

Versions of relevant libraries:
[pip3] numpy==1.24.3
[pip3] pytorch-triton==2.1.0+3c400e7818
[pip3] torch==2.1.0a0+gitfb376f8
[pip3] torchaudio==2.1.0a0+cf53a48
[conda] blas 1.0 mkl
[conda] cudatoolkit 11.8.0 h6a678d5_0
[conda] magma-cuda118 2.6.1 1 pytorch
[conda] mkl 2023.1.0 h6d00ec8_46342
[conda] mkl-include 2023.1.0 h06a4308_46342
[conda] mkl-service 2.4.0 py311h5eee18b_1
[conda] mkl_fft 1.3.6 py311ha02d727_1
[conda] mkl_random 1.2.2 py311ha02d727_1
[conda] numpy 1.24.3 py311h08b1b3b_1
[conda] numpy-base 1.24.3 py311hf175353_1
[conda] pytorch-triton 2.1.0+3c400e7818 pypi_0 pypi
[conda] torch 2.1.0a0+gitfb376f8 pypi_0 pypi
[conda] torchaudio 2.1.0a0+cf53a48 dev_0

cc @ezyang @msaroufim @bdhirsh @anijain2305 @zou3519 @chauhang @wconstab

[ezyang]

When I run your example with TORCH_LOGS=dynamic I get

[2023-07-15 18:09:58,188] torch.fx.experimental.symbolic_shapes: [INFO] 0.0: create_env
[2023-07-15 18:09:58,480] torch.fx.experimental.symbolic_shapes: [INFO] 0.0: create_symbol s0 = 14 for L['input'].size()[0]
[2023-07-15 18:09:58,517] torch.fx.experimental.symbolic_shapes: [INFO] 0.0: create_symbol s1 = 32 for L['input'].size()[1]
[2023-07-15 18:09:58,519] torch.fx.experimental.symbolic_shapes: [INFO] 0.0: create_symbol s2 = 140 for L['input'].size()[2]
[2023-07-15 18:09:58,601] torch.fx.experimental.symbolic_shapes: [INFO] 0.0: eval Eq(s2, 140) [guard added] at rn.py:18 in RNNScript (_meta_registrations.py:1650 in meta_mm)
[2023-07-15 18:09:58,626] torch.fx.experimental.symbolic_shapes: [INFO] 0.0: eval Eq(s0, 14) [guard added] at rn.py:20 in RNNScript (_dynamo/variables/tensor.py:651 in evaluate_expr)
[2023-07-15 18:09:58,687] torch.fx.experimental.symbolic_shapes: [INFO] 0.0: eval Eq(s1, 32) [guard added] at rn.py:22 in RNNScript (_subclasses/fake_tensor.py:674 in infer_size)

There are few things going on here.

The specialization at line 18 is due to a matrix multiply with a Parameter, which we assume is always static. Since you're varying parameter sizes, this assumption is not good for you. You can un-Parameter'ify these tensors, but we should also add a knob that overrides parameter specialization.

The specialization at line 20 is more difficult to deal with. It comes from

    for wt_inp, wt_rec in zip(Wx_inp, Wx_rec):

In other words, you are iterating over the batch dimension of Wx_inp/Wx_rec. This will necessarily specialize, because we always have to specialize loop iteration counts to do loops in Python, since our strategy is always to unroll them. To avoid specialization here, you'll have to rewrite your code without loops. Vmap can help write your code in the single batch style without an explicit loop (though cc @zou3519, I don't remember if you have to do anything extra to make it work in torch.compile.)

[bohnstingl]

Thank you so much for your very quick and detailed insights here! They are very helpful!

When you say un-Parameter'ify, you mean using Tensor(..., require_grad=True) instead of a Parameter(...)?

As you might have noticed, my code should emulate the function of a simple RNN unit. In the static case this works perfectly fine, but the dynamic case is giving me trouble. Therefore, for wt_inp, wt_rec in zip(Wx_inp, Wx_rec): iterates over the time dimension of the multiplication result (The shape of the input is time x batch x features). I would be more than happy to learn to rewrite my code in a way that it works with torch.compile(..., dynamic=True) as I think this features is just amazing.

[vadimkantorov]

@ezyang btw this might me related to the torch.scan/loop discussion. If we get one dynamic variant to work, it could be an example or even torch.scan could be implemented using just a loop (if we have tests that it compiles into sth nice). Custom RNNs is a prototypical use case for this...

[bohnstingl]

@ezyang I have modified my example using torch.Tensor instead of Parameter and the guards that are added in that case look fine to me.

However, the for-loop is something that I cannot work around.
I tried all sorts of variations for the for-loop as well as a while loops with branches, but the guards that are added are always fully constraint to the specific time dimension of the tensor. @vadimkantorov do you have an example of the torch.scan/loop where a dynamic number of iterations is supported?

[vadimkantorov]

No working example, just the discussion here: #50688

If your use case can be implemented by efficiently by the relevant functions in jax, I would say it's one more argument to have similar functions in pytorch

And if existing compile fails for this dynamic case, I would say that your usecase is an important test case

[bohnstingl]

I, see. With Jax, custom RNNs should be simple to implement using jax.lax.scan as illustrated here. Also this was done efficiently in TF, I believe. To me in general having the capability to evaluate a function on a per-timestep basis would be important. This is also similar to #103588, I think.

As I mentioned, torch.compile(), even with the dynamic mode, introduces guards that trigger a recompilation of the custom RNN in PyTorch. So the for me dynamic input lengths are not working.

[bohnstingl]

Could this pull request #104442 resolve this issue?

[bohnstingl]

@ezyang I just wanted to follow-up on the topic of the parameter specialization you mentioned in case of the matrix multiplication. Is there a plan to change this, such that the the matmul with Parameters behaves similarly as the matmul with Tensors as far as the compilation is concerned?

I don't have a good overview of the PT2 parts, but if there is something I can help with, I would be happy to

[ezyang]

Look at tensor_always_has_static_shape in torch/_dynamo/utils.py. Add a config to bypass this and always return False, None. Send a PR and add me as reviewer. Should be fairly simple. Don't forget to add a test.

[bohnstingl]

@ezyang Since parameter part of this issue is in the making, is there a way to address the second part of this issue?
I looked at torch.vmap, but I wouldn't know how to implement RNNs with it, because if one wants to use it to run a function over the time dimension, it does not allow for a carry, right? Something like scan would then be needed, which has demanded also a lot, see here and here.

Alternatively to that, would it be possible that the python for-loop is not always unrolled, but translated into a for loop inside the triton kernel? This might be similarly effective as the unrolling, but more flexible to allow dynamic shapes?

[ezyang]

If you want to do an device scan you will have to wait for @peterbell10 to get to it alas