[FA-2] Add Flash Attention to Phi

[susnato]

What does this PR do?

This PR adds Flash Attention to Phi.

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.

cc: @younesbelkada, @ArthurZucker

[susnato]

All the FA tests pass except the test_flash_attn_2_generate_padding_right.

This is odd given that the flash_attn_2_inference_padding_right test is passing as does the test_flash_attn_2_generate_left_padding test.

[younesbelkada]

@susnato can you try to run that test multiple times? sometimes it is flaky - apart from that the changes look great on my end !

[susnato]

Hi @younesbelkada, I ran that test 30 times and every time it failed!

Shouldn't the inference test fail too, if the generation test fails? 😅

[younesbelkada]

Hmm yes correct, what I did for llama was to overwrite the test as can be see here: https://github.com/huggingface/transformers/blob/main/tests/models/llama/test_modeling_llama.py#L392 using a real checkpoint. It would be great if you can do the same and test the next 10 tokens are the same (make sure to use do_sample=False)

[susnato]

Hi @younesbelkada, thanks a lot for the advice! All the flash attention tests are passing now. 🤗

[younesbelkada]

Truly amazing work @susnato ! Thanks a lot for this great contribution

[younesbelkada]

Suggested change

	# in fp32. (LlamaRMSNorm handles it correctly)
	# in fp32. (PhiRMSNorm handles it correctly)

[susnato]

Hey, we don't have a PhiRMSNorm, only nn.LayerNorm in the Attention Layer. So removing this part of the line (PhiRMSNorm handles it correctly).

[younesbelkada]

ok makes sense!

[ArthurZucker]

LGTM other than the small comments

[ArthurZucker]

that's not something we want to remove, we should also have flash attention support in persimmon should be the same as this

[susnato]

Actually, there is another PR which is adding the FA support for Persimmon.

Should I add the self.causal=True in the PersimmonAttention so that we can keep this # Copied from statement?

[ArthurZucker]

Yeah! 🤗

[susnato]

Done @ArthurZucker.

[younesbelkada]

Thanks ! We should be able to merge after accepting the suggestion below

[younesbelkada]

Suggested change

	### Expected speedups
	Below is an expected speedup diagram that compares pure inference time between the native implementation in transformers using `susnato/phi-1_dev` checkpoint and the Flash Attention 2 version of the model using a sequence length of 2048.
	<div style="text-align: center">
	<img src="https://huggingface.co/datasets/ybelkada/documentation-images/resolve/main/phi_1_speedup_plot.jpg">
	</div>

[susnato]

Hi @younesbelkada, I have pushed the suggestion you asked.

[susnato]

BTW when is the next release date for transformers?

[younesbelkada]

Hi @susnato sorry one last thing before merging, can you apply changes similar to this commit: 93fe356 and merge with upstream main branch?

[susnato]

Just force-pushed the branch along with the changes. @younesbelkada

[younesbelkada]

Thanks !
We'll probably be able to make it for the next release wdyt @ArthurZucker

[younesbelkada]

Nice!

[ArthurZucker]

LGTM cc @fxmarty for your sdpa PR that will need rebasing I think

[fxmarty]

yep