Community contribution: Adding Flash Attention 2 support for more architectures
Author: younesbelkada
Pull request: no
Labels: name: Good Second Issue

### Feature request

Flash Attention 2 is a library that provides attention operation kernels for faster and more memory efficient inference and training: https://github.com/Dao-AILab/flash-attention

![Screenshot 2023-09-22 at 17 49 18](https://github.com/huggingface/transformers/assets/49240599/1395f962-26ca-4728-a8d0-085792295c28)


Let's try to add Flash Attention 2 support for more architectures! Currently supported architectures are

- [x] Llama
- [x] Falcon

It would be great to add the support for more architectures such as

- [x] Bark
- [x] Bart
- [ ] BERT | @sorenmc 
- [ ] CLIP https://github.com/huggingface/transformers/pull/27444/
- [x] DistilBERT
- [x] GPT-2
- [x] GPT-J
- [x] GPTBigCode (Starcoder) | @susnato 
- [x] GPT-neo
- [x] GPT-neo-x | @younesbelkada #26463
- [x] OPT | @susnato #26414 
- [x] Llava
- [x] VipLlava
- [x] mBART
- [x] Mistral
- [x] Mixtral
- [ ] MPT | @rajveer43 
- [ ] T5
- [ ] Persimmon | @jeromeku 
- [x] Phi
- [x] Whisper
- [x] Qwen2

... and many more

Adding this feature would require to follow the same protocol as in https://github.com/huggingface/transformers/pull/25598 
 . First create a new module inside the corresponding modeling file termed as `xxxFlashAttention` that inherits from `xxxAttention` and override the foward method to use the public methods from `flash-attn`. Make sure to have access to a GPU that supports Flash Attention 2. 

Given the slight challenge of the issue, labelling it as a good second issue!

If you are interested to take up the challenge, comment below with the architecture name you want to integrate and open a PR!

Once you open a PR, feel free to ping @LysandreJik @ArthurZucker @amyeroberts @younesbelkada @fxmarty @SunMarc @pacman100 for a review

### Motivation

Making LLMs more memory efficient and faster ! 

### Your contribution

Reviewing PRs and possibly adding the support for more models

Comments:
Hi @younesbelkada - I want to work on adding Flash Attention 2 support for GPTBigCode (Starcoder). Can I take this task? Can you please assign this task to me?


Will definitely take a look next week 
Great to see it merged now 💪

I would like to work on `MPT` @younesbelkada 

I would like to work on `OPT`.

Is it possible to add FlashAttention2 to GPT2 models?

@sahilbhosale63 @flozi00 @rajveer43 @susnato thanks very much for your interest! Indeed it would be great if you could help us!
Before assigning you to this issue can you confirm you have access to a GPU that does support Flash Attention 2: https://github.com/Dao-AILab/flash-attention#installation-and-features in order to be able to run the tests ?
@ZeusFSX , yes I think that it is possible, I'll update the list accodingly

@younesbelkada Yes I have

OK perfect, I will assign you to MPT ! Feel free to let me know if you need any help or if you have any question, as a starting point, I would recommend to have a look at #25598 and see if you can replicate the PR for MPT. For running flash attention tests you can just run (once PR is ready):
```bash
RUN_SLOW=1 pytest -m flash_attn_test tests/models/mpt/
```

@younesbelkada yes I have.

Thanks @susnato , perfect then, let me know whenever you start the PR and if you have any question ! Check out my instructions above for more details

@younesbelkada Unfortunately, My GPU is not supported

> OK perfect, I will assign you to MPT ! Feel free to let me know if you need any help or if you have any question, as a starting point, I would recommend to have a look at #25598 and see if you can replicate the PR for MPT. For running flash attention tests you can just run (once PR is ready):
> 
> ```shell
> RUN_SLOW=1 pytest -m flash_attn_tests tests/models/mpt/
> ```

Sure I will work on it!

@younesbelkada Would like to work on Persimmon.  I have access to A4000, A5000, and A6000, which I believe should be compatible with FA2.

Perfect sounds great, thanks for your help, I will assign you to Persimmon !

Since @sahilbhosale63 is not working on ` GPTBigCode (Starcoder)`(as he said [here](https://github.com/huggingface/transformers/issues/26350#issuecomment-1733685691)) can I take that @younesbelkada?

Yes no problem, thanks very much for proposing your help on this ! As a starting point you can have a look at @pacman100 's implementation here: https://github.com/pacman100/DHS-LLM-Workshop/blob/main/personal_copilot/training/starcoder_flash_attn_monkey_patch.py

@younesbelkada I would like to implement it for BERT if it hasn't already been done? A lot of the models topping MTEB are still relying on this architecture! I have tested that i can run flash attention 2 on my nvidia geforce RTX 3060 TI.

Awesome, thanks a lot for your help, ok I will assign you to BERT then!

Hi everyone, I would like to help implement this with GPT2 if you want.

@younesbelkada

I have a working version for `Persimmon` that passes the `flash_attn_v2` tests except for `generate_padding_right` as the original `PersimmonFlashAttention` does not have `padding_mask` as a kw input (as opposed to the `Llama` and `Falcon` flash implementations).  Is this something that needs to be changed in both Persimmon Flash v1 and v2? 

Also, any plans on incorporating additional optimizations, e.g., `Flash Attention` repo has fused layers for `dense`, `rotary`, and `layer norm` for[ faster training](https://github.com/Dao-AILab/flash-attention/blob/main/README.md#full-model-code-and-training-script); and [Triton](https://github.com/Dao-AILab/flash-attention/blob/main/README.md#triton-implementation-of-flashattention) kernels, more generally?  Happy to investigate more!

Also, would like to help with `Mistral-7b` (just released).  They use `xformers` `memory efficient attention` in their released implementation but also mention Tri Dao's FA in the blogpost. 

Hi @DougTrajano 
Awesome! Can you confirm you have access to a hardware that is supported by FA-2? 

![Screenshot 2023-09-28 at 11 23 36](https://github.com/huggingface/transformers/assets/49240599/e509a60b-f011-40de-a450-cf7d5d291671)

@jeromeku awesome thanks! Can you move forward for Persimmon by opening a PR so that I can have a look?

> Also, any plans on incorporating additional optimizations, e.g., Flash Attention repo has fused layers for dense, rotary, and layer norm for[ faster training](https://github.com/Dao-AILab/flash-attention/blob/main/README.md#full-model-code-and-training-script); and [Triton](https://github.com/Dao-AILab/flash-attention/blob/main/README.md#triton-implementation-of-flashattention) kernels, more generally? Happy to investigate more!

If that is something that can nicely fit into the API without any breaking behaviour that would be great !

> Also, would like to help with Mistral-7b (just released). They use xformers memory efficient attention in their released implementation but also mention Tri Dao's FA in the blogpost.

I think Mistral's attention has been released in the latest version of FA-2 --> Would you be happy to open a PoC PR so that I can play with it and see what we can do?

Again thanks a lot!

Hi @jeromeku 
I had to check internally for Mistral, given the very recent release and the urgency, we'll take this over (https://github.com/huggingface/transformers/pull/26464); if you have started a PR, I'm very happy to start from it or to add you as a co-author to the PR ! 
We might also refactor things a bit to support Local attention introduced by Mistral, so that needs further investigation, I'll keep you posted

@younesbelkada what is the expected deadline to complete `MPT`, I have other issues to tackle on so I can plan accordingly

Hi @younesbelkada , I am talking this up for `GPT-neo`.

Awesome @susnato ! Thanks ! 
@rajveer43 thanks for taking up MPT, will check it out!

> Hi @DougTrajano Awesome! Can you confirm you have access to a hardware that is supported by FA-2?
> 
> ![Screenshot 2023-09-28 at 11 23 36](https://user-images.githubusercontent.com/49240599/271235959-e509a60b-f011-40de-a450-cf7d5d291671.png)
> 

Yes, I'll work on AWS SageMaker.



Would love to take on GPT2!

Thanks for confirming @DougTrajano !
@marcasty thanks a lot for your interest, @DougTrajano has taken up GPT2, would be happy taking another model? 🙏 
Can you also confirm you have access to a hardware that support FA-2 ?

Hi @younesbelkada, I am taking this up for `DistillBERT`.

@younesbelkada what about T5? I have access to compatible hardware 

Hey @younesbelkada, are other param-sizes of Llama architecture supported, such as 1.1B and 34B? I tried using flash-attention-2 with [PY007/TinyLlama-1.1B-intermediate-step-240k-503b](https://huggingface.co/PY007/TinyLlama-1.1B-intermediate-step-240k-503b) and [codellama/CodeLlama-34b-hf](https://huggingface.co/codellama/CodeLlama-34b-hf), however I run into shape issues.

For ex, I get the following error for codellama-34b.
```py
File "/home/ray/anaconda3/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 489, in forward
    attn_output = self._flash_attention_forward(
  File "/home/ray/anaconda3/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 526, in _flash_attention_forward
    query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
  File "/home/ray/anaconda3/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 562, in _upad_input
    query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
RuntimeError: shape '[1024, 8, 128]' is invalid for input of size 8388608
```

Hi @younesbelkada  adding support for GPT2

> Hi @younesbelkada adding support for GPT2

well, I'll stop working on it. good work!

> > Hi @younesbelkada adding support for GPT2
> 
> well, I'll stop working on it. good work!

We can collaborate on it if you want. The PR is in draft, and  I am still working it 

> > > Hi @younesbelkada adding support for GPT2
> > 
> > 
> > well, I'll stop working on it. good work!
> 
> We can collaborate on it if you want. The PR is in draft, and I am still working it

That's fine buddy. I'll help with another issue. I just added one comment in the PR since I saw the changes. Thank you!

Hi @marcasty  - sure yes ! that sounds great!
@nidhishs  would you mind opening a new ticket for the issue with a reproducible snippet? 🙏 I think the issue might be fixed in https://github.com/huggingface/transformers/pull/26490 - can you try again on main? 


> Hey @younesbelkada, are other param-sizes of Llama architecture supported, such as 1.1B and 34B? I tried using flash-attention-2 with [PY007/TinyLlama-1.1B-intermediate-step-240k-503b](https://huggingface.co/PY007/TinyLlama-1.1B-intermediate-step-240k-503b) and [codellama/CodeLlama-34b-hf](https://huggingface.co/codellama/CodeLlama-34b-hf), however I run into shape issues.
> 
> For ex, I get the following error for codellama-34b.
> 
> ```python
> File "/home/ray/anaconda3/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 489, in forward
>     attn_output = self._flash_attention_forward(
>   File "/home/ray/anaconda3/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 526, in _flash_attention_forward
>     query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
>   File "/home/ray/anaconda3/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 562, in _upad_input
>     query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
> RuntimeError: shape '[1024, 8, 128]' is invalid for input of size 8388608
> ```

Hi @younesbelkada, just wanted to bring your attention to this.

Hi @nidhishs please see my message above (sorry I have initially tagged the wrong person)

> @nidhishs would you mind opening a new ticket for the issue with a reproducible snippet? 🙏 I think the issue might be fixed in https://github.com/huggingface/transformers/pull/26490 - can you try again on main?

> @younesbelkada I would like to implement it for BERT if it hasn't already been done? A lot of the models topping MTEB are still relying on this architecture! I have tested that i can run flash attention 2 on my nvidia geforce RTX 3060 TI.

Hi @sorenmc ! I would like to ask if some help is needed on BERT FA2-integration, though I am not a very skillful programmer) I'd like to try, cause I am facing an issue at work, where this feature might accelerate some studies

Hi @MrTimmy89 haven't had too much free time the last week so I have only been working a little bit on it every night. I almost have a working version and am planning to get the PR up to night! It will only support standard attention for now $$Softmax\left(\frac{QK^T}{\sqrt{d}}\right)V$$ as opposed to huggingfaces standard bert implementation that gives you the option of eg. adding relative positional encoding as $$Softmax\left(\frac{QK^T + E_{position}}{\sqrt{d}}\right)V$$
As I would have to write custom fused kernels for that to work

@younesbelkada 

Are we also considering multimodal (i.e., `IDEFICS`) and VL models (i.e., `Nougat`) for this issue?

Also, would be interested in learning about and contributing to any larger efforts within HF on integrating optimized components across all libraries (`transformers`, `optimum`, `accelerate`, etc.) -- i.e., designing a more modular system / API for integrating emerging kernels / ops from `xformers`, `triton`, and any other reasonably mature external libraries.

> OK perfect, I will assign you to MPT ! Feel free to let me know if you need any help or if you have any question, as a starting point, I would recommend to have a look at #25598 and see if you can replicate the PR for MPT. For running flash attention tests you can just run (once PR is ready):
> 
> ```shell
> RUN_SLOW=1 pytest -m flash_attn_test tests/models/mpt/
> ```

@younesbelkada I need to add support for MPT only then why replicating the PR?

> ### Feature request
> Flash Attention 2 is a library that provides attention operation kernels for faster and more memory efficient inference and training: https://github.com/Dao-AILab/flash-attention

I was also thinking of adding `FA 2` in `Nought`, 

I would also like to give a model a shot. Since all of the models from the list are taken, are there any recommendations for additional models for which this would be beneficial to implement? @younesbelkada 
Maybe ALBERT?

> I would also like to give a model a shot. Since all of the models from the list are taken, are there any recommendations for additional models for which this would be beneficial to implement? @younesbelkada Maybe ALBERT?

Would you like to work on with me on on adding `FA 2` for `MPT`> @Rubikalubi .here is the PR: #26471 

> Would you like to work on with me on on adding `FA 2` for `MPT`> @Rubikalubi .here is the PR: #26471

Yeah sure @rajveer43 just let me know how i can help.


Hi, sure yes, you can take MPT ! Let us know when the PR is ready for review

@younesbelkada I am not able to provide time to `MPT` you can assign it to someone else.

Hey, @rajveer43  I am taking this up for `MPT` then.

@susnato i already started working on MPT last week (see comment above), but hadn't had the time to finish it yet. If you want, we can work on it together

@marcasty , are you tackling T5? (🙏)

@Taytay had to shelf it this week, someone else can take over if they have the time 

More top chinese models? like https://github.com/OpenLLMAI/OpenRLHF/issues/116

@younesbelkada 

Thoughts on prior [comment](https://github.com/huggingface/transformers/issues/26350?utm_source=tldrai#issuecomment-1745613947)?  Also Persimmon [PR](https://github.com/huggingface/transformers/pull/26482) awaiting review.  

:100: on the presentation on PEFT at PyTorch conf  

Thanks a lot @jeromeku ! I just left some comments on your PR :) And thanks for your comment wrt PT conf!
Regarding the other comment, what would you have exactly in mind? Happy to brainstorm a bit!

Thanks @marcasty.
Anyone have the skillz/desire to add FA2 to T5?

> Thanks @marcasty. Anyone have the skillz/desire to add FA2 to T5?

I can give it a crack 

Hi @younesbelkada. If this issue is still open for DistilBert, I would like to take it up. I have a compatible GPU :) 

Thanks @strankid ! OK yes that would be awesome!
EDIT: indeed it has been already added, see @fxmarty message below

@strankid distilbert support was added in this PR: https://github.com/huggingface/transformers/pull/26489. I updated the list in the first post.

@younesbelkada @fxmarty in that case I could also work on BERT if FA2 hasn't been integrated there yet. Seems to be open in the list. 

Hi @strankid BERT is also WIP currently by another PR 

@younesbelkada If any integration is pending, I would like to take it up. Thanks

@younesbelkada if there are any integrations left, I'd like to contribute! Maybe T5?

> Is it possible to add FlashAttention2 to GPT2 models?

i use class GPTNeoXFlashAttention2 apply for GPT2FlashAttention2  and this work well

> @susnato i already started working on MPT last week (see comment above), but hadn't had the time to finish it yet. If you want, we can work on it together

Any idea when it will be available for MPT?

Did you make any progress on MPT? @Rubikalubi 

clip?

> Thanks @marcasty. Anyone have the skillz/desire to add FA2 to T5?

does someone know if T5 models are blocked by bias support in flash attention ?
https://github.com/Dao-AILab/flash-attention/pull/617

> does someone know if T5 models are blocked by bias support in flash attention ?
> [Dao-AILab/flash-attention#617](https://github.com/Dao-AILab/flash-attention/pull/617)

Yes, you do need attention biases (and their gradients if you want to train/finetune) for the T5, which is exactly why I made this PR in FA2 (I have a custom optimized implementation of T5 using this patch and few other tricks). Once this is supported, the implementation should be quite trivial.

Thank you so much @b-albar , I see that you did already most of the implementation, kudos !! 
can we help you in any way?

Hi @younesbelkada, I'd like to inquire if I may proceed with work on GPT-J. Actually, I already have a draft in progress. 

Moreover, I'm wondering if we should implement the tests accordingly on our own once the implementation is done.

I also confirm that I have access to GPUs compatible with FA2.


Hi @sorenmc, are you still working on a PR for BERT?

> Hi @sorenmc, are you still working on a PR for BERT?

No, you can take it if you want to. The issue is that you have to change the same code several places 

> Moreover, I'm wondering if we should implement the tests accordingly on our own once the implementation is done.

The tests associated with flash attention2 can be found at `tests/test_modeling_common.py`, decorated with `mark.flash_attn_test`.




Hi @bytebarde thanks for working on GPT-J, I saw that you have opened a PR for it, I will have a look !

@b-albar, I'm intrigued by your statement: 
> (I have a custom optimized implementation of T5 using this patch and few other tricks). 

That's not open source by any chance, is it? 🥺 

Any chance of adding support for ESM? (PRoteins)?

It's just that flash attention with bias is not part of the kernels that were develop. Not sure about the status but neither flash nor sdpa support this 🤗 

I think FA authors need to merge: https://github.com/Dao-AILab/flash-attention/pull/617 to add that support

Hi @younesbelkada 👋🏻 @sorenmc is not working anymore on BERT.

> > Hi @sorenmc, are you still working on a PR for BERT?
> 
> No, you can take it if you want to. The issue is that you have to change the same code several places

I would be happy to **try** to work on this but I might need some guidance. I think that @strankid also previously mentioned that he might be willing to help with BERT too.

> @b-albar, I'm intrigued by your statement:
> 
> > (I have a custom optimized implementation of T5 using this patch and few other tricks).
> 
> That's not open source by any chance, is it? 🥺

@Taytay It is now : https://github.com/catie-aq/flashT5
I mention it here because if the PR with the biases for FA2 finally gets merged :crossed_fingers: , it can probably help to make the PR for the T5 here.

Working on `GPT2` here is the PR https://github.com/huggingface/transformers/pull/29226

Hello @younesbelkada i can take on GPT-2 if no one else taken it?

> Hello @younesbelkada i can take on GPT-2 if no one else taken it?

There's already an open PR here https://github.com/huggingface/transformers/pull/29226

Hey thank you for letting me know on the existing PR @EduardoPach , @younesbelkada  & @ArthurZucker are there any available architectures left to implement? I see there are "many more" that are not included on this list. Can you give me an example of one?

It seems that T5 is still open?

@jprivera44 yes it seems so, cannot find any FA/sdpa version. Would be great if you could get that working. I am also looking at integrating sdpa into ESM @ddofer or others would love your help if there is interest!

Hello! My classmate @DELTA-DoubleWise and I are trying to write a project proposal for a course and we would like to extend this to RAG model (which is not listed above). Would you mind assigning this to us?
Thank you very much! @EduardoPach, @younesbelkada  & @ArthurZucker

@rubenweitzman  I wish I could help, but I'm only familiar with keras, not pytorch :\ 

Hey @William-WANG2, we usually don't assign issues, and rather let the code talk: if a PR is open and pinned then that means someone is working on something and the entire community can check the progress 😉 we try to prevent work duplication with this!

Does adding FA2 to CLIP make any sense? 

@sayakpaul Yes! It would be great to add - had a draft in #27444 but was hitting some large differences on just some of the elements which I wasn't able to track. I don't have bandwidth atm so very happy for anyone to take it up! 

Integration question:
I'd love to see this implementation be available for TPU backends. Is there anything we can leverage from this line of work (out of the box), or shall TorchXLA focus on a Pallas implementation under `optimum-tpu`?

cc @philschmid @alanwaketan @shauheen @allenwang28 

Adding @mfuntowicz @tengomucho to answer this

@miladm for now there is a check on models to see if flash attention is available, but `flash_attn` is implemented in cuda. To allow it to run smoothly on TPU, IMO the cleanest option would be to contribute to `flash_attn` to implement a Torch XLA (or pallas) alternative when available. Another alternative would be to implement it directly on `optimum-tpu`, but that would mean we would need to patch models to use that. If you choose the second path, I will be happy to help you to integrate your contribution.

For anyone who is interested in optimized T5 version, I just finished my project on creating flash attention version with fused attention bias calculation. It allows to fix the major drawbacks of T5 and allow to run it on 100k sequences on single L4 GPU (22.5 GB). Check it  [here](https://github.com/Knowledgator/TurboT5). 

Can we update the current state of the list of models in the issue description? For instance, `GPT2` has already been merged

Hi @LysandreJik @ArthurZucker @amyeroberts @younesbelkada @fxmarty @SunMarc @pacman100 ,

I want to *try* to add flash attention 2 to xlmr-large.

Do you have any guidelines?

@michaelshekasta I'd recommend referring to other PRs which have added this for models e.g. [for GPT2](https://github.com/huggingface/transformers/pull/29226) and [reading the contributing guidelines](https://github.com/huggingface/transformers/blob/0ba15cedbcbe9a9db2d5cde76ea4bb712c311934/CONTRIBUTING.md#L4)