Repository: pytorch
Issue: Torch Compile Dynamic fails on sample on diffusers VAE
Labels: triaged,oncall: pt2,module: dynamic shapes

### 🐛 Describe the bug

Torch compile fails on dIffusers VAE code when calling .sample() in the randn_tensor helper function. Disabling dynamic compilation mode solves the issue, but slows other things down. I finally got the minifier to output a reproducer. It occurs on this line: https://github.com/mosaicml/diffusion/blob/cc26b39ad18ed0956a69f035096588507980b572/diffusion/models/stable_diffusion.py#L183

### Error logs

```
Error executing job with overrides: []
Traceback (most recent call last):
  File "/mnt/workdisk/aaron/diffusion-main/run.py", line 22, in main
    return train(config)
  File "/mnt/workdisk/aaron/diffusion-main/diffusion/train.py", line 139, in train
    return eval_and_then_train()
  File "/mnt/workdisk/aaron/diffusion-main/diffusion/train.py", line 137, in eval_and_then_train
    trainer.fit()
  File "/usr/lib/python3/dist-packages/composer/trainer/trainer.py", line 1876, in fit
    self._train_loop()
  File "/usr/lib/python3/dist-packages/composer/trainer/trainer.py", line 2051, in _train_loop
    total_loss_dict = self._train_batch(use_grad_scaling)
  File "/usr/lib/python3/dist-packages/composer/trainer/trainer.py", line 2239, in _train_batch
    self._train_microbatches(microbatches, total_loss_dict)
  File "/usr/lib/python3/dist-packages/composer/trainer/trainer.py", line 2339, in _train_microbatches
    microbatch_loss_dict = self._train_microbatch(use_grad_scaling, current_batch_size, is_final_microbatch)
  File "/usr/lib/python3/dist-packages/composer/trainer/trainer.py", line 2403, in _train_microbatch
    self.state.outputs = self.state.model(self.state.batch)
  File "/usr/lib/python3/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/torch/_dynamo/eval_frame.py", line 328, in _fn
    return fn(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/torch/_dynamo/eval_frame.py", line 490, in catch_errors
    return callback(frame, cache_entry, hooks, frame_state)
  File "/usr/lib/python3/dist-packages/torch/_dynamo/convert_frame.py", line 641, in _convert_frame
    result = inner_convert(frame, cache_size, hooks, frame_state)
  File "/usr/lib/python3/dist-packages/torch/_dynamo/convert_frame.py", line 133, in _fn
    return fn(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/torch/_dynamo/convert_frame.py", line 389, in _convert_frame_assert
    return _compile(
  File "/usr/lib/python3/dist-packages/torch/_dynamo/convert_frame.py", line 569, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/usr/lib/python3/dist-packages/torch/_dynamo/utils.py", line 189, in time_wrapper
    r = func(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/torch/_dynamo/convert_frame.py", line 491, in compile_inner
    out_code = transform_code_object(code, transform)
  File "/usr/lib/python3/dist-packages/torch/_dynamo/bytecode_transformation.py", line 1028, in transform_code_object
    transformations(instructions, code_options)
  File "/usr/lib/python3/dist-packages/torch/_dynamo/convert_frame.py", line 458, in transform
    tracer.run()
  File "/usr/lib/python3/dist-packages/torch/_dynamo/symbolic_convert.py", line 2074, in run
    super().run()
  File "/usr/lib/python3/dist-packages/torch/_dynamo/symbolic_convert.py", line 724, in run
    and self.step()
  File "/usr/lib/python3/dist-packages/torch/_dynamo/symbolic_convert.py", line 688, in step
    getattr(self, inst.opname)(inst)
  File "/usr/lib/python3/dist-packages/torch/_dynamo/symbolic_convert.py", line 439, in wrapper
    self.output.compile_subgraph(self, reason=reason)
  File "/usr/lib/python3/dist-packages/torch/_dynamo/output_graph.py", line 857, in compile_subgraph
    self.compile_and_call_fx_graph(tx, pass2.graph_output_vars(), root)
  File "/usr/lib/python3.10/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/usr/lib/python3/dist-packages/torch/_dynamo/output_graph.py", line 957, in compile_and_call_fx_graph
    compiled_fn = self.call_user_compiler(gm)
  File "/usr/lib/python3/dist-packages/torch/_dynamo/utils.py", line 189, in time_wrapper
    r = func(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/torch/_dynamo/output_graph.py", line 1024, in call_user_compiler
    raise BackendCompilerFailed(self.compiler_fn, e).with_traceback(
  File "/usr/lib/python3/dist-packages/torch/_dynamo/output_graph.py", line 1009, in call_user_compiler
    compiled_fn = compiler_fn(gm, self.example_inputs())
  File "/usr/lib/python3/dist-packages/torch/_dynamo/repro/after_dynamo.py", line 96, in debug_wrapper
    run_fwd_maybe_bwd(compiled_gm, example_inputs)
  File "/usr/lib/python3/dist-packages/torch/_dynamo/debug_utils.py", line 327, in run_fwd_maybe_bwd
    out = gm(args)
  File "/usr/lib/python3/dist-packages/torch/_functorch/aot_autograd.py", line 1482, in g
    return f(*args)
  File "/usr/lib/python3/dist-packages/torch/_dynamo/eval_frame.py", line 328, in _fn
    return fn(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/torch/_dynamo/external_utils.py", line 17, in inner
    return fn(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/torch/_functorch/aot_autograd.py", line 3905, in forward
    return compiled_fn(full_args)
  File "/usr/lib/python3/dist-packages/torch/_functorch/aot_autograd.py", line 1482, in g
    return f(*args)
  File "/usr/lib/python3/dist-packages/torch/_functorch/aot_autograd.py", line 2533, in runtime_wrapper
    all_outs = call_func_with_args(
  File "/usr/lib/python3/dist-packages/torch/_functorch/aot_autograd.py", line 1506, in call_func_with_args
    out = normalize_as_list(f(args))
  File "/usr/lib/python3/dist-packages/torch/_functorch/aot_autograd.py", line 1594, in rng_functionalization_wrapper
    return compiled_fw(args)
  File "/usr/lib/python3/dist-packages/torch/_inductor/codecache.py", line 374, in __call__
    return self.get_current_callable()(inputs)
  File "/usr/lib/python3/dist-packages/torch/_inductor/compile_fx.py", line 628, in run
    return model(new_inputs)
  File "/usr/lib/python3/dist-packages/torch/_inductor/codecache.py", line 401, in _run_from_cache
    return compiled_graph.compiled_artifact(inputs)
  File "/tmp/torchinductor_root/ic/cickzxbroxiht2qr44bogc2wzvnihmoqpfjeapb4cfbfqplh77s3.py", line 65, in call
    buf0 = empty_strided((s0, 3, s2, s2), (3*(s2*s2), s2*s2, s2, 1), device='cuda', dtype=torch.float16)
torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
RuntimeError: aten/src/ATen/RegisterCUDA.cpp:7014: SymIntArrayRef expected to contain only concrete integers


from user code:
   File "/mnt/workdisk/aaron/diffusion-main/diffusion/models/stable_diffusion.py", line 183, in <resume in forward>
    latents = self.vae.encode(inputs.half()).latent_dist.sample()
  File "/usr/lib/python3/dist-packages/diffusers/models/vae.py", line 661, in sample
    sample = randn_tensor(
  File "/usr/lib/python3/dist-packages/diffusers/utils/torch_utils.py", line 79, in randn_tensor
    latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype, layout=layout).to(device)


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True
```

