Skip to content

[ESM] Add flash-attention-2 backend for ESM-2 #38023

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 4 commits into from
May 16, 2025
Merged

[ESM] Add flash-attention-2 backend for ESM-2 #38023

merged 4 commits into from
May 16, 2025
+188 −8

Conversation

pstjohn
Copy link
Contributor

@pstjohn pstjohn commented May 8, 2025 •
edited
Loading

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, 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.

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • Did you read the contributor guideline,
    Pull Request section?
  • Was this discussed/approved via a Github issue or the forum? 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, and
    here are tips on formatting docstrings.
  • 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.

@pstjohn
Add flash-attention-2 backend for ESM-2
0d4ec1e 
Signed-off-by: Peter St. John <pstjohn@nvidia.com>
@github-actions github-actions bot marked this pull request as draft May 8, 2025 14:16
@github-actions GitHub Actions
Copy link
Contributor

github-actions bot commented May 8, 2025

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.

@pstjohn pstjohn marked this pull request as ready for review May 8, 2025 14:16
@github-actions github-actions bot requested review from ArthurZucker and ydshieh May 8, 2025 14:17
@pstjohn
Copy link
Contributor Author

pstjohn commented May 9, 2025

Tagging @ArthurZucker and @ydshieh for review. Thanks!

@pstjohn
update extended_attention_mask for fa2
88149df 
Signed-off-by: Peter St. John <pstjohn@nvidia.com>
@pstjohn pstjohn mentioned this pull request May 13, 2025
Open
5 tasks
@pstjohn pstjohn changed the title Add flash-attention-2 backend for ESM-2 [ESM] Add flash-attention-2 backend for ESM-2 May 13, 2025
@pstjohn
Copy link
Contributor Author

pstjohn commented May 14, 2025

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

@Rocketknight1
Copy link
Member

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

@pstjohn
add test_flash_attn_2_equivalence test
4c5bddb 
Signed-off-by: Peter St. John <pstjohn@nvidia.com>
@pstjohn
Copy link
Contributor Author

pstjohn commented May 15, 2025

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:

@Rocketknight1
Copy link
Member

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?

Rocketknight1
Rocketknight1 approved these changes May 16, 2025
View reviewed changes
Copy link
Member

@Rocketknight1 Rocketknight1 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

PR looks good to me now, though, so merging!

@Rocketknight1 Rocketknight1 merged commit d69945e into huggingface:main May 16, 2025
14 checks passed
@HuggingFaceDocBuilderDev
Copy link

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@pstjohn
Copy link
Contributor Author

pstjohn commented May 16, 2025

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@Rocketknight1 Rocketknight1 Rocketknight1 approved these changes

@ArthurZucker ArthurZucker Awaiting requested review from ArthurZucker

@ydshieh ydshieh Awaiting requested review from ydshieh

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@pstjohn @Rocketknight1 @HuggingFaceDocBuilderDev
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy