Repository: pytorch
Issue: `torch.export.export` infers dynamic_shape as constant
Labels: oncall: pt2,module: dynamic shapes,oncall: export

### 🐛 Describe the bug

Install most recent transformers and pytorch.

```
pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu124
pip install transformers
```

<details>
<summary>Repro</summary>

```
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from torch.export import Dim


def get_example_inputs(prompt: str, tokenizer: AutoTokenizer) -> torch.tensor:
    """
    These arbitrary example inputs were observed by adding a debugger (`import pdb; pdb.set_trace()`)
    to the line corresponding this https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L1197
    inside local venv site-packages.
    """
    example_inputs = tokenizer(
        prompt,
        return_tensors="pt"
    ).to("cuda")
    seq_len = example_inputs["input_ids"].shape[1]
    example_inputs["position_ids"] = torch.arange(seq_len).unsqueeze(0).to("cuda")
    example_inputs["past_key_values"] = None
    example_inputs["inputs_embeds"] = None
    example_inputs["use_cache"] = False
    example_inputs["output_attentions"] = False
    example_inputs["output_hidden_states"] = False
    example_inputs["return_dict"] = True
    example_inputs["cache_position"] = torch.arange(seq_len).to("cuda")
    return example_inputs

@torch.no_grad()
def load_model():
    model = AutoModelForCausalLM.from_pretrained(
        "meta-llama/Llama-3.2-3B-Instruct", torch_dtype=torch.float16
    ).to("cuda")
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    return model, tokenizer

def aot_compile(path, model, **sample_kwargs):
    """
    torch.export.export + torch._inductor.aoti_compile_and_package the model,
    using given sample inputs.
    """
    seq_len_dim = Dim("seq_len", min=1, max=128)
    dynamic_shapes = {
        "input_ids": {1: seq_len_dim},
        "attention_mask": {1: seq_len_dim},
        "position_ids": {1: seq_len_dim},
        "past_key_values": None,
        "inputs_embeds": None,
        "use_cache": None,
        "output_attentions": None,
        "output_hidden_states": None,
        "return_dict": None,
        "cache_position": {1: Dim("cache_position", min=1, max=128)},
    }
    exported_program = torch.export.export(
        model.model,
        (),
        sample_kwargs,
        dynamic_shapes=dynamic_shapes
    )

    return torch._inductor.aoti_compile_and_package(
        exported_program,
        (),
        sample_kwargs,
        package_path=path,
    )

def aot_load(path):
    return torch._inductor.aoti_load_package(path)

if __name__ == "__main__":
    model, tokenizer = load_model()

    prompt = "What is a compiler?"

    inputs1 = get_example_inputs(prompt, tokenizer)
    compile_path = aot_compile('llama3.pt2', model, **inputs1)
    print(f"AoT compiled path {compile_path}")
```
</details>

<details>
<summary>Error</summary>

