Add FA2 and `sdpa` support for SigLIP
Author: qubvel
Pull request: yes
Labels: name: run-slow

# What does this PR do?

Add flash attention 2 and sdpa (`torch.nn.functional.scaled_dot_product_attention`) attention implementations for SigLIP model.

<!--
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 -->

Fixes #31138


## Before submitting
- [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.
- [x] 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).
- [x] 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.


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

Nice! We can probably combine this with CLIP cc @sayakpaul 

For reference, there's a FA2 siglip implementation for IDEFICS2, but I'm not sure how much testing was done for the equivalence between the eager and FA2 classes

@amyeroberts there are some discrepancies with the attention mask, I am digging deeper to the equivalence testing

Running test locally  


### Flash Attention

```bash
RUN_SLOW=1 python -m pytest --verbose -m flash_attn_test \
    tests/models/siglip/test_modeling_siglip.py
```
<img width="1390" alt="Screenshot 2024-06-24 at 15 31 23" src="https://github.com/huggingface/transformers/assets/31920396/2b981867-fac0-4b0a-a9e4-dc73f50001e3">

### SDPA
```bash
RUN_SLOW=1 python -m pytest --verbose \
    tests/models/siglip/test_modeling_siglip.py -k "sdpa"
```
<img width="1387" alt="Screenshot 2024-06-24 at 15 31 06" src="https://github.com/huggingface/transformers/assets/31920396/b3598a8f-782a-41fc-9f28-edd8fcfc1f39">


@amyeroberts @molbap please review if you have time

@amyeroberts 

> Main comment is about the attention selection, we should instead be instantiating the model components with from_config (or possibly _from_config?) and passing in `attn_implementation=config.attn_implementation).

I tried to make it with `_from_config`, however, internal model components are not inherited from `PretrainedModel`, they are just pure `nn.Module` and don't have `_from_config` method. I changed the inheritance, but that led to other issues: models have to be included in docs and some fields have to be specified. 

Not sure we can change internal model components too, for example, change `SiglipVisionTransformer(nn.Module)` to  `SiglipVisionModel(SiglipPreTrainedModel)`. This will lead to incompatible checkpoints.

See implementation https://github.com/huggingface/transformers/pull/31499/commits/669c5378f4758b05345d6faa8fd51328ad57d32c

~~Do you have any thoughts on that?~~

P.S. I am looking at https://github.com/huggingface/transformers/pull/30390 with similar questions 👀 

>Could we extend this to CLIP and add both at the same time?

I hope we can, I will check this!

P.S. Given the work made in https://github.com/huggingface/transformers/pull/30390, probably will be better to continue work on this, rather than merge both PRs

>There should be a section added to the model doc page + benchmarks showing some times for improvements for SDPA and FA2 e.g. [like here for mistral](https://huggingface.co/docs/transformers/v4.41.3/en/model_doc/mistral#expected-speedups).

Addressed in https://github.com/huggingface/transformers/pull/31499/commits/d41955d5ac12a25c3f64d445dab23668e6a2b34f

![speedup_plot](https://github.com/huggingface/transformers/assets/31920396/607c1bb8-eefa-45c4-94d9-6712dbd50aac)


Regarding the CLIP part: https://github.com/huggingface/transformers/pull/30390. Feel free to cherry-pick commits if that help. I got stuck in some FLAX tests that I never got time to get resolved. 

> I tried to make it with _from_config, however, internal model components are not inherited from PretrainedModel, they are just pure nn.Module and don't have _from_config method

Hi, regarding this part, even if an internal component is just `nn.Module`, we can still passing `config` to its `__init__`, just like what is done in `GemmaDecoderLayer`. Hope this could help.

```python
class GemmaDecoderLayer(nn.Module):
    def __init__(self, config: GemmaConfig, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size

        self.self_attn = GEMMA_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)
```

I changed the attention implementation propagation as follows (commit https://github.com/huggingface/transformers/pull/31499/commits/23457a2233bf62cf3c3e7c24d6d0b08f07c78864):

1. Initialize the `*Model` class instead of the `*Transformer` module to utilize the `_from_config` method.
2. Use only the submodule of the `*Model` to maintain the overall model structure and weights for backward compatibility.

```python
# First, initialize the text and vision models with proper attention implementation
text_model = SiglipTextModel._from_config(text_config, attn_implementation=config._attn_implementation)
vision_model = SiglipVisionModel._from_config(vision_config, attn_implementation=config._attn_implementation)

# Second, get the text and vision submodules (for backward compatibility)
self.text_model = text_model.text_model
self.vision_model = vision_model.vision_model
```

With this approach, the underlying modules will exhibit the same behavior for the `attn_implementation` setting.

The disadvantage of this method is that the `post_init()` method is called twice: once for each `*Model` and again for the parent model.

@amyeroberts please let me know what you think.

Thanks for the work! 

Which latest version on pypi would support this feature?

Hi @lucasjinreal, it's going to be included in the next release, most probably 4.43.0. You can try it now by installing transformers from the source

```bash
pip install -U git+https://github.com/huggingface/transformers.git
```

I will waiting for the 4.43, my server unable to access network, any estimate time will 4.43 out?

@lucasjinreal Just to confirm, you can't access github from your server? If it's just relating to internet, you would also need access for installing the latest release from pypi. 

We typically release on a monthly schedule. You can see the list of [releases here](https://github.com/huggingface/transformers/releases). The next minor release will probably be in 2-3 weeks.

thanks, will try update weeks later.