Repository: pytorch
Issue: [export] Llama3 export with dynamic shapes fails with constraint violations
Labels: onnx-needs-info,module: dynamic shapes,oncall: export

### 🐛 Describe the bug

I'm trying to export Llama3-8B model with dynamic shapes and encountered the following error. The same error occurs for Llama-2 export as well. Any suggestions ? 

```py
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct").eval().cuda()
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
base_prompt = "How many hours are in a day?"
base_inputs = tokenizer(base_prompt, return_tensors="pt").to("cuda:0")
input_ids = base_inputs.input_ids

seq_len = torch.export.Dim("seq_len", min=1, max=1024)
with torch.no_grad():
    ep = torch.export.export(model, (input_ids,), dynamic_shapes=({1: seq_len},))
```
The above core results in the following error 
```py
orch._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 <= 1024 satisfy the generated guard Eq(Mod(L['input_ids'].size()[1]**2 + L['input_ids'].size()[1], 8), 0).
  - Not all values of seq_len = L['input_ids'].size()[1] in the specified range seq_len <= 1024 satisfy the generated guard Ne(Mod(L['input_ids'].size()[1] + 1, 8), 0).
  - Not all values of seq_len = L['input_ids'].size()[1] in the specified range seq_len <= 1024 satisfy the generated guard 8 - Mod(L['input_ids'].size()[1], 8) > 0.

Suggested fixes:
  _seq_len = Dim('_seq_len', min=1, max=128)  # 2 <= 8*_seq_len <= 1024
  seq_len = 8*_seq_len
```
After applying the suggested fixes, I changed the line 
```py
# Previous line
seq_len = torch.export.Dim("seq_len", min=1, max=1024)
# Modified as follows as per torch dynamo suggestions 
_seq_len = torch.export.Dim('_seq_len', min=1, max=128) 
seq_len = 8*_seq_len
```
this results in another error:
```py
 ep = torch.export.export(model, (input_ids,), dynamic_shapes=({1: seq_len},))
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/export/__init__.py", line 174, in export
    return _export(
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/export/_trace.py", line 635, in wrapper
    raise e
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/export/_trace.py", line 618, in wrapper
    ep = fn(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/export/exported_program.py", line 83, in wrapper
    return fn(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/export/_trace.py", line 860, in _export
    gm_torch_level = _export_to_torch_ir(
  File "/root/.pyenv/versions/3.10.13/lib/python3.10/site-packages/torch/export/_trace.py", line 359, in _export_to_torch_ir
    raise UserError(UserErrorType.CONSTRAINT_VIOLATION, str(e))  # noqa: TRY200
torch._dynamo.exc.UserError: 8 not in range [16, 1024]
```
Note : Exporting model on cpu works (if I remove .eval().cuda())


### Versions

Versions of relevant libraries:
transformers==4.40.0
[pip3] numpy==1.26.4
[pip3] torch==2.3.0+cu121
[pip3] torch_tensorrt==2.3.0.dev0+4323e3646
[pip3] torchvision==0.18.0+cu121
[pip3] triton==2.3.0

cc: @angelayi 


cc @ezyang @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4