LLama can't use `torch.compile()`
Author: ydshieh
Pull request: no

### System info

- `transformers` version: 4.43.0.dev0
- Platform: Linux-5.10.0-30-cloud-amd64-x86_64-with-glibc2.29
- Python version: 3.8.10
- Huggingface_hub version: 0.23.4
- Safetensors version: 0.4.3
- Accelerate version: 0.32.0.dev0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.3.0+cu121 (True)
- Tensorflow version (GPU?): 2.13.0 (False)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA A100-SXM4-80GB

### Reproduction

```python
import torch
from transformers import GemmaForCausalLM, LlamaForCausalLM, AutoModelForCausalLM
from datasets import load_dataset

import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"


if __name__ == "__main__":
    from transformers import AutoTokenizer

    ckpt = "google/gemma-2b"
    ckpt = "meta-llama/Meta-Llama-3-8B"

    #model = AutoModelForCausalLM.from_pretrained(ckpt)
    model = AutoModelForCausalLM.from_pretrained(ckpt)

    tokenizer = AutoTokenizer.from_pretrained(ckpt)
    model.to(dtype=torch.bfloat16, device="cuda")
    model.train()

    model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)

    dataset = load_dataset("tiiuae/falcon-refinedweb", streaming=True, split='train')
    dataset_iter = iter(dataset)

    for _ in range(2):

        sample = next(dataset_iter)
        prompt = sample["content"]
        inputs = tokenizer(prompt, return_tensors="pt", max_length=64, truncation=True)

        with torch.no_grad():
            outputs = model.forward(**inputs)
```

### Error 1 (this fails even without `model.train()`, and `Gemma` has no such issue)
```bash
    raise Unsupported(msg)
torch._dynamo.exc.Unsupported: call_method UserDefinedObjectVariable(Logger) warning_once [ConstantVariable()] {}

from user code:
   File "/transformers/src/transformers/models/llama/modeling_llama.py", line 1069, in forward
    outputs = self.model(
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/transformers/src/transformers/models/llama/modeling_llama.py", line 832, in forward
    logger.warning_once(

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information
```

### Error 2 (after removing ` logger.warning_once` in the code)

```bash
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 "/transformers/src/transformers/models/llama/modeling_llama.py", line 1065, in forward
    outputs = self.model(
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/transformers/src/transformers/models/llama/modeling_llama.py", line 841, in forward
    causal_mask = self._update_causal_mask(
  File "/transformers/src/transformers/models/llama/modeling_llama.py", line 932, in _update_causal_mask
    if AttentionMaskConverter._ignore_causal_mask_sdpa(
  File "/transformers/src/transformers/modeling_attn_mask_utils.py", line 279, in _ignore_causal_mask_sdpa
    elif (is_training or not is_tracing) and torch.all(attention_mask == 1):
```

### Expected behavior

It should be able to work with `torch.compile`

Comments:
Ah, I see why! I am not using `generate`, and the cache object is not Static Cache, so it enters

```
        # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
        if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions:
            if AttentionMaskConverter._ignore_causal_mask_sdpa(
                attention_mask,
                inputs_embeds=input_tensor,
                past_key_values_length=past_seen_tokens,
                is_training=self.training,
            ):
```
then failed inside `_ignore_causal_mask_sdpa`.

I will close this issue.

Close as it not a real issue. For training, keep in mind to setup the static cache if you are using `torch.compile`