```
E1107 12:34:58.861135 2008591 torch/_guards.py:300] [0/0] Error while creating guard:
E1107 12:34:58.861135 2008591 torch/_guards.py:300] [0/0] Name: ''
E1107 12:34:58.861135 2008591 torch/_guards.py:300] [0/0]     Source: shape_env
E1107 12:34:58.861135 2008591 torch/_guards.py:300] [0/0]     Create Function: SHAPE_ENV
E1107 12:34:58.861135 2008591 torch/_guards.py:300] [0/0]     Guard Types: None
E1107 12:34:58.861135 2008591 torch/_guards.py:300] [0/0]     Code List: None
E1107 12:34:58.861135 2008591 torch/_guards.py:300] [0/0]     Object Weakref: None
E1107 12:34:58.861135 2008591 torch/_guards.py:300] [0/0]     Guarded Class Weakref: None
E1107 12:34:58.861135 2008591 torch/_guards.py:300] [0/0] Traceback (most recent call last):
E1107 12:34:58.861135 2008591 torch/_guards.py:300] [0/0]   File "/home/venv/lib/python3.9/site-packages/torch/_guards.py", line 298, in create
E1107 12:34:58.861135 2008591 torch/_guards.py:300] [0/0]     return self.create_fn(builder, self)
E1107 12:34:58.861135 2008591 torch/_guards.py:300] [0/0]   File "/home/venv/lib/python3.9/site-packages/torch/_dynamo/guards.py", line 1766, in SHAPE_ENV
E1107 12:34:58.861135 2008591 torch/_guards.py:300] [0/0]     code_parts, verbose_code_parts = output_graph.shape_env.produce_guards_verbose(
E1107 12:34:58.861135 2008591 torch/_guards.py:300] [0/0]   File "/home/venv/lib/python3.9/site-packages/torch/fx/experimental/symbolic_shapes.py", line 5072, in produce_guards_verbose
E1107 12:34:58.861135 2008591 torch/_guards.py:300] [0/0]     raise ConstraintViolationError(
E1107 12:34:58.861135 2008591 torch/_guards.py:300] [0/0] torch.fx.experimental.symbolic_shapes.ConstraintViolationError: Constraints violated (seq_len)! For more information, run with TORCH_LOGS="+dynamic".
E1107 12:34:58.861135 2008591 torch/_guards.py:300] [0/0]   - Not all values of seq_len = L['input_ids'].size()[1] in the specified range seq_len <= 128 are valid because seq_len was inferred to be a constant (6).
E1107 12:34:58.861135 2008591 torch/_guards.py:300] [0/0]   - Not all values of seq_len = L['attention_mask'].size()[1] in the specified range seq_len <= 128 are valid because seq_len was inferred to be a constant (6).
E1107 12:34:58.861135 2008591 torch/_guards.py:300] [0/0]   - Not all values of seq_len = L['position_ids'].size()[1] in the specified range seq_len <= 128 are valid because seq_len was inferred to be a constant (6).
E1107 12:34:58.862866 2008591 torch/_guards.py:302] [0/0] Created at:
E1107 12:34:58.862866 2008591 torch/_guards.py:302] [0/0]   File "/home/venv/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 642, in transform
E1107 12:34:58.862866 2008591 torch/_guards.py:302] [0/0]     tracer = InstructionTranslator(
E1107 12:34:58.862866 2008591 torch/_guards.py:302] [0/0]   File "/home/venv/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py", line 2807, in __init__
E1107 12:34:58.862866 2008591 torch/_guards.py:302] [0/0]     output=OutputGraph(
E1107 12:34:58.862866 2008591 torch/_guards.py:302] [0/0]   File "/home/venv/lib/python3.9/site-packages/torch/_dynamo/output_graph.py", line 319, in __init__
E1107 12:34:58.862866 2008591 torch/_guards.py:302] [0/0]     self.init_ambient_guards()
E1107 12:34:58.862866 2008591 torch/_guards.py:302] [0/0]   File "/home/venv/lib/python3.9/site-packages/torch/_dynamo/output_graph.py", line 468, in init_ambient_guards
E1107 12:34:58.862866 2008591 torch/_guards.py:302] [0/0]     self.guards.add(ShapeEnvSource().make_guard(GuardBuilder.SHAPE_ENV))
Traceback (most recent call last):
  File "/home/venv/lib/python3.9/site-packages/torch/export/_trace.py", line 660, in _export_to_torch_ir
    gm_torch_level, _ = torch._dynamo.export(
  File "/home/venv/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 1584, in inner
    raise constraint_violation_error
  File "/home/venv/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 1539, in inner
    result_traced = opt_f(*args, **kwargs)
  File "/home/venv/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/venv/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/venv/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 556, in _fn
    return fn(*args, **kwargs)
  File "/home/venv/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/venv/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/venv/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 1423, in __call__
    return self._torchdynamo_orig_callable(
  File "/home/venv/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 549, in __call__
    return _compile(
  File "/home/venv/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 977, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/home/venv/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 708, in compile_inner
    return _compile_inner(code, one_graph, hooks, transform)
  File "/home/venv/lib/python3.9/site-packages/torch/_utils_internal.py", line 95, in wrapper_function
    return function(*args, **kwargs)
  File "/home/venv/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py", line 840, in _compile_inner
    check_fn = CheckFunctionManager(
  File "/home/venv/lib/python3.9/site-packages/torch/_dynamo/guards.py", line 2183, in __init__
    guard.create(builder)
  File "/home/venv/lib/python3.9/site-packages/torch/_guards.py", line 298, in create
    return self.create_fn(builder, self)
  File "/home/venv/lib/python3.9/site-packages/torch/_dynamo/guards.py", line 1766, in SHAPE_ENV
    code_parts, verbose_code_parts = output_graph.shape_env.produce_guards_verbose(
  File "/home/venv/lib/python3.9/site-packages/torch/fx/experimental/symbolic_shapes.py", line 5072, in produce_guards_verbose
    raise ConstraintViolationError(
torch.fx.experimental.symbolic_shapes.ConstraintViolationError: Constraints violated (seq_len)! For more information, run with TORCH_LOGS="+dynamic".
  - Not all values of seq_len = L['input_ids'].size()[1] in the specified range seq_len <= 128 are valid because seq_len was inferred to be a constant (6).
  - Not all values of seq_len = L['attention_mask'].size()[1] in the specified range seq_len <= 128 are valid because seq_len was inferred to be a constant (6).
  - Not all values of seq_len = L['position_ids'].size()[1] in the specified range seq_len <= 128 are valid because seq_len was inferred to be a constant (6).

Suggested fixes:
  seq_len = 6

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/workspaces/torch-serve-amd/examples/compile_tests/aot_llama/github_issue_code.py", line 96, in <module>
    compile_path = aot_compile('llama3.pt2', model, **inputs1)
  File "/workspaces/torch-serve-amd/examples/compile_tests/aot_llama/github_issue_code.py", line 56, in aot_compile
    exported_program = torch.export.export(
  File "/home/venv/lib/python3.9/site-packages/torch/export/__init__.py", line 368, in export
    return _export(
  File "/home/venv/lib/python3.9/site-packages/torch/export/_trace.py", line 1031, in wrapper
    raise e
  File "/home/venv/lib/python3.9/site-packages/torch/export/_trace.py", line 1004, in wrapper
    ep = fn(*args, **kwargs)
  File "/home/venv/lib/python3.9/site-packages/torch/export/exported_program.py", line 122, in wrapper
    return fn(*args, **kwargs)
  File "/home/venv/lib/python3.9/site-packages/torch/export/_trace.py", line 1957, in _export
    export_artifact = export_func(  # type: ignore[operator]
  File "/home/venv/lib/python3.9/site-packages/torch/export/_trace.py", line 1251, in _strict_export
    return _strict_export_lower_to_aten_ir(
  File "/home/venv/lib/python3.9/site-packages/torch/export/_trace.py", line 1279, in _strict_export_lower_to_aten_ir
    gm_torch_level = _export_to_torch_ir(
  File "/home/venv/lib/python3.9/site-packages/torch/export/_trace.py", line 677, in _export_to_torch_ir
    raise UserError(UserErrorType.CONSTRAINT_VIOLATION, str(e))  # noqa: B904
torch._dynamo.exc.UserError: Constraints violated (seq_len)! For more information, run with TORCH_LOGS="+dynamic".
  - Not all values of seq_len = L['input_ids'].size()[1] in the specified range seq_len <= 128 are valid because seq_len was inferred to be a constant (6).
  - Not all values of seq_len = L['attention_mask'].size()[1] in the specified range seq_len <= 128 are valid because seq_len was inferred to be a constant (6).
  - Not all values of seq_len = L['position_ids'].size()[1] in the specified range seq_len <= 128 are valid because seq_len was inferred to be a constant (6).

Suggested fixes:
  seq_len = 6
```
</details>

