Add `sdpa` and  FA2 for CLIP
Author: qubvel
Pull request: yes
Labels: name: Vision

name: run-slow

# What does this PR do?

Add flash attention 2 and sdpa (`torch.nn.functional.scaled_dot_product_attention`) attention implementations for CLIP model.

This is a continuation of work done in:
 - #30390 by @sayakpaul 
 - #31499

## 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_31940). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Benchmarking results for the CLIP models (`"openai/clip-vit-large-patch14"` checkpoint) in `float16` with the following env:

```
- `transformers` version: 4.43.0.dev0
- Platform: Linux-6.5.0-1020-aws-x86_64-with-glibc2.35
- Python version: 3.10.12
- PyTorch version (GPU?): 2.3.1+cu121 (True)
- GPU type: NVIDIA A10G
```

Code:
https://gist.github.com/qubvel/ac691a54e54f9fae8144275f866a7ff8

### CLIPModel

|   Image batch size |   Num text labels |   Eager (s/iter) |   FA2 (s/iter) |   FA2 speedup |   SDPA (s/iter) |   SDPA speedup |
|-------------------:|------------------:|-----------------:|---------------:|--------------:|----------------:|---------------:|
|                  1 |                 4 |            0.025 |          0.026 |         0.954 |           0.02  |          1.217 |
|                  1 |                16 |            0.026 |          0.028 |         0.918 |           0.02  |          1.287 |
|                  1 |                32 |            0.026 |          0.034 |         0.767 |           0.022 |          1.173 |
|                  1 |                64 |            0.042 |          0.046 |         0.906 |           0.036 |          1.167 |
|                  4 |                 4 |            0.028 |          0.033 |         0.849 |           0.024 |          1.189 |
|                  4 |                16 |            0.034 |          0.035 |         0.955 |           0.029 |          1.169 |
|                  4 |                32 |            0.043 |          0.042 |         1.025 |           0.036 |          1.175 |
|                  4 |                64 |            0.059 |          0.055 |         1.072 |           0.05  |          1.179 |
|                 16 |                 4 |            0.096 |          0.088 |         1.091 |           0.078 |          1.234 |
|                 16 |                16 |            0.102 |          0.09  |         1.129 |           0.083 |          1.224 |
|                 16 |                32 |            0.111 |          0.096 |         1.15  |           0.091 |          1.222 |
|                 16 |                64 |            0.127 |          0.11  |         1.157 |           0.105 |          1.218 |
|                 32 |                 4 |            0.185 |          0.159 |         1.157 |           0.149 |          1.238 |
|                 32 |                16 |            0.19  |          0.162 |         1.177 |           0.154 |          1.233 |
|                 32 |                32 |            0.199 |          0.168 |         1.188 |           0.162 |          1.231 |
|                 32 |                64 |            0.216 |          0.181 |         1.19  |           0.176 |          1.228 |

### CLIPTextModel

|   Num text labels |   Eager (s/iter) |   FA2 (s/iter) |   FA2 speedup |   SDPA (s/iter) |   SDPA speedup |
|------------------:|-----------------:|---------------:|--------------:|----------------:|---------------:|
|                 4 |            0.009 |          0.012 |         0.737 |           0.007 |          1.269 |
|                16 |            0.009 |          0.014 |         0.659 |           0.008 |          1.187 |
|                32 |            0.018 |          0.021 |         0.862 |           0.016 |          1.142 |
|                64 |            0.034 |          0.034 |         1.001 |           0.03  |          1.163 |
|               128 |            0.063 |          0.058 |         1.09  |           0.054 |          1.174 |

![clip_text_model_viz_3](https://github.com/user-attachments/assets/e9826b43-4e66-4f4c-952b-af4d90bd38eb)

### CLIPVisionModel

|   Image batch size |   Eager (s/iter) |   FA2 (s/iter) |   FA2 speedup |   SDPA (s/iter) |   SDPA speedup |
|-------------------:|-----------------:|---------------:|--------------:|----------------:|---------------:|
|                  1 |            0.016 |          0.013 |         1.247 |           0.012 |          1.318 |
|                  4 |            0.025 |          0.021 |         1.198 |           0.021 |          1.202 |
|                 16 |            0.093 |          0.075 |         1.234 |           0.075 |          1.24  |
|                 32 |            0.181 |          0.147 |         1.237 |           0.146 |          1.241 |

![clip_image_model_viz_3](https://github.com/user-attachments/assets/50a36206-e3b9-4adc-ac8e-926b8b071d63)



Note:

CLIPTextModel and CLIPModel as well, are not supposed to be used with left padding:
1) text model pooled output extracts the last relevant token by the index of the first `pad token`
2) with left padding and causal mask we are going to have fully-masked rows, which might lead to `nan` values for `sdpa` with a memory-efficient backend (relevant issues https://github.com/huggingface/transformers/pull/31208, https://github.com/pytorch/pytorch/issues/103749)


## Tests on SSH runner

### Flash attention
<img width="1066" alt="Screenshot 2024-07-15 at 17 32 12" src="https://github.com/user-attachments/assets/8f842ffd-bc1e-4029-9668-a63c011664ac">

### SDPA
<img width="1054" alt="Screenshot 2024-07-15 at 17 30 10" src="https://github.com/user-attachments/assets/456d37f7-caab-4bb5-9307-c0ca78dff0a4">


@sayakpaul in case you have time, can you please review this PR?

Slow test error:
```
FAILED tests/models/clip/test_tokenization_clip.py::CLIPTokenizationTest::test_np_encode_plus_sent_to_model - AttributeError: 'CLIPTokenizationTest' object has no attribute 'skip'
```

addressed in https://github.com/huggingface/transformers/pull/32012

@amyeroberts whenever you have time can you please review this PR 🙂 

Great work, Pavel. I think this is ready to be shipped. 

FYI) @sayakpaul I think you are the managing `diffusers` library so I left some comment here. Various stable diffusion models use CLIP encoder, so this PR change particulary using SDPA attention might be cause error if you update your diffusers library to the most recent one, when we convert this model into ONNX or other platform such as TensorRT.  I think we should encourage diffusers user to use this method `attn_implementation="eager"` in conversion.

Cc @DN6 @yiyixuxu ^

https://github.com/huggingface/transformers/issues/28610