[`FA2`] Fix flash attention 2 fine-tuning with Falcon
Author: younesbelkada
Pull request: yes

# What does this PR do?

Fixes https://github.com/huggingface/trl/issues/832 
Fixes https://github.com/huggingface/trl/issues/875
and all issues related with FA-2 + Falcon fine-tuning

Before this PR we were passing a `nn.Dropout` to the flash attention forward method, leading to an error since the dropout argument is expected to be a float. 

Also modified the test a bit to cover that usecase for future models

cc @ArthurZucker @lewtun 


Comments:
_The documentation is not available anymore as the PR was closed or merged._

@younesbelkada so the status right now is that falcon 180 b should work because it uses RoPE but 7b won't work until the alibi issue is fixed?

Falcon 40b instruct also seems to be rope, so I expected that to work but I get this error:
```
File /usr/local/lib/python3.10/dist-packages/transformers/modeling_utils.py:1270, in PreTrainedModel._check_and_enable_flash_attn_2(cls, config, torch_dtype, device_map)
   1255 """
   1256 If you don't know about Flash Attention, check out the official repository of flash attention:
   1257 https://github.com/Dao-AILab/flash-attention
   (...)
   1267 can initialize the correct attention module
   1268 """
   1269 if not cls._supports_flash_attn_2:
-> 1270     raise ValueError(
   1271         "The current architecture does not support Flash Attention 2.0. Please open an issue on GitHub to "
   1272         "request support for this architecture: https://github.com/huggingface/transformers/issues/new"
   1273     )
   1275 if not is_flash_attn_2_available():
   1276     raise ImportError(
   1277         "Flash Attention 2 is not available. Please refer to the documentation of https://github.com/Dao-AILab/flash-attention for"
   1278         " installing it. Make sure to have at least the version 2.1.0"
   1279     )

ValueError: The current architecture does not support Flash Attention 2.0. Please open an issue on GitHub to request support for this architecture: https://github.com/huggingface/transformers/issues/new
```
after trying:
```
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True, #adds speed with minimal loss of quality.
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

config = transformers.AutoConfig.from_pretrained(model_name_A, trust_remote_code=True)
config.max_seq_len = 4096 # (input + output) tokens can now be up to 4096

## Model A
model_A = AutoModelForCausalLM.from_pretrained(
    model_name_A,
    config=config,
    quantization_config=bnb_config,
    # rope_scaling={"type": "linear", "factor": 2.0},
    device_map='auto',
    trust_remote_code=True,
    use_flash_attention_2=True, # works with Llama models and reduces memory reqs
    cache_dir=cache_dir)
```

Hi @RonanKMcGovern 
The error is expected, you cannot use FA-2 + `trust_remote_code=True`. Can you please try again without `trust_remote_code=True`

Thanks @younesbelkada my mistake, that worked.

Thanks @RonanKMcGovern !

@younesbelkada this issue appears to be back now.

I had to revert to end of day yesterday commit to fix the issue:
```
!pip install git+https://github.com/huggingface/transformers.git@d39352d12c3f6a301e6c9d99218d1770eb513aca
```

BTW, seems like there are a lot of commits happening on main. Is that the usual approach? I would have thought the latest release would be somehow stable and new commits go to a dev branch or version.

Hmmm that's strange, I just checked and on the main branch https://github.com/huggingface/transformers/blob/main/src/transformers/models/falcon/modeling_falcon.py#L600 the falcon model contains the fix I introduced in this PR. Can you try:

```bash
pip install -U git+https://github.com/huggingface/transformers.git
```

Thanks, that's working. Perhaps what was causing the issue was me using:
```
pip install -q -U transformers
```
which probably only includes the last release but not commits on top?

Yes this forces you to use the last version on pypi which do not includes the commits we push on main