Error has very clear suggestion: the `seq_len`, that corresponds the second dimension of the input should be a constant number. That is the input sentence length. However the input sentence lengths of the Llama3.2 model differ. 

This sounds like a perfect use case for [dynamic_shapes](https://pytorch.org/docs/stable/export.html#expressing-dynamism) of `torch.export.export`, but it does not work. What is the way to make compiled model accept inputs of various length?

### Versions

<details>
<summary>collect_env.py output</summary>

```
Collecting environment information...
PyTorch version: 2.6.0.dev20241106+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.10.12 (main, Sep 11 2024, 15:47:36) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-121-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 4090
Nvidia driver version: 550.120
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.1.4
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
Address sizes:                        39 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               4
On-line CPU(s) list:                  0-3
Vendor ID:                            GenuineIntel
Model name:                           Intel(R) Core(TM) i5-7640X CPU @ 4.00GHz
CPU family:                           6
Model:                                158
Thread(s) per core:                   1
Core(s) per socket:                   4
Socket(s):                            1
Stepping:                             9
CPU max MHz:                          4200.0000
CPU min MHz:                          800.0000
BogoMIPS:                             7999.96
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single pti ssbd ibrs ibpb stibp tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid mpx rdseed adx smap clflushopt intel_pt xsaveopt xsavec xgetbv1 xsaves dtherm ida arat pln pts hwp hwp_notify hwp_act_window hwp_epp md_clear flush_l1d arch_capabilities
Virtualization:                       VT-x
L1d cache:                            128 KiB (4 instances)
L1i cache:                            128 KiB (4 instances)
L2 cache:                             1 MiB (4 instances)
L3 cache:                             6 MiB (1 instance)
NUMA node(s):                         1
NUMA node0 CPU(s):                    0-3
Vulnerability Gather data sampling:   Mitigation; Microcode
Vulnerability Itlb multihit:          KVM: Mitigation: VMX disabled
Vulnerability L1tf:                   Mitigation; PTE Inversion; VMX conditional cache flushes, SMT disabled
Vulnerability Mds:                    Mitigation; Clear CPU buffers; SMT disabled
Vulnerability Meltdown:               Mitigation; PTI
Vulnerability Mmio stale data:        Mitigation; Clear CPU buffers; SMT disabled
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Mitigation; 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; IBRS; IBPB conditional; STIBP disabled; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                  Mitigation; Microcode
Vulnerability Tsx async abort:        Mitigation; TSX disabled

Versions of relevant libraries:
[pip3] numpy==2.1.3
[pip3] nvidia-cublas-cu12==12.4.5.8
[pip3] nvidia-cuda-cupti-cu12==12.4.127
[pip3] nvidia-cuda-nvrtc-cu12==12.4.127
[pip3] nvidia-cuda-runtime-cu12==12.4.127
[pip3] nvidia-cudnn-cu12==9.1.0.70
[pip3] nvidia-cufft-cu12==11.2.1.3
[pip3] nvidia-curand-cu12==10.3.5.147
[pip3] nvidia-cusolver-cu12==11.6.1.9
[pip3] nvidia-cusparse-cu12==12.3.1.170
[pip3] nvidia-cusparselt-cu12==0.6.2
[pip3] nvidia-nccl-cu12==2.21.5
[pip3] nvidia-nvjitlink-cu12==12.4.127
[pip3] nvidia-nvtx-cu12==12.4.127
[pip3] pytorch-triton==3.1.0+cf34004b8a
[pip3] torch==2.6.0.dev20241106+cu124
[conda] Could not collect
```
</details>

cc @ezyang @chauhang @penguinwu @bobrenjc93 @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4