[CLIP] add: sdpa support to clip.
Author: sayakpaul
Pull request: yes

# What does this PR do?

CLIP is heavily used in the diffusion modeling space for text encoding. So, having native SDPA support for CLIP would be beneficial for diffusion models both for training and inference. 

The test failures can be tackled later once we match the logits to non-SDPA. Here's my test script:

```python
from transformers import AutoTokenizer, CLIPTextModel
import torch

model_sdpa = CLIPTextModel.from_pretrained("openai/clip-vit-base-patch32", attn_implementation="sdpa").to("cuda")
tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32")

inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt")
inputs = {k: v.to("cuda") for k, v in inputs.items()}

with torch.no_grad():
    outputs = model_sdpa(**inputs)
    last_hidden_state_sdpa = outputs.last_hidden_state

model = CLIPTextModel.from_pretrained("openai/clip-vit-base-patch32", attn_implementation="eager").to("cuda")

with torch.no_grad():
    outputs = model(**inputs)
    last_hidden_state = outputs.last_hidden_state

print(last_hidden_state_sdpa[0, :3, -1].flatten())
print(last_hidden_state[0, :3, -1].flatten())
print(torch.allclose(last_hidden_state_sdpa, last_hidden_state, rtol=1e-2, atol=1e-2))
```

These don't pass:

```bash
tensor([ 0.1013,  0.0805, -0.8429], device='cuda:0')
tensor([ 0.1013, -0.9801, -1.9444], device='cuda:0')
False
```

I inspected the `key_states`, `value_states`, and the `query_states` in `CLIPAttention` and `CLIPSdpaAttention`, respectively. In both classes, these matrices are of same value. The differences start arising from `attn_output`. 

```bash
SDPA attn_output[0, :3, -1]=tensor([0.3397, 0.2862, 0.5571], device='cuda:0')
SDPA attn_output[0, :3, -1]=tensor([0.0143, 0.0158, 0.0971], device='cuda:0')
SDPA attn_output[0, :3, -1]=tensor([0.0082, 0.0130, 0.1147], device='cuda:0')

Non SDPA attn_output[0, :3, -1]=tensor([0.3397, 0.2576, 0.7800], device='cuda:0')
Non SDPA attn_output[0, :3, -1]=tensor([ 0.0143,  0.0069, -0.0258], device='cuda:0')
Non SDPA attn_output[0, :3, -1]=tensor([ 0.0082,  0.0046, -0.0023], device='cuda:0')
```

I suspect this is happening because of how masking is handled in `CLIPAttention`. We first apply causal mask:

https://github.com/huggingface/transformers/blob/8b02bb6e748842cfb74ec9cba808201b893fa0af/src/transformers/models/clip/modeling_clip.py#L285

And then we apply attention mask:

https://github.com/huggingface/transformers/blob/8b02bb6e748842cfb74ec9cba808201b893fa0af/src/transformers/models/clip/modeling_clip.py#L295

But I think it deviates a bit from `CLIPSdpaAttention`. Not very sure, though. Hence I am opening this PR for seeking feedback. 

I have added some comments in line to provide further clarification on some points. 

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

Thanks for working on this @sayakpaul! Let us know when it's ready for review i.e. when all tests are passing and there's no more changes to be made. 

Oh I thought the PR description made it clear. Sorry if I didn't. 

> The test failures can be tackled later once we match the logits to non-SDPA.

> Not very sure, though. Hence I am opening this PR for seeking feedback.

So, concretely, it would be great to have some initial reviews first so that we can localize why the basic logit assertion test isn't passing. 

> IMO we should avoid the complexe prepare_4d etc, and just use the _update_causal_mask. Now that's not super possible as would require deprecating, but anyways, not against this PR, 

@ArthurZucker if it’s easier/benefitting for the library, happy to use the methods you are suggesting. But I don’t get why things would need deprecating, etc. If you could provide more reference, that would be helpful. 

I disabled `is_causal` in `F.scaled_dot_product_attention()` and decided to handle the masking with the factory methods from the library itself. This is because CLIP applies BOTH causal mask and regular attention mask:

https://github.com/huggingface/transformers/blob/092f1fdaa4224fdd88c616dc9678e6fcb37bfffd/src/transformers/models/clip/modeling_clip.py#L287C1-L303C85

However, this leads to:

