ValueError: MllamaForConditionalGeneration does not support Flash Attention 2.0 yet
Author: yxchng
Pull request: no
Labels: name: Feature request

name: 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

Comments:
Indeed Mllama currently has no FA2 support. We are currently in progress of changing the API for different attention implementations and I believe Mllama will support FA2 out-of-the-box when the changes from https://github.com/huggingface/transformers/pull/35235 are applied

Mmm it's a bit more complicated than that, as you need both cross and self to support flash! But yeah we'll find a way to do that! 

The `mllama` code is also not super super efficient at the moment as it is not flexible. 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.