### Minified repro

```
from math import inf
import torch
from torch import tensor, device
import torch.fx as fx
import torch._dynamo
from torch._dynamo.testing import rand_strided
from torch._dynamo.debug_utils import run_fwd_maybe_bwd

import torch._dynamo.config
import torch._inductor.config
import torch._functorch.config
torch._dynamo.config.assume_static_by_default = False







from torch.nn import *
class Repro(torch.nn.Module):
    def __init__(self):
        super().__init__()



    def forward(self, s0 : torch.SymInt, s2 : torch.SymInt, L_batch_image_ : torch.Tensor, s3 : torch.SymInt, L_batch_captions_ : torch.Tensor):
        l_batch_image_ = L_batch_image_
        l_batch_captions_ = L_batch_captions_
        size = l_batch_captions_.size()
        getitem_1 = size[1];  size = None
        view = l_batch_captions_.view(-1, getitem_1);  l_batch_captions_ = getitem_1 = None
        _enter_autocast = torch.amp.autocast_mode._enter_autocast('cuda', torch.float16, False, True)
        half = l_batch_image_.half();  l_batch_image_ = None
        _exit_autocast = torch.amp.autocast_mode._exit_autocast(_enter_autocast);  _enter_autocast = None
        return (half, view)


mod = Repro()

def load_args(reader):
    reader.symint(64)  # s0
    reader.symint(256)  # s2
    buf0 = reader.storage('ea62cd3108ac82f857285316ed9a6be4f3040d2a', 402653184, device=device(type='cuda', index=1))
    reader.tensor(buf0, (64, 3, 256, 256), is_leaf=True)  # L_batch_image_
    reader.symint(77)  # s3
    buf1 = reader.storage('9d850de4f79e8252556788997d9e61eb08e160ce', 315392, device=device(type='cuda', index=1), dtype_hint=torch.int64)
    reader.tensor(buf1, (64, 77), dtype=torch.int64, is_leaf=True)  # L_batch_captions_
load_args._version = 0

if __name__ == '__main__':
    from torch._dynamo.repro.after_dynamo import run_repro
    run_repro(mod, load_args, accuracy=False, command='minify',
        save_dir='/mnt/workdisk/aaron/diffusion-main/torch_compile_debug/run_2023_10_16_19_54_51_021280-pid_42543/minifier/checkpoints', autocast=True, backend='inductor')
```

### Versions

```
PyTorch version: 2.1.0+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: version 3.26.3
Libc version: glibc-2.31

Python version: 3.10.13 (main, Aug 25 2023, 13:20:03) [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.4.0-137-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 12.1.105
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA A100-SXM4-40GB
GPU 1: NVIDIA A100-SXM4-40GB
GPU 2: NVIDIA A100-SXM4-40GB
GPU 3: NVIDIA A100-SXM4-40GB

Nvidia driver version: 515.48.07
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.0
/usr/lib/x86_64-linux-gnu/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:                   48 bits physical, 48 bits virtual
CPU(s):                          64
On-line CPU(s) list:             0-63
Thread(s) per core:              1
Core(s) per socket:              32
Socket(s):                       2
NUMA node(s):                    8
Vendor ID:                       AuthenticAMD
CPU family:                      25
Model:                           1
Model name:                      AMD EPYC 7513 32-Core Processor
Stepping:                        1
Frequency boost:                 enabled
CPU MHz:                         1497.016
CPU max MHz:                     2600.0000
CPU min MHz:                     1500.0000
BogoMIPS:                        5200.13
Virtualization:                  AMD-V
L1d cache:                       2 MiB
L1i cache:                       2 MiB
L2 cache:                        32 MiB
L3 cache:                        256 MiB
NUMA node0 CPU(s):               0-7
NUMA node1 CPU(s):               8-15
NUMA node2 CPU(s):               16-23
NUMA node3 CPU(s):               24-31
NUMA node4 CPU(s):               32-39
NUMA node5 CPU(s):               40-47
NUMA node6 CPU(s):               48-55
NUMA node7 CPU(s):               56-63
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Not affected
Vulnerability Retbleed:          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; Retpolines, IBPB conditional, IBRS_FW, STIBP disabled, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected
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 pcid sse4_1 sse4_2 x2apic 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 invpcid_single hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 invpcid 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 arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca

Versions of relevant libraries:
[pip3] numpy==1.26.0
[pip3] pytorch-ranger==0.1.1
[pip3] torch==2.1.0+cu121
[pip3] torch-optimizer==0.3.0
[pip3] torchmetrics==1.0.3
[pip3] torchvision==0.16.0+cu121
[pip3] triton==2.1.0
```

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305