[`FA2`] Add flash attention for opt
Author: susnato
Pull request: yes

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Adds Flash Attention 2 for `OPT` as discussed in in this issue - #26350 .


## 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](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] 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 and @younesbelkada
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @pacman100

Integrations:

- deepspeed: HF Trainer/Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc and @younesbelkada

Documentation: @stevhliu and @MKhalusova

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: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->
cc : @younesbelkada

Comments:
The tests were done using a RTX 3060 (Ampere) which supports Flash Attention 2.
![Screenshot from 2023-09-26 16-23-55](https://github.com/huggingface/transformers/assets/56069179/58562bc0-37e7-45f5-a60c-6dc816cab0da)


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

Hi @younesbelkada, I have added the model to the list. Also I have merged the [susnato#3](https://github.com/susnato/transformers/pull/3), please let me know if anymore changes are needed or not.

Hi @younesbelkada , done.

Thanks for the reply @younesbelkada, in meantime I would like to work on adding flash attention for another model.

Thanks @susnato , sure yes ok, let me know on #26350 which model would you take, perhaps you can try your hands on Starcoder if you are interested. Let me know!

yes I am interested in that! @younesbelkada 

Thanks very much ! Looking forward to your PR ! 

Also please assign me in the list :). @younesbelkada 

Yes, just assigned you on the list !

@susnato @younesbelkada why not use `F.scaled_dot_product_attention` instead as it support Flash-Attention-2 and able to use with `torch.compile` ?. I used to use the `flash-attention` on official repo like this PR but now I moved to the `F.scaled_dot_product_attention`.

Hi @susnato 
Thanks for your patience,
I had a look at the issue with @michaelbenayoun and the solution is the following: 

1- define a method on the top level of the modeling opt file that checks if there is any padding token inside the attention mask
```python

if is_torch_fx_available():
    @torch.fx.wrap
    def check_padding_in_attention_mask(attention_mask):
        if 0 in attention_mask:
            return attention_mask
        return None
else:
    def check_padding_in_attention_mask(attention_mask):
        if 0 in attention_mask:
            return attention_mask
        return None
```
And add the `@torch.fx.wrap` decorator in the method to make sure it is compatible with the PT versions we support. 
2- revert until the commit [689f599](https://github.com/huggingface/transformers/pull/26414/commits/689f599132f121a3716cf4497570cbbda5c01737) and replace the simple logic by a call to that method 

That way FX tests should hopefully pass

I will later take care of moving `check_padding_in_attention_mask` in `pytorch_utils` file so that models that support FX tracing can use that method for future integrations.



Hi @dathudeptrai, `transformers` maintain a single modeling file policy, where we want to make sure that someone who is reading the code is properly able to understand the modeling code without having to switch to any other file. That's why even the commonly used `dot product attention` which is used by many modes are defined for each one. We define `F.scaled_dot_product_attention` so that the user does not have to visit the torch website to understand the implementation. 

I believe that in future when flash attention becomes more stable the maintainers will try to implement the function within the library itself(given that it's possible) to reduce the dependency. 

Also @younesbelkada could give you a better explanation and correct me if I am wrong here :). 

Hi @younesbelkada, I have added the solution that you suggested and it's working!
Let me know if any more changes are required or not.

@dathudeptrai Thanks for your question
We are currently thinking of migrating SDPA natively into transformers in the next weeks, we will keep you posted with @fxmarty 

Hi @younesbelkada the `test_flash_attn_2_generate_use_cache` is failing for `llama`, `falcon` and `opt` models :confused: .

![Screenshot from 2023-09-28 15-51-01](https://github.com/huggingface/transformers/assets/56069179/71bc2ba4-73df-4a6e-b21f-8c646e025b09)
![Screenshot from 2023-09-28 15-52-25](https://github.com/huggingface/transformers/assets/56069179/19145257-f713-4844-925e-126884b163cd)
![Screenshot from 2023-09-28 15-54-18](https://github.com/huggingface/transformers/assets/56069179/b88a6e58-51c1-4aa8-9592-2f09f329da8e)


Trying to figure out the reason.



Hi @susnato 
The test seem to pass on my end

![Screenshot 2023-09-28 at 12 28 01](https://github.com/huggingface/transformers/assets/49240599/63150845-c7ed-4898-9a70-64cd46bd4db2)

Can you try to merge with main and re-run the tests again?

@younesbelkada I already did :(.

@susnato can you share the full traceback?

Hi @younesbelkada ,

<details>
<summary>Full Traceback for Llama</summary>


self = <tests.models.llama.test_modeling_llama.LlamaModelTest testMethod=test_flash_attn_2_generate_use_cache>

    @require_flash_attn
    @require_torch_gpu
    @mark.flash_attn_test
    @slow
    def test_flash_attn_2_generate_use_cache(self):
        import torch
    
        for model_class in self.all_generative_model_classes:
            if not model_class._supports_flash_attn_2:
                return
    
            config, _ = self.model_tester.prepare_config_and_inputs_for_common()
            model = model_class(config)
    
            with tempfile.TemporaryDirectory() as tmpdirname:
                model.save_pretrained(tmpdirname)
    
                dummy_input = torch.LongTensor([[0, 2, 3, 4], [0, 2, 3, 4]]).to(torch_device)
                dummy_attention_mask = torch.LongTensor([[1, 1, 1, 1], [1, 1, 1, 0]]).to(torch_device)
    
                model = model_class.from_pretrained(
                    tmpdirname, torch_dtype=torch.float16, use_flash_attention_2=True, low_cpu_mem_usage=True
                ).to(torch_device)
    
                # Just test that a large cache works as expected
>               _ = model.generate(
                    dummy_input, attention_mask=dummy_attention_mask, max_new_tokens=30, do_sample=False
                )

tests/test_modeling_common.py:2936: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
../../anaconda3/envs/transformers/lib/python3.9/site-packages/torch/autograd/grad_mode.py:27: in decorate_context
    return func(*args, **kwargs)
src/transformers/generation/utils.py:1606: in generate
    return self.greedy_search(
src/transformers/generation/utils.py:2454: in greedy_search
    outputs = self(
../../anaconda3/envs/transformers/lib/python3.9/site-packages/torch/nn/modules/module.py:1194: in _call_impl
    return forward_call(*input, **kwargs)
src/transformers/models/llama/modeling_llama.py:1034: in forward
    outputs = self.model(
../../anaconda3/envs/transformers/lib/python3.9/site-packages/torch/nn/modules/module.py:1194: in _call_impl
    return forward_call(*input, **kwargs)
src/transformers/models/llama/modeling_llama.py:921: in forward
    layer_outputs = decoder_layer(
../../anaconda3/envs/transformers/lib/python3.9/site-packages/torch/nn/modules/module.py:1194: in _call_impl
    return forward_call(*input, **kwargs)
src/transformers/models/llama/modeling_llama.py:631: in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
../../anaconda3/envs/transformers/lib/python3.9/site-packages/torch/nn/modules/module.py:1194: in _call_impl
    return forward_call(*input, **kwargs)
src/transformers/models/llama/modeling_llama.py:489: in forward
    attn_output = self._flash_attention_forward(
src/transformers/models/llama/modeling_llama.py:546: in _flash_attention_forward
    attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
../../anaconda3/envs/transformers/lib/python3.9/site-packages/flash_attn/bert_padding.py:208: in pad_input
    output = index_put_first_axis(hidden_states, indices, batch * seqlen)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

ctx = <torch.autograd.function.IndexPutFirstAxisBackward object at 0x7f1fb51b9040>
values = tensor([[[-2.3666e-02,  2.1423e-02, -7.4646e-02,  3.6530e-02, -7.8979e-02,
          -3.4546e-02,  1.0236e-01,  4.3549...4504e-03, -9.3384e-02,
          -4.5532e-02, -5.5847e-02,  4.0253e-02]]], device='cuda:0',
       dtype=torch.float16)
indices = tensor([0, 1], device='cuda:0', dtype=torch.int32), first_axis_dim = 2

    @staticmethod
    def forward(ctx, values, indices, first_axis_dim):
        ctx.save_for_backward(indices)
        assert indices.ndim == 1
        assert values.ndim >= 2
        output = torch.zeros(
            first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype
        )
        # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
>       output[indices] = values
E       IndexError: tensors used as indices must be long, byte or bool tensors

../../anaconda3/envs/transformers/lib/python3.9/site-packages/flash_attn/bert_padding.py:51: IndexError


</details>

BTW I updated the flash attention library and merged to main but the error is still there.

I have pushed the docstring change.

If you don't mind @younesbelkada , could you please checkout this branch and run the tests on your end and let me know the results?
If this passes on your machine, then It could be due to some unknown error on my end.

Sure @susnato no problem, will run the tests locally tomorrow and let you know how it goes

Hi @younesbelkada, did you manage to check if the tests run successfully on your local machine? 

Yes the tests seemed to pass the time I ran them, however since #26792 being merged you need to remove the support for `padding_mask` and follow whqat has been done in that PR. Let me know if you need help on this!

After I finish updating the `gpt_bigcode` I will take this up.

OK awesome, thanks @susnato !

Hi @younesbelkada, I have pushed the changes.

All the flash attention tests are running fine except `test_flash_attn_2_generate_use_cache` but I believe it's due to some problem with my local installation since the model generates outputs quite well when I run separately.


Hello @younesbelkada , are you going to check the speed-up for this model too?  

I am asking this out of curiosity - Are you going to check speed-ups for every model, FlashAttention is added to or is this benchmark needed for popular ones only? 

Hi @amyeroberts, I have pushed the changes you suggested, in addition to that I have also updated the model_doc file and added speedup graphs. 

Yes, I just checked and all tests are passing! @younesbelkada!

Flash Attention tests - 

![Screenshot from 2023-11-23 01-04-43](https://github.com/huggingface/transformers/assets/56069179/d8be700e-46e6-409c-9232-c32bd55278b4)

I found only one Integration test[`pt`] related to OPT - 

![Screenshot from 2023-11-23 01-09-20](https://github.com/huggingface/transformers/assets/56069179/63978a47-a7e4-41ab-87e8-9d92834cff60)