[`core` ] Integrate Flash attention 2 in most used models
Author: younesbelkada
Pull request: yes

# What does this PR do?

Based on flash attention official repository and code snippets shared internally by @pacman100 & @fxmarty and also based on some internal discussion with @LysandreJik, made a PoC of what a Flash Attention 2 integration would look like in transformers. 

We should restrict the integration to Flash Attention 2 only for now, as there is a way to run Flash-Attention 1 through torch.SDPA + `BetterTransformer` API that is explained here: https://github.com/huggingface/transformers/pull/25265 

I added it only for Llama for now but could be easily extended for other architectures (though I think Alibi is not supported but not sure).

Note that the performance with this integration will not be optimal in this case as flash attention shines in a batched setting when the KV caching is done in a specific format. 

Draft for now.

## API 

Currently the API is very simple:

<details><summary>API example</summary>

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, LlamaForCausalLM

model_id = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=True)

model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    torch_dtype=torch.bfloat16, 
    use_auth_token=True, 
    low_cpu_mem_usage=True,
    use_flash_attention_2=True,
).to(0)

text = [
    "Hello how are you?",
    "hi"
]

tokenizer.pad_token = tokenizer.eos_token

inputs = tokenizer(text, return_tensors="pt", padding=True).to(0)

outputs = model.generate(inputs["input_ids"], attention_mask=inputs["attention_mask"], use_cache=False, max_new_tokens=30, do_sample=False)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
```

</details>

## TODOs:

- [x] Benchmarks
- [x] Add padded input support
- [x] Stronger UX: block users if no GPU available, check if no device_map --> warn users to set the model on GPU / warn if no half precision, etc. 
- [x] Tests
- [ ] Docs
- [x] add it for falcon
- [ ] add it for starcoder

cc @fxmarty @pacman100 @LysandreJik 

Comments:
I think attempting to plug Hazy-flash in transformers and properly benchmarking against vanilla transformers and against BetterTransformer is a **very good thing** to possibly motivate:
* Upstreaming SDPA or Hazy-flash natively in transformers.
* a refactorization of the KV cache to rely on indexing rather than padding, or make the KV cache implementation more modular (against transformers philosophy though).

However, there should in my opinion be a serious internal discussion about what goes natively in transformers and whatnot when it comes to hardware optimization. Hazy-flash relies heavily on CUTLASS, that can not be transpiled for AMD devices.

There is [this fork](https://github.com/ROCmSoftwarePlatform/flash-attention) of flash for RoCm. If we are to integrate such optimizations in transformers natively, would we be comfortable doing so for a variety of hardware providers?

In my opinion, an approach a la BetterTransformer replacing modules or methods can make sense as well. The issue with the code injection approach is that it can make it more difficult to combine several injections (for example BT + bitsandbytes - though I hear they currently work smoothly together).

Anyway, very keen to help to benchmark!

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

🚀 Like the new design a lot better

Current state for batched generation (with padding, which is the common case for batched generation):

![image](https://github.com/huggingface/transformers/assets/9808326/dbec148b-c40b-4da0-bf05-303a112aa538)

![image](https://github.com/huggingface/transformers/assets/9808326/21d0f442-68c1-46d3-bb5a-acbddced3915)

IMO it is unrealistic to use FA for batched generation without using the new https://github.com/Dao-AILab/flash-attention/blob/37c6e0540650658e466e42b4d2e15925b8cfbe24/flash_attn/flash_attn_interface.py#L796. I would say let's just disable FA2 when using batched generation.

Current state for training (or use_cache=False) with padding: The current implementation unpad/pad many times which is not necessary and could be done only once at the beginning/end. This adds some overhead as well, I'll profile.

Just ran one training, the bettertransformer implementation took 9.5 hours per epoch with llama v2
This PR takes 6 hours per epoch
Batch size 1, 4 bit bnb, 7b model, 4096 seq length, 4.5% trainable params using peft and large lora depth
RTX A6000

Nice speedup then with respect to BT ! That's great to hear @flozi00 , we will run final benchmarks, try to see if we can further optimize this PR and mark it as ready. Out of curiosity, do you use any padding token in the training data ?

As BT does not support training with padding tokens I assume you ran your experiment without padding tokens (i.e. `packing=True` in case you are using `SFTTrainer` from TRL library

> 

Yeah, no padding.
I used an similiar function to the packing function of trl library

Awesome, thanks for confirming @flozi00 

@flozi00 Were you using torch 2.0.1? I wonder if the difference comes just from flash v2 (this PR) vs flash v1 (torch 2.0.1), or if there is more to it (some overlead in the BT + SDPA path).

Yes, the latest pytorch release (2.0.1), cuda 11.8
optimum and transformers from the latest release

Update: with the latest commit 7f06af6226ad6e50164c59ce7a7201d12cd3acfe we avoid some unnecessary operations. However we do still unpad/pad at each layer which remains expensive.


> Update: with the latest commit [7f06af6](https://github.com/huggingface/transformers/commit/7f06af6226ad6e50164c59ce7a7201d12cd3acfe) we avoid some unnecessary operations. However we do still unpad/pad at each layer which remains expensive.

Would you mind making an issue on PyTorch for this so that we can track enablement?

> As BT does not support training with padding tokens I assume you ran your experiment without padding tokens (i.e. `packing=True` in case you are using `SFTTrainer` from TRL library

Could you also make an issue for this so that we can track enablement:
https://github.com/pytorch/pytorch/pull/97485 fwiw


cc @jbschlosser

Thank you @drisspg! Opened https://github.com/pytorch/pytorch/issues/108770. For the second point, it is indeed the PR you linked + BT not implementing a path using nested tensors.

I think the PR is in a nice state, requesting some reviews before adding the documentation! 
![Screenshot 2023-09-11 at 19 16 06](https://github.com/huggingface/transformers/assets/49240599/27f5d79f-2dd4-447b-9839-37e3b2713b98)
Speedups I get for falcon-7b model using this PR with pure forward! 

cc @ArthurZucker @patrickvonplaten @sgugger @pacman100 



FYI out of curiosity I gave a try to remove the unpad/pad overhead when `use_cache=False` (e.g. padded training), this overhead is actually rather large: https://github.com/younesbelkada/transformers/pull/5

Thanks all for your comments and review, the PR is ready for a final review!

This PR broke my custom attention module in AutoAWQ because of the introduction of a new `padding_mask` argument. This should be in the release notes if the intention is to keep the `padding_mask` as an input argument to the attention module - for better visibility. Fixed by accepting `*args, **kwargs` in custom modules and leaving implementation of Flash attention for later. cc @younesbelkada