torch.export fails for llama model 
Author: mreso
Pull request: no

### System Info

- `transformers` version: 4.38.0.dev0
- Platform: Linux-5.12.0-0_fbk15_zion_7536_ged335dd7211d-x86_64-with-glibc2.34
- Python version: 3.11.7
- Huggingface_hub version: 0.20.3
- Safetensors version: 0.4.2
- Accelerate version: 0.27.2
- Accelerate config:    not found
- PyTorch version (GPU?): 2.3.0.dev20240221+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: no
- Using distributed or parallel set-up in script?: no

### Who can help?

@ArthurZucker @gante

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

```
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

llama = AutoModelForCausalLM.from_pretrained(
        "meta-llama/Llama-2-7b-chat-hf", low_cpu_mem_usage=True
    )
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")

llama.eval()
inputs = tokenizer("How do you", return_tensors="pt")

exported_llama = torch.export.export(llama, args=(inputs["input_ids"],inputs["attention_mask"]))
```

### Expected behavior

Successfully export of the model but fails with:
```
torch._dynamo.exc.UserError: Dynamic control flow is not supported at the moment. Please use functorch.experimental.control_flow.cond to explicitly capture the control flow. For more information about this error, see: https://pytorch.org/docs/main/generated/exportdb/index.html#cond-operands

from user code:
   File "/home/mreso/.conda/envs/pippy/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py", line 1176, in forward
    outputs = self.model(
  File "/home/mreso/.conda/envs/pippy/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/mreso/.conda/envs/pippy/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py", line 990, in forward
    causal_mask = self._update_causal_mask(attention_mask, inputs_embeds)
  File "/home/mreso/.conda/envs/pippy/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py", line 1090, in _update_causal_mask
    if not is_tracing and attention_mask is not None and torch.any(attention_mask != 1):
```
I worked around this issue by guarding with torch._dynamo_is_compiling() in these two places:
https://github.com/huggingface/transformers/blob/fc37f38915372c15992b540dfcbbe00a916d4fc6/src/transformers/models/llama/modeling_llama.py#L1080
```
is_tracing = torch.jit.is_tracing() or isinstance(input_tensor, torch.fx.Proxy) or (hasattr(torch,"_dynamo") and not torch._dynamo.is_compiling())
```
https://github.com/huggingface/transformers/blob/fc37f38915372c15992b540dfcbbe00a916d4fc6/src/transformers/models/llama/modeling_llama.py#L131
```
if hasattr(torch,"_dynamo") and not torch._dynamo.is_compiling():
    self._cos_cached = cos
    self._sin_cached = sin
```
Happy to create a PR if this is a viable solution.

Comments:
cc @fxmarty as well

@mreso This should be fixed (at least the first part) in https://github.com/huggingface/transformers/pull/29211.

I did not hit an issue with `self._cos_cached = cos`, with your repro, maybe we are not using the same commit? Can you try on main again?

Maybe related https://github.com/huggingface/transformers/pull/29109 & https://github.com/huggingface/transformers/pull/29198

Thanks @fxmarty for fixing the first part! Just reinstalled from source using commit 2cc8cf6 and rechecked. Still running into the second issue with the code snippet from above:
```
  File "/home/mreso/.conda/envs/pippy/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2266, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mreso/.conda/envs/pippy/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2380, in inline_call_
    tracer.run()
  File "/home/mreso/.conda/envs/pippy/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 793, in run
    and self.step()
        ^^^^^^^^^^^
  File "/home/mreso/.conda/envs/pippy/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 756, in step
    getattr(self, inst.opname)(inst)
  File "/home/mreso/.conda/envs/pippy/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1308, in STORE_ATTR
    assert (
AssertionError: Mutating module attribute _cos_cached during export.

from user code:
   File "/home/mreso/.conda/envs/pippy/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py", line 1175, in forward
    outputs = self.model(
  File "/home/mreso/.conda/envs/pippy/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/mreso/.conda/envs/pippy/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py", line 1011, in forward
    layer_outputs = decoder_layer(
  File "/home/mreso/.conda/envs/pippy/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/mreso/.conda/envs/pippy/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py", line 735, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
  File "/home/mreso/.conda/envs/pippy/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/mreso/.conda/envs/pippy/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py", line 641, in forward
    cos, sin = self.rotary_emb(value_states, position_ids)
  File "/home/mreso/.conda/envs/pippy/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/mreso/.conda/envs/pippy/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py", line 131, in forward
    self._cos_cached = cos

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information
```


