feat: add flash_attn 2 to bert

[chiennv2000]

Feat: Add flash attention option for BERT
Usage:
model = BertModel.from_pretrained('bert-base-uncased', torch_dtype=torch.bfloat16, use_flash_attention_2=True)

• 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?

@ArthurZucker and @younesbelkada

[younesbelkada]

Thanks a lot for your PR! In principle this looks great!
Many architecture uses BertAttention with # Copied from, therefore all these architectures could benefit from FA-2 for free, however you will need to apply _supports_flash_attn_2 = True on all these architectures. You need to
1- run make fix-copies
2- on modified architectures add the flag above + copy paste the BertFlashAttention class on all on them (with modified names).
Would you be happy to address these changes? Otherwise happy to help you!

[chiennv2000]

Thanks a lot for your review and your suggestions @younesbelkada.
But I don't really familiar with make fix-copies command. Can you guide me on how to do that?

[chiennv2000]

2. I appreciate your feedback. I'm happy to receive your assistance in implementing these changes.
   If you could help me with other architectures, that would be fantastic. Additionally, I'm open to collaborating on extending this to the Roberta and XLMR model. @younesbelkada

[younesbelkada]

Perfect thanks!
As a first step, can you simply run make fix-copies and push the changes here? Then we'll take it over from there !

[chiennv2000]

Thanks @younesbelkada , I did it

[ArthurZucker]

cc @younesbelkada

[younesbelkada]

Didn't had time to properly look into it, will do it asap!

[kevinhu]

Any updates on getting this PR merged?

[hackyon]

Hello there!

I'm working on integrating scaled_dot_product_attention to BERT #28802, and there might be some merge conflicts with this change. Mostly, if my changes go through, then we can get rid of most of the downstream dependencies from fix-copies.

Let me know if you have any questions. Happy to discuss and/or chat on the best way forward if necessary.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[loswald]

This would be a lifesaver for me, I hope merging this is prioritized! cc @younesbelkada
@chiennv2000 what kind of speedups are you observing with this?

[hackyon]

This would be a lifesaver for me, I hope merging this is prioritized! cc @younesbelkada @chiennv2000 what kind of speedups are you observing with this?

@loswald - you can see a quick estimate of the speedups in #28802. The pytorch SDPA implementation uses FA2 under the hood (if your hardware supports it). The PR is ready but we're just waiting on the HG team to merge it.