FalconForCausalLM does not support Flash Attention 2.0 yet
Author: menouarazib
Pull request: no

I attempted to use Flash Attention with the Falcon-7B model, but encountered the following error: **ValueError: FalconForCausalLM does not support Flash Attention 2.0 yet.**

This error was raised by the transformers/modeling_utils.py:

```
if not cls._supports_flash_attn_2:
    raise ValueError(
        f"{cls.__name__} does not support Flash Attention 2.0 yet. Please request to add support where"
        f" the model is hosted, on its model hub page: https://huggingface.co/{config._name_or_path}/discussions/new"
        " or in the Transformers GitHub repo: https://github.com/huggingface/transformers/issues/new"
    )
```

I installed the Transformers library from the GitHub repository using the following command:

`pip install git+https://github.com/huggingface/transformers`

Here is the code I used:

```
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

# Hugging Face Falcon-7B model ID
model_id = "tiiuae/falcon-7b"

# BitsAndBytesConfig for 4-bit integers
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True, 
    bnb_4bit_use_double_quant=True, 
    bnb_4bit_quant_type="nf4", 
    bnb_4bit_compute_dtype=torch.bfloat16
)

# Load the model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    device_map="auto",
    attn_implementation="flash_attention_2",
    torch_dtype=torch.bfloat16,
    quantization_config=bnb_config
)
```

Comments:
Hi @menouarazib, thanks for raising this issue! 

This is because the model being loading with this checkpoint, is from [code on the hub](https://huggingface.co/tiiuae/falcon-7b/blob/main/modeling_falcon.py) -- [mapping here](https://huggingface.co/tiiuae/falcon-7b/blob/898df1396f35e447d5fe44e0a3ccaaaa69f30d36/config.json#L14). 

Whereas the code in the library [does support FA2](https://github.com/huggingface/transformers/blob/f40b87de0ca234df61f76928956c4a2118c0b548/src/transformers/models/falcon/modeling_falcon.py#L946). 

This divergence is happening because the model was originally on the hub and then was ported into this library. The good news is, [FA has already been implemented for this model](https://github.com/huggingface/transformers/blob/f40b87de0ca234df61f76928956c4a2118c0b548/src/transformers/models/falcon/modeling_falcon.py#L540) and so it should be pretty easy to add to the hub model. I'd suggest opening a discussion there and requesting the addition with the model repo owners. 

cc @Rocketknight1 who know more about the intended hub vs transformers code usage and the recommended way to map between the two. 

If you want to use FA2 directly with Falcon, you can import the model from transformers directly instead of using the auto model: 

```py
import torch
from transformers import AutoTokenizer, BitsAndBytesConfig, FalconForCausalLM

# Hugging Face Falcon-7B model ID
model_id = "tiiuae/falcon-7b"

# BitsAndBytesConfig for 4-bit integers
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

# Load the model and tokenizer
model = FalconForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    device_map="auto",
    attn_implementation="flash_attention_2",
    torch_dtype=torch.bfloat16,
    quantization_config=bnb_config
)
```

Thanks, @amyeroberts, for this clarification. It works when using `FalconForCausalLM from transformers `. 
I have created a discussion to add Flash Attention to the model hosted on the hub:
https://huggingface.co/tiiuae/falcon-7b/discussions/98

Thanks