this will be fixed by https://github.com/huggingface/transformers/pull/29198, there is no longer `self._cos_cached = cos`

@ArthurZucker  I have seen similar issues with  mistralai/Mixtral-8x7B-v0.1 model, I have seen there are 13 graph breaks because of conditional statements and call backs to python, 

is there any plan to support dynamo full_graph support for maximum of models ?

Closing as https://github.com/huggingface/transformers/pull/29198 is merged

@tbaggu that would be the goal indeed.

Hi, is this still supported? I'm getting an error running the same code snippet / script:
```
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:05<00:00,  2.92s/it]
---------------------------------------------------------------------------
Unsupported                               Traceback (most recent call last)
File ~/projects/tracer/src/neural_tracer/llama3.py:14
     11 llama.eval()
     12 inputs = tokenizer("How do you", return_tensors="pt")
---> 14 exported_llama = torch.export.export(llama, args=(inputs["input_ids"],inputs["attention_mask"]))

File ~/projects/tracer/.venv/lib/python3.12/site-packages/torch/export/__init__.py:360, in export(mod, args, kwargs, dynamic_shapes, strict, preserve_module_call_signature)
    354 if isinstance(mod, torch.jit.ScriptModule):
    355     raise ValueError(
    356         "Exporting a ScriptModule is not supported. "
    357         "Maybe try converting your ScriptModule to an ExportedProgram "
    358         "using `TS2EPConverter(mod, args, kwargs).convert()` instead."
    359     )
--> 360 return _export(
    361     mod,
    362     args,
    363     kwargs,
    364     dynamic_shapes,
    365     strict=strict,
    366     preserve_module_call_signature=preserve_module_call_signature,
    367     pre_dispatch=True,
    368 )
File ~/projects/tracer/.venv/lib/python3.12/site-packages/torch/export/_trace.py:1092, in _log_export_wrapper.<locals>.wrapper(*args, **kwargs)
   1085     else:
   1086         log_export_usage(
   1087             event="export.error.unclassified",
   1088             type=error_type,
   1089             message=str(e),
   1090             flags=_EXPORT_FLAGS,
   1091         )
-> 1092     raise e
   1093 finally:
   1094     _EXPORT_FLAGS = None

File ~/projects/tracer/.venv/lib/python3.12/site-packages/torch/export/_trace.py:1065, in _log_export_wrapper.<locals>.wrapper(*args, **kwargs)
   1063 try:
   1064     start = time.time()
-> 1065     ep = fn(*args, **kwargs)
   1066     end = time.time()
   1067     log_export_usage(
   1068         event="export.time",
   1069         metrics=end - start,
   1070         flags=_EXPORT_FLAGS,
   1071         **get_ep_stats(ep),
   1072     )

File ~/projects/tracer/.venv/lib/python3.12/site-packages/torch/export/exported_program.py:121, in _disable_prexisiting_fake_mode.<locals>.wrapper(*args, **kwargs)
    118 @functools.wraps(fn)
    119 def wrapper(*args, **kwargs):
    120     with unset_fake_temporarily():
--> 121         return fn(*args, **kwargs)

File ~/projects/tracer/.venv/lib/python3.12/site-packages/torch/export/_trace.py:2112, in _export(mod, args, kwargs, dynamic_shapes, strict, preserve_module_call_signature, pre_dispatch, allow_complex_guards_as_runtime_asserts, _is_torch_jit_trace)
   2104 # NOTE Export training IR rollout
   2105 # Old export calls export._trace(pre_dispatch=True)
   2106 # and there are still lot of internal/OSS callsites that
   (...)   2109 # export_training_ir_rollout_check returns True in OSS
   2110 # while internally it returns False UNLESS otherwise specified.
   2111 if pre_dispatch and export_training_ir_rollout_check():
-> 2112     ep = _export_for_training(
   2113         mod,
   2114         args,
   2115         kwargs,
   2116         dynamic_shapes,
   2117         strict=strict,
   2118         preserve_module_call_signature=preserve_module_call_signature,
   2119     )
   2120     dtrace_structured("exported_program", payload_fn=lambda: str(ep))
   2121     return e

```

There are more lines in the error, but I feel the first part is telling, i.e.
```
    356         "Exporting a ScriptModule is not supported. "
    357         "Maybe try converting your ScriptModule to an ExportedProgram "
    358         "using `TS2EPConverter(mod, args, kwargs).convert()` instead."
```