Adding Flash Attention 2 Support for GPT2
Author: EduardoPach
Pull request: yes

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes https://github.com/huggingface/transformers/issues/26350


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

Hey, @younesbelkada added flash attention 2 support for `GPT2`. The only thing missing is the Expected speedups, could you share the code you used for the other models you added support to keep consistency?

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @pacman100

Integrations:

- deepspeed: HF Trainer/Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc and @younesbelkada

Documentation: @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
> Wow thanks for the great work ! At a quick glance it seems you took care very well of the copy mechanism which is quite a challenge for GPT2 ! Please find the benchmarking script: https://gist.github.com/younesbelkada/02f35734da906cc0f2389ae4f665c58f I suggest to try it out for prefill only on large sequence length - let us know with @ArthurZucker @fxmarty how it goes

Hey, I don't have a GPU and I was renting in RunPod an RTX 3090 to work on this PR, is it a problem to use the 3090 to benchmark or should I switch to an A100 (which I believe it was the GPU used in the other benchmarks at least the ones I've seen)?

Thanks @EduardoPach for getting back, I think using a 3090 is fine !

@ArthurZucker I believe it should be ready for review

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_29226). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@EduardoPach thanks again, what @ArthurZucker meant is an integration test similar as: https://github.com/huggingface/transformers/blob/63caa370e6c618dbe7d3fd4cbf545cc32eca1a15/tests/models/llama/test_modeling_llama.py#L400 for GPT2 only, would you be happy to work on that? 🙏 

> @EduardoPach thanks again, what @ArthurZucker meant is an integration test similar as: https://github.com/huggingface/transformers/blob/63caa370e6c618dbe7d3fd4cbf545cc32eca1a15/tests/models/llama/test_modeling_llama.py#L400 for GPT2 only, would you be happy to work on that? 🙏 

Yeah, I will add the test in the following hours

Following hours became more like following days haha, but should be good now @ArthurZucker 

> Thanks for your great work @EduardoPach ! Can you merge this branch with upstream main? after that we can merge i think

Done! 😄 

Thanks so much @EduardoPach ! can you just run `make fix-copies` ! this should fix the failing CI :D 

> Thanks again ! We just merged some fixes on main - could you rebase again 🙏 then we should finally merge :D sorry for all the iterations !

No worries! Done

Thanks ! Hmm I can't see the rebase commit on the history, perhaps can you try again ?

> Thanks ! Hmm I can't see the rebase commit on the history, perhaps can you try again ?

I've done `git fetch upstream && git merge upstream/main` here https://github.com/huggingface/transformers/pull/29226/commits/dcde56c66907cdcd74d011fc7b1919446f8240c3

Thank you for your hard work! Many of us are excited about the GPT-2 model supporting flash attention. May I ask when the PR is expected to be merged?

> Thank you for your hard work! Many of us are excited about the GPT-2 model supporting flash attention. May I ask when the PR is expected to be merged?

Hey, I believe if @amyeroberts agrees with my latest message it should get merged right away 🤞 

It seems everything is okay. May I kindly request to merge this PR? I am really looking forward to speeding up my GPT-2. If my request has added to your workload, I apologize for any inconvenience.

> It seems everything is okay. May I kindly request to merge this PR? I am really looking forward to speeding up my GPT-2. If my request has added to your workload, I apologize for any inconvenience.

c.c. @amyeroberts