[WIP] Add FA2 for all Bart-like
Author: patrickvonplaten
Pull request: yes

# What does this PR do?

Add FA2 to all Bart-like models

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_26722). All of your documentation changes will be reflected on that endpoint.

Verified that FA2 works by checking Whisper. Bart's attention is exactly the same as Whisper so it should as well. I will run some better benchmarks later.

@ArthurZucker @younesbelkada could you do a first review here, just for the following files:
- `tests/test_modeling_common.py`
- `src/transformers/models/bart/modeling_bart.py`
[ignore Whisper completely now, please]

It  would be nice to agree on these files before running `make fix-copies` which will change 10+ other modeling files.
The implementation was pretty straight-forward as I can more or less copy-paste all the code from Llama (nice job @younesbelkada!)

Some comments:
- 1.) The flash attention tests are very nicely implemented in `tests/test_modeling_common.py`, but it looks like the tolerance is too high to catch any incorrect masking or other settings. E.g. in the beginning for BART I had an incorrect scaling factor in the attention and all tests passed anyways here. We might want to look into this.
- 2.) I'm not super happy about passing around both `attention_mask` and `padding_mask` all the time. This makes the code really difficult to read and is quite confusing (what is the difference between the two?!). As far as I understand it the two masks are the same - the only reason we use `padding_mask` in addition to `attention_mask` is because the attention_mask is expanded and thus can't be used for FA. I wonder whether we should do a bigger refactor here actually and instead of expanding the attention_mask in the beginning we only expand it right before the attention so that we don't have to pass around both `padding_mask` and `attention_mask`. We could even cache the expanded mask if needed for speed. I would be strongly in favor of **not** having both a padding mask and an attention mask. Also cc @fxmarty here. 
- 3.) Can we make the automatic conversion to intended FA precision a bit more robust. E.g. see here: https://github.com/huggingface/transformers/pull/26722/files#r1353462371 . Aren't there use cases where the user would like to train in bfloat16, but might have a layer norm in fp32? cc @younesbelkada 



Let's also make sure to update the `.md` once you have the expected performance gains and benchmark the other models where we copy from (could help debug / make sure everything is working properly).

Thanks for the comments:

> 2- I think that it is the right approach indeed to cache the expanded mask. The attention mask format that FA-2 expects is the raw 'padding mask' --> 1 for attended tokens and 0 for padding tokens. It is possible to convert the expanded mask to padding mask but creates a huge overhead that makes it not interesting for long sequences (with FA-2 being slower than native). However I know caching some tensors as a buffer can cause some weird issues with accelerate, I would make sure that the accelerate tests pass in case we decide to go for that approach. RUN_SLOW=1 pytest -m accelerate_tests tests/models/bart/test_modeling_bart.py

I see, I'm not sure how easy it'll be to get this working in a clean way then

Before continuing with this PR, I'll wait for https://github.com/huggingface/transformers/pull/26560 to be merged and also will see if we find a cleaner solution for 2.) 

OK I'll work on #26560 ASAP and let you know once that PR gets merged

https://github.com/huggingface/transformers/pull/26846 is now merged, let me know if you want me to help you on this PR by updating the necessary changes + benchmarking

**Update**:

The PR now works for Bart. @ArthurZucker @younesbelkada @fxmarty @LysandreJik could you give the design chosen for BART a look here and if ok, I'll apply it to all other Bart-like models.

**Please only review `modeling_bart.py` !!!**