[ESM] Add flash-attention-2 backend for ESM-2
Author: pstjohn
Pull request: yes

# What does this PR do?

Implements the flash-attention-2 attention backend for ESM-2. 

ESM-2 is a BERT-based protein language model from facebook research. While the model is undoubtedly a bit old at this point - there are now ESM-3 and ESM-C models available, they remain very popular on the huggingface hub with ESM-2 3B having nearly 6M downloads last month. Despite this, there are no efficient attention backends for ESM-2 available via huggingface. There have been various drop-in replacements for ESM-2 that have been created to fill this need ([FastESM](https://huggingface.co/Synthyra/FastESM2_650), [FAPLM](https://github.com/pengzhangzhi/faplm)), but a native transformers implementation should speed up a lot of workflows that currently rely on the native pytorch attention.

Addresses #26350 

Related to this PR; it would be great if we could merge some of the remaining automated `safetensor` conversions for the ESM models on huggingface hub, I noticed @Rocketknight1 merged some for the other model scales previously.

* [ESM-2 3B](https://huggingface.co/facebook/esm2_t36_3B_UR50D/discussions/18)
* [ESM-2 15B](https://huggingface.co/facebook/esm2_t48_15B_UR50D/discussions/1)




## 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.
- [ ] 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).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- 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
- vision models: @amyeroberts, @qubvel
- speech models: @eustlb
- graph models: @clefourrier

Library:

- flax: @gante and @Rocketknight1
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Rocketknight1
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @zach-huggingface and @SunMarc
- chat templates: @Rocketknight1

Integrations:

- deepspeed: HF Trainer/Accelerate: @SunMarc @zach-huggingface
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc @MekkCyber

Documentation: @stevhliu

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: @Rocketknight1
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the `Ready for review` button (at the bottom of the PR page). This will assign reviewers and trigger CI.

Tagging @ArthurZucker and @ydshieh for review. Thanks!

@Rocketknight1, @ArthurZucker; gentle ping for review. Thanks!

Sorry for the delay! The changes look good, but can you add an equivalence test like `test_flash_attn_2_equivalence`? It's fine to copy it from another model

Added! Thanks!
Also while I have your attention I just wanted to echo the request above to merge safetensors versions for the two largest ESM-2 models (unless there's a reason to omit them I'm not aware of); it should make loading these bigger models faster:
 * https://huggingface.co/facebook/esm2_t36_3B_UR50D/discussions/18
 * https://huggingface.co/facebook/esm2_t48_15B_UR50D/discussions/1

Unfortunately, those are Meta's checkpoints! We generally don't intervene in user or organization repos unless it's absolutely critical, but you can try pinging them there?

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

> Unfortunately, those are Meta's checkpoints! We generally don't intervene in user or organization repos unless it's absolutely critical, but you can try pinging them there?

Right, that's fair. The original authors don't work there anymore though, so I expect those are abandoned 😀. I only asked because I noticed you merged similar PRs for the smaller models sizes