Repository: transformers
Issue: ValueError: MllamaForConditionalGeneration does not support Flash Attention 2.0 yet
Labels: bug

### System Info

- `transformers` version: 4.47.1
- Platform: Linux-4.18.0-553.22.1.el8_10.x86_64-x86_64-with-glibc2.28
- Python version: 3.10.15
- Huggingface_hub version: 0.26.3
- Safetensors version: 0.4.5
- Accelerate version: 0.34.2
- Accelerate config:    not found
- PyTorch version (GPU?): 2.5.1+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 distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA H200

### Who can help?

@amyeroberts, @qubvel

### Information

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

### Tasks

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

### Reproduction

code:
```
import requests
import torch
from PIL import Image
from transformers import MllamaForConditionalGeneration, AutoProcessor

model_id = "meta-llama/Llama-3.2-11B-Vision-Instruct"
model = MllamaForConditionalGeneration.from_pretrained(model_id, device_map="auto", torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2")
processor = AutoProcessor.from_pretrained(model_id)

messages = [ 
    [   
        {
            "role": "user", 
            "content": [
                {"type": "image"},
                {"type": "text", "text": "What does the image show?"}
            ]
        }
    ],  
]
text = processor.apply_chat_template(messages, add_generation_prompt=True)

url = "https://llava-vl.github.io/static/images/view.jpg"
image = Image.open(requests.get(url, stream=True).raw)

inputs = processor(text=text, images=image, return_tensors="pt").to(model.device)
output = model.generate(**inputs, max_new_tokens=25)
print(processor.decode(output[0]))
```

output:
Traceback (most recent call last):
  File "/home/user/reasoning/test_mllama.py", line 8, in <module>
    model = MllamaForConditionalGeneration.from_pretrained(model_id, device_map="auto", torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2")                                                                                                           
  File "/home/user/cache/conda/envs/openinstruct/lib/python3.10/site-packages/transformers/modeling_utils.py", line 4124, in from_pretrained                                                                     
    config = cls._autoset_attn_implementation(
  File "/home/user/cache/conda/envs/openinstruct/lib/python3.10/site-packages/transformers/modeling_utils.py", line 1586, in _autoset_attn_implementation                                                        
    cls._check_and_enable_flash_attn_2(
  File "/home/user/cache/conda/envs/openinstruct/lib/python3.10/site-packages/transformers/modeling_utils.py", line 1707, in _check_and_enable_flash_attn_2                                                      
    raise ValueError(
ValueError: MllamaForConditionalGeneration does not support Flash Attention 2.0 yet. Please request to add support where the model is hosted, on its model hub page: https://huggingface.co/meta-llama/Llama-3.2-11B-Vision-Instruct/discussions/new or in the Transformers GitHub repo: https://github.com/huggingface/transformers/issues/new

### Expected behavior

Support flash attention 2