```
=========================== short test summary info ============================
FAILED tests/models/clip/test_modeling_clip.py::CLIPVisionModelTest::test_eager_matches_sdpa_inference_1_bfloat16
FAILED tests/models/clip/test_modeling_clip.py::CLIPTextModelTest::test_eager_matches_sdpa_inference_1_bfloat16
FAILED tests/models/clip/test_modeling_clip.py::CLIPTextModelTest::test_eager_matches_sdpa_inference_2_float32
FAILED tests/models/clip/test_modeling_clip.py::CLIPTextModelTest::test_sdpa_can_dispatch_on_flash
FAILED tests/models/clip/test_modeling_clip.py::CLIPModelTest::test_eager_matches_sdpa_inference_0_float16
FAILED tests/models/clip/test_modeling_clip.py::CLIPModelTest::test_eager_matches_sdpa_inference_1_bfloat16
FAILED tests/models/clip/test_modeling_clip.py::CLIPModelTest::test_eager_matches_sdpa_inference_2_float32
FAILED tests/models/clip/test_modeling_clip.py::CLIPModelTest::test_sdpa_can_dispatch_on_flash
FAILED tests/models/clip/test_modeling_clip.py::CLIPForImageClassificationModelTest::test_eager_matches_sdpa_inference_1_bfloat16
========== 9 failed, 208 passed, 159 skipped, 59 warnings in 40.40s ===========
```

Amongst these, failures for `test_sdpa_can_dispatch_on_flash()` are of particular interest. Would appreciate some pointers on approaching this. Would be also helpful to validate the treatment I am giving to the masks when using `CLIPSdpaAttention`. 

This passes:

```python
from transformers import AutoTokenizer, CLIPTextModel
import torch

model_sdpa = CLIPTextModel.from_pretrained("openai/clip-vit-base-patch32", attn_implementation="sdpa").to("cuda")
model = CLIPTextModel.from_pretrained("openai/clip-vit-base-patch32", attn_implementation="eager").to("cuda")
tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32")

inputs = tokenizer(["a photo of a cat", "a photo of a dog"], return_tensors="pt")
print(inputs["attention_mask"].tolist())
inputs = {k: v.to("cuda") for k, v in inputs.items()}

with torch.no_grad():
    outputs_sdpa = model_sdpa(**inputs)
    last_hidden_state_sdpa = outputs_sdpa.last_hidden_state
    outputs_eager = model(**inputs)
    last_hidden_state = outputs_eager.last_hidden_state


print(last_hidden_state_sdpa[0, :3, -1].flatten())
print(last_hidden_state[0, :3, -1].flatten())
print(torch.allclose(last_hidden_state_sdpa, last_hidden_state, rtol=1e-3, atol=1e-3))
```

@ArthurZucker done :) 

@ArthurZucker a friendly ping :)

@amyeroberts thanks for your reviews. I think I have addressed everything!

@amyeroberts in the light of https://github.com/huggingface/transformers/pull/30390#discussion_r1596645236, LMK if the current changes make sense, otherwise I don't see a good way to propagate the attention implementation method. 

Once that is sorted, I will work on https://github.com/huggingface/transformers/pull/30390#discussion_r1592956020. 

@amyeroberts am I doing something wrong when running `make style && make quality`? 

Here's what I have done:

* Created a fresh environment for the formatting related libs. 
* Ran `pip install -e ".[quality]"` inside the env.
* Then from the root of `transformers`, ran `make style && make quality`. 

@sayakpaul Could you try running `make fixup`? 

Leads to: 

<img width="1210" alt="image" src="https://github.com/huggingface/transformers/assets/22957388/b7c44677-9965-47ed-9bf8-e7db5518bd28">


@sayakpaul Huh, weird, I haven't seen that before. I'm going to try and re-trigger a fresh CI run, as it doesn't seem to be anything related to this PR 

@amyeroberts I had to touch a couple of loading utilities to make sure the equivalence tests pass. LMK if they should have been approached differently. 

Ah oh, the FLAX tests still fail the equivalence. I tried a bunch of state dict rejigging in order for the PT related changes to propagate in the FLAX model but none of them worked out. Would appreciate some guidance. 

Pinging @sanchit-gandhi here - who knows most of the intricacies of our flax models :) 

Closing in favor of https://github.com/huggingface/transformers/pull/31940.