From 1c37e8c1a6274e6e87b45c6319eb190757214c2a Mon Sep 17 00:00:00 2001 From: Pavel Iakubovskii Date: Thu, 18 Jul 2024 06:00:37 +0100 Subject: [PATCH] Add `sdpa` and FA2 for CLIP (#31940) * Squashed commit of the following: commit 102842cd477219b9f9bcb23a0bca3a8b92bd732f Author: Pavel Iakubovskii Date: Fri Jul 12 18:23:52 2024 +0000 Add model-specific sdpa tests commit 60e4c88581abf89ec098da84ed8e92aa904c997d Author: Pavel Iakubovskii Date: Fri Jul 12 18:20:53 2024 +0000 Add fallback to eager (expensive operation) commit c29033d30e7ffde4327e8a15cbbc6bee37546f80 Author: Pavel Iakubovskii Date: Thu Jul 11 17:09:55 2024 +0000 Fix attn_implementation propagation commit 783aed05f0f38cb2f99e758f81db6838ac55b9f8 Author: sayakpaul Date: Sat May 25 09:05:27 2024 +0530 style commit e77e703ca75d00447cda277eca6b886cd32bddc0 Author: sayakpaul Date: Sat May 25 09:04:57 2024 +0530 add comment to explain why I had to touch forbidden codebase. commit ab9d8849758e7773a31778ccba71588d18552623 Author: sayakpaul Date: Sat May 25 09:03:02 2024 +0530 fix: flax attribute access. commit c570fc0abf9d1bd58c291aae3c7e384f995996d2 Author: sayakpaul Date: Sat May 25 08:23:54 2024 +0530 fix tensorflow attribute name. commit 32c812871cfdb268d8a6e3e2c61c5c925c8ed47e Author: sayakpaul Date: Sat May 25 07:57:10 2024 +0530 fix attribute access. commit 4f41a0138b6c417aed9c9332278f8bcd979cb7c2 Author: sayakpaul Date: Sat May 25 07:44:02 2024 +0530 _from_config. commit 35aed64ff602422adcf41d7f677a0a24bd9eccae Author: sayakpaul Date: Fri May 24 18:46:52 2024 +0530 propagation of attn_implementation. commit 4c25c19845438b1dc1d35a5adf9436151c8c5940 Author: sayakpaul Date: Fri May 24 09:24:36 2024 +0530 style again commit 5f7dc5c5015c0f8116408f737e8c318d1802c80c Author: sayakpaul Date: Fri May 24 09:19:05 2024 +0530 use from_config. commit b70c409956d0359fa6ae5372275d2a20ba7e3389 Author: sayakpaul Date: Fri May 24 09:13:43 2024 +0530 quality commit a7b63beff53d0fc754c6564e2a7b51731ddee49d Author: sayakpaul Date: Fri May 10 14:35:10 2024 +0200 add benchmark numbers commit 455b0eaea50862b8458c8f422b60fe60ae40fdcb Author: sayakpaul Date: Fri May 10 13:50:16 2024 +0200 Revert "reflect feedback more" This reverts commit dc123e71eff60aae74d5f325f113d515d0d71117. commit ca674829d28787349c2a9593a14e0f1d41f04ea4 Author: sayakpaul Date: Fri May 10 13:50:05 2024 +0200 Revert "fix" This reverts commit 37a1cb35b87acdc4cf7528b8b1ed6da27d244e52. commit fab2dd8576c099eb1a3464958cb206a664d28247 Author: sayakpaul Date: Fri May 10 13:47:46 2024 +0200 fix commit fbc6ae50fd6f2d36294d31e191761631b701d696 Author: sayakpaul Date: Fri May 10 13:38:30 2024 +0200 reflect feedback more commit 87245bb020b2d60a89afe318a951df0159404fc9 Author: sayakpaul Date: Fri May 3 08:54:34 2024 +0530 fixes commit 1057cc26390ee839251e7f8b3326c4207595fb23 Author: sayakpaul Date: Fri May 3 07:49:03 2024 +0530 don't explicit set attn_implementation in tests commit e33f75916fc8a99f516b1cf449dbbe9d3aabda81 Author: sayakpaul Date: Fri May 3 07:43:54 2024 +0530 explicitly override attn_implementation in the towers. commit 4cf41cb1bc885c39df7cb8f2a0694ebf23299235 Author: sayakpaul Date: Fri May 3 07:38:42 2024 +0530 import in one-line. commit f2cc447ae9e74ccfacb448140cdf88259d4afc8c Author: sayakpaul Date: Fri May 3 07:34:58 2024 +0530 move sdpa mention to usage tips. commit 92884766c64dbb456926a3a84dd427be1349fa95 Author: sayakpaul Date: Mon Apr 29 10:58:26 2024 +0530 fix: memory allocation problem. commit d7ffbbfe12f7750b7d0a361420f35c13e0ea787d Author: sayakpaul Date: Mon Apr 29 09:56:59 2024 +0530 fix-copies commit 8dfc3731cedd02e36acd3fe56bb2e6d61efd25d8 Author: sayakpaul Date: Fri Apr 26 20:16:12 2024 +0530 address arthur's comments. commit d2ed7b4ce4ff15ae9aa4d3d0500f1544e3dcd9e9 Author: Sayak Paul Date: Fri Apr 26 20:08:15 2024 +0530 Apply suggestions from code review Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> commit 46e04361f37ded5c522ff05e9f725b9f82dce40e Author: sayakpaul Date: Wed Apr 24 09:55:27 2024 +0530 add to docs. commit 831629158ad40d34d8983f209afb2740ba041af2 Author: sayakpaul Date: Wed Apr 24 09:33:10 2024 +0530 styling.g commit d263a119c77314250f4b4c8469caf42559197f22 Author: sayakpaul Date: Wed Apr 24 09:15:20 2024 +0530 up commit d44f9d3d7633d4c241a737a1bc317f791f6aedb3 Author: sayakpaul Date: Tue Apr 23 18:40:42 2024 +0530 handle causal and attention mask commit 122f1d60153df6666b634a94e38d073f3f260926 Author: sayakpaul Date: Tue Apr 23 15:18:21 2024 +0530 test fixes. commit 4382d8cff6fa1dee5dbcf0d06b3e2841231e36f5 Author: sayakpaul Date: Tue Apr 23 09:39:25 2024 +0530 fix: scaling inside sdpa. commit 0f629989efc48b7315cf19405a81e02955efe7e5 Author: Sayak Paul Date: Tue Apr 23 08:14:58 2024 +0530 Update src/transformers/models/clip/modeling_clip.py Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> commit 14367316877dc27ea40f767ad1aee38bbc97e4ce Author: sayakpaul Date: Mon Apr 22 16:21:36 2024 +0530 add: sdpa support to clip. * Remove fallback for empty attention mask (expensive operation) * Fix typing in copies * Add flash attention * Add flash attention tests * List CLIP in FA docs * Fix embeddings attributes and tf * [run-slow] clip * Update clip documentation * Remove commented code, skip compile dynamic for CLIPModel * Fix doc * Fix doc 2 * Remove double transpose * Add torch version check for contiguous() * Add comment to test mixin * Fix copies * Add comment for mask * Update docs * [run-slow] clip --- docs/source/en/model_doc/clip.md | 117 ++++++ docs/source/en/perf_infer_gpu_one.md | 2 + .../models/altclip/modeling_altclip.py | 5 +- src/transformers/models/clip/modeling_clip.py | 204 ++++++++++- .../models/clipseg/modeling_clipseg.py | 13 +- src/transformers/models/git/modeling_git.py | 10 +- .../models/groupvit/modeling_groupvit.py | 4 +- src/transformers/models/idefics/vision.py | 6 +- .../models/kosmos2/modeling_kosmos2.py | 8 +- .../models/owlv2/modeling_owlv2.py | 2 +- .../models/owlvit/modeling_owlvit.py | 2 +- .../models/siglip/modeling_siglip.py | 2 +- .../models/x_clip/modeling_x_clip.py | 6 +- tests/models/clip/test_modeling_clip.py | 345 +++++++++++++++++- 14 files changed, 682 insertions(+), 44 deletions(-) diff --git a/docs/source/en/model_doc/clip.md b/docs/source/en/model_doc/clip.md index 692ea08371..f0829f484a 100644 --- a/docs/source/en/model_doc/clip.md +++ b/docs/source/en/model_doc/clip.md @@ -79,6 +79,123 @@ encode the text and prepare the images. The following example shows how to get t >>> probs = logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities ``` + +### Combining CLIP and Flash Attention 2 + +First, make sure to install the latest version of Flash Attention 2. + +```bash +pip install -U flash-attn --no-build-isolation +``` + +Make also sure that you have a hardware that is compatible with Flash-Attention 2. Read more about it in the official documentation of flash-attn repository. Make also sure to load your model in half-precision (e.g. `torch.float16`) + + + +For small batch sizes, you might notice a slowdown in your model when using flash attention. Refer to the section [Expected speedups with Flash Attention and SDPA](#Expected-speedups-with-Flash-Attention-and-SDPA) below and select an appropriate attention implementation. + + + +To load and run a model using Flash Attention 2, refer to the snippet below: + +```python +>>> import torch +>>> import requests +>>> from PIL import Image + +>>> from transformers import CLIPProcessor, CLIPModel + +>>> device = "cuda" +>>> torch_dtype = torch.float16 + +>>> model = CLIPModel.from_pretrained( +... "openai/clip-vit-base-patch32", +... attn_implementation="flash_attention_2", +... device_map=device, +... torch_dtype=torch_dtype, +... ) +>>> processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32") + +>>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" +>>> image = Image.open(requests.get(url, stream=True).raw) + +>>> inputs = processor(text=["a photo of a cat", "a photo of a dog"], images=image, return_tensors="pt", padding=True) +>>> inputs.to(device) + +>>> with torch.no_grad(): +... with torch.autocast(device): +... outputs = model(**inputs) + +>>> logits_per_image = outputs.logits_per_image # this is the image-text similarity score +>>> probs = logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities +>>> print(probs) +tensor([[0.9946, 0.0052]], device='cuda:0', dtype=torch.float16) +``` + + +### Using Scaled Dot Product Attention (SDPA) + +PyTorch includes a native scaled dot-product attention (SDPA) operator as part of `torch.nn.functional`. This function +encompasses several implementations that can be applied depending on the inputs and the hardware in use. See the +[official documentation](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html) +or the [GPU Inference](https://huggingface.co/docs/transformers/main/en/perf_infer_gpu_one#pytorch-scaled-dot-product-attention) +page for more information. + +SDPA is used by default for `torch>=2.1.1` when an implementation is available, but you may also set +`attn_implementation="sdpa"` in `from_pretrained()` to explicitly request SDPA to be used. + +```python +from transformers import CLIPModel + +model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32", torch_dtype=torch.float16, attn_implementation="sdpa") +``` + +For the best speedups, we recommend loading the model in half-precision (e.g. `torch.float16` or `torch.bfloat16`). + +### Expected speedups with Flash Attention and SDPA + +On a local benchmark (NVIDIA A10G, PyTorch 2.3.1+cu121) with `float16`, we saw the following speedups during inference for `"openai/clip-vit-large-patch14"` checkpoint ([code](https://gist.github.com/qubvel/ac691a54e54f9fae8144275f866a7ff8)): + +#### 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) + +#### 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 | 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 | 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 | 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 | 64 | 0.216 | 0.181 | 1.19 | 0.176 | 1.228 | + ## Resources A list of official Hugging Face and community (indicated by 🌎) resources to help you get started with CLIP. diff --git a/docs/source/en/perf_infer_gpu_one.md b/docs/source/en/perf_infer_gpu_one.md index 396c7bc2a9..b0109a0e8d 100644 --- a/docs/source/en/perf_infer_gpu_one.md +++ b/docs/source/en/perf_infer_gpu_one.md @@ -40,6 +40,7 @@ FlashAttention-2 is currently supported for the following architectures: * [Bark](https://huggingface.co/docs/transformers/model_doc/bark#transformers.BarkModel) * [Bart](https://huggingface.co/docs/transformers/model_doc/bart#transformers.BartModel) * [Chameleon](https://huggingface.co/docs/transformers/model_doc/chameleon#transformers.Chameleon) +* [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPModel) * [Cohere](https://huggingface.co/docs/transformers/model_doc/cohere#transformers.CohereModel) * [Dbrx](https://huggingface.co/docs/transformers/model_doc/dbrx#transformers.DbrxModel) * [DistilBert](https://huggingface.co/docs/transformers/model_doc/distilbert#transformers.DistilBertModel) @@ -200,6 +201,7 @@ For now, Transformers supports SDPA inference and training for the following arc * [Bart](https://huggingface.co/docs/transformers/model_doc/bart#transformers.BartModel) * [Bert](https://huggingface.co/docs/transformers/model_doc/bert#transformers.BertModel) * [Chameleon](https://huggingface.co/docs/transformers/model_doc/chameleon#transformers.Chameleon) +* [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPModel) * [Cohere](https://huggingface.co/docs/transformers/model_doc/cohere#transformers.CohereModel) * [Dbrx](https://huggingface.co/docs/transformers/model_doc/dbrx#transformers.DbrxModel) * [DeiT](https://huggingface.co/docs/transformers/model_doc/deit#transformers.DeiTModel) diff --git a/src/transformers/models/altclip/modeling_altclip.py b/src/transformers/models/altclip/modeling_altclip.py index 6bffdc70a5..10c9e10491 100755 --- a/src/transformers/models/altclip/modeling_altclip.py +++ b/src/transformers/models/altclip/modeling_altclip.py @@ -749,7 +749,7 @@ class AltCLIPAttention(nn.Module): attention_mask: Optional[torch.Tensor] = None, causal_attention_mask: Optional[torch.Tensor] = None, output_attentions: Optional[bool] = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """Input shape: Batch x Time x Channel""" bsz, tgt_len, embed_dim = hidden_states.size() @@ -838,7 +838,6 @@ class AltCLIPMLP(nn.Module): return hidden_states -# Copied from transformers.models.clip.modeling_clip.CLIPEncoderLayer with CLIP->AltCLIP class AltCLIPEncoderLayer(nn.Module): def __init__(self, config: AltCLIPConfig): super().__init__() @@ -889,7 +888,6 @@ class AltCLIPEncoderLayer(nn.Module): return outputs -# Copied from transformers.models.clip.modeling_clip.CLIPEncoder with CLIP->AltCLIP class AltCLIPEncoder(nn.Module): """ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a @@ -1080,7 +1078,6 @@ class AltCLIPPreTrainedModel(PreTrainedModel): module.weight.data[module.padding_idx].zero_() -# Copied from transformers.models.clip.modeling_clip.CLIPVisionTransformer with CLIPVisionTransformer->AltCLIPVisionTransformer,CLIPVisionConfig->AltCLIPVisionConfig,CLIPVisionEmbeddings->AltCLIPVisionEmbeddings,CLIPEncoder->AltCLIPEncoder,CLIP_VISION_INPUTS_DOCSTRING->ALTCLIP_VISION_INPUTS_DOCSTRING class AltCLIPVisionTransformer(nn.Module): def __init__(self, config: AltCLIPVisionConfig): super().__init__() diff --git a/src/transformers/models/clip/modeling_clip.py b/src/transformers/models/clip/modeling_clip.py index 48e6dfa849..b96acfc093 100644 --- a/src/transformers/models/clip/modeling_clip.py +++ b/src/transformers/models/clip/modeling_clip.py @@ -26,17 +26,24 @@ from ...activations import ACT2FN from ...modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ImageClassifierOutput from ...modeling_utils import PreTrainedModel +from ...pytorch_utils import is_torch_greater_or_equal_than_2_2 from ...utils import ( ModelOutput, add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, + is_flash_attn_2_available, + is_flash_attn_greater_or_equal_2_10, logging, replace_return_docstrings, ) from .configuration_clip import CLIPConfig, CLIPTextConfig, CLIPVisionConfig +if is_flash_attn_2_available(): + from ...modeling_flash_attention_utils import _flash_attention_forward + + logger = logging.get_logger(__name__) # General docstring @@ -254,7 +261,7 @@ class CLIPAttention(nn.Module): attention_mask: Optional[torch.Tensor] = None, causal_attention_mask: Optional[torch.Tensor] = None, output_attentions: Optional[bool] = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """Input shape: Batch x Time x Channel""" bsz, tgt_len, embed_dim = hidden_states.size() @@ -327,6 +334,173 @@ class CLIPAttention(nn.Module): return attn_output, attn_weights_reshaped +class CLIPFlashAttention2(CLIPAttention): + """ + CLIPAttention flash attention module. This module inherits from `CLIPAttention` as the weights of the module stays + untouched. The only required change would be on the forward pass where it needs to correctly call the public API of + flash attention and deal with padding tokens in case the input contains any of them. + """ + + # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__ + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. + # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. + # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). + self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() + + # Adapted from transformers.models.llama.modeling_llama.LlamaFlashAttention2.forward + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + causal_attention_mask: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = False, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + output_attentions = False + + batch_size, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + # Flash attention requires the input to have the shape + # batch_size x seq_length x head_dim x hidden_dim + # therefore we just need to keep the original shape + query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim) + key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim) + value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim) + + dropout_rate = self.dropout if self.training else 0.0 + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in the correct dtype just to be sure everything works as expected. + # This might slowdown training & inference so it is recommended to not cast the LayerNorms + # in fp32. + + input_dtype = query_states.dtype + if input_dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self.q_proj.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + attn_output = _flash_attention_forward( + query_states, + key_states, + value_states, + attention_mask, + q_len, + dropout=dropout_rate, + is_causal=causal_attention_mask is not None, + use_top_left_mask=self._flash_attn_uses_top_left_mask, + ) + + attn_output = attn_output.reshape(batch_size, q_len, self.embed_dim).contiguous() + attn_output = self.out_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights + + +class CLIPSdpaAttention(CLIPAttention): + """ + SDPA attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from + `CLIPAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to + SDPA API. + """ + + # Adapted from CLIPAttention.forward + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + causal_attention_mask: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = False, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + if output_attentions: + # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented. + logger.warning_once( + "CLIPModel is using CLIPSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not " + "support `output_attentions=True`. Falling back to the manual attention implementation, but specifying " + "the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can " + 'be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + return super().forward( + hidden_states=hidden_states, + attention_mask=attention_mask, + causal_attention_mask=causal_attention_mask, + output_attentions=output_attentions, + ) + + # CLIP text model uses both `causal_attention_mask` and `attention_mask` + if attention_mask is not None and causal_attention_mask is not None: + attn_mask = attention_mask + causal_attention_mask + elif causal_attention_mask is not None: + attn_mask = causal_attention_mask + else: + attn_mask = attention_mask + + bsz, tgt_len, embed_dim = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, -1, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, -1, self.num_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, -1, self.num_heads, self.head_dim).transpose(1, 2) + + # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, + # Reference: https://github.com/pytorch/pytorch/issues/112577. + if not is_torch_greater_or_equal_than_2_2 and query_states.device.type == "cuda" and attn_mask is not None: + query_states = query_states.contiguous() + key_states = key_states.contiguous() + value_states = value_states.contiguous() + + # CLIP text model uses both `causal_attention_mask` and `attention_mask` sequentially. + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states, + key_states, + value_states, + attn_mask=attn_mask, + dropout_p=self.dropout if self.training else 0.0, + scale=self.scale, + ) + + attn_output = attn_output.transpose(1, 2) + attn_output = attn_output.reshape(bsz, tgt_len, embed_dim) + + attn_output = self.out_proj(attn_output) + + return attn_output, None + + +CLIP_ATTENTION_CLASSES = { + "eager": CLIPAttention, + "sdpa": CLIPSdpaAttention, + "flash_attention_2": CLIPFlashAttention2, +} + + class CLIPMLP(nn.Module): def __init__(self, config): super().__init__() @@ -346,7 +520,7 @@ class CLIPEncoderLayer(nn.Module): def __init__(self, config: CLIPConfig): super().__init__() self.embed_dim = config.hidden_size - self.self_attn = CLIPAttention(config) + self.self_attn = CLIP_ATTENTION_CLASSES[config._attn_implementation](config) self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) self.mlp = CLIPMLP(config) self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) @@ -401,6 +575,8 @@ class CLIPPreTrainedModel(PreTrainedModel): config_class = CLIPConfig base_model_prefix = "clip" supports_gradient_checkpointing = True + _supports_sdpa = True + _supports_flash_attn_2 = True def _init_weights(self, module): """Initialize the weights""" @@ -668,6 +844,9 @@ class CLIPTextTransformer(nn.Module): # For `pooled_output` computation self.eos_token_id = config.eos_token_id + # For attention mask, it differs between `flash_attention_2` and other attention implementations + self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" + @add_start_docstrings_to_model_forward(CLIP_TEXT_INPUTS_DOCSTRING) @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=CLIPTextConfig) def forward( @@ -702,8 +881,9 @@ class CLIPTextTransformer(nn.Module): causal_attention_mask = _create_4d_causal_attention_mask( input_shape, hidden_states.dtype, device=hidden_states.device ) + # expand attention_mask - if attention_mask is not None: + if attention_mask is not None and not self._use_flash_attention_2: # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) @@ -957,8 +1137,11 @@ class CLIPModel(CLIPPreTrainedModel): self.text_embed_dim = text_config.hidden_size self.vision_embed_dim = vision_config.hidden_size - self.text_model = CLIPTextTransformer(text_config) - self.vision_model = CLIPVisionTransformer(vision_config) + text_model = CLIPTextModel._from_config(text_config, attn_implementation=config._attn_implementation) + self.text_model = text_model.text_model + + vision_model = CLIPVisionModel._from_config(vision_config, attn_implementation=config._attn_implementation) + self.vision_model = vision_model.vision_model self.visual_projection = nn.Linear(self.vision_embed_dim, self.projection_dim, bias=False) self.text_projection = nn.Linear(self.text_embed_dim, self.projection_dim, bias=False) @@ -1173,7 +1356,8 @@ class CLIPTextModelWithProjection(CLIPPreTrainedModel): def __init__(self, config: CLIPTextConfig): super().__init__(config) - self.text_model = CLIPTextTransformer(config) + text_model = CLIPTextModel._from_config(config, attn_implementation=config._attn_implementation) + self.text_model = text_model.text_model self.text_projection = nn.Linear(config.hidden_size, config.projection_dim, bias=False) @@ -1253,7 +1437,8 @@ class CLIPVisionModelWithProjection(CLIPPreTrainedModel): def __init__(self, config: CLIPVisionConfig): super().__init__(config) - self.vision_model = CLIPVisionTransformer(config) + vision_model = CLIPVisionModel._from_config(config, attn_implementation=config._attn_implementation) + self.vision_model = vision_model.vision_model self.visual_projection = nn.Linear(config.hidden_size, config.projection_dim, bias=False) @@ -1332,7 +1517,10 @@ class CLIPForImageClassification(CLIPPreTrainedModel): super().__init__(config) self.num_labels = config.num_labels - self.vision_model = CLIPVisionTransformer(config.vision_config) + vision_model = CLIPVisionModel._from_config( + config.vision_config, attn_implementation=config._attn_implementation + ) + self.vision_model = vision_model.vision_model # Classifier head self.classifier = ( diff --git a/src/transformers/models/clipseg/modeling_clipseg.py b/src/transformers/models/clipseg/modeling_clipseg.py index 24d4b2322e..af7b94a10f 100644 --- a/src/transformers/models/clipseg/modeling_clipseg.py +++ b/src/transformers/models/clipseg/modeling_clipseg.py @@ -266,7 +266,7 @@ class CLIPSegAttention(nn.Module): attention_mask: Optional[torch.Tensor] = None, causal_attention_mask: Optional[torch.Tensor] = None, output_attentions: Optional[bool] = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """Input shape: Batch x Time x Channel""" bsz, tgt_len, embed_dim = hidden_states.size() @@ -355,7 +355,7 @@ class CLIPSegMLP(nn.Module): return hidden_states -# Copied from transformers.models.clip.modeling_clip.CLIPEncoderLayer with CLIP->CLIPSeg +# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoderLayer with AltCLIP->CLIPSeg class CLIPSegEncoderLayer(nn.Module): def __init__(self, config: CLIPSegConfig): super().__init__() @@ -554,7 +554,7 @@ CLIPSEG_INPUTS_DOCSTRING = r""" """ -# Copied from transformers.models.clip.modeling_clip.CLIPEncoder with CLIP->CLIPSeg +# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoder with AltCLIP->CLIPSeg class CLIPSegEncoder(nn.Module): """ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a @@ -653,7 +653,6 @@ class CLIPSegEncoder(nn.Module): class CLIPSegTextTransformer(nn.Module): - # Copied from transformers.models.clip.modeling_clip.CLIPTextTransformer.__init__ with CLIP->CLIPSeg def __init__(self, config: CLIPSegTextConfig): super().__init__() self.config = config @@ -667,7 +666,7 @@ class CLIPSegTextTransformer(nn.Module): @add_start_docstrings_to_model_forward(CLIPSEG_TEXT_INPUTS_DOCSTRING) @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=CLIPSegTextConfig) - # Copied from transformers.models.clip.modeling_clip.CLIPTextTransformer.forward with clip->clipseg, CLIP->CLIPSeg + # Adapted from transformers.models.clip.modeling_clip.CLIPTextTransformer.forward with clip->clipseg, CLIP->CLIPSeg def forward( self, input_ids: Optional[torch.Tensor] = None, @@ -806,7 +805,7 @@ class CLIPSegTextModel(CLIPSegPreTrainedModel): class CLIPSegVisionTransformer(nn.Module): - # Copied from transformers.models.clip.modeling_clip.CLIPVisionTransformer.__init__ with CLIP->CLIPSeg + # Copied from transformers.models.altclip.modeling_altclip.AltCLIPVisionTransformer.__init__ with AltCLIP->CLIPSeg def __init__(self, config: CLIPSegVisionConfig): super().__init__() self.config = config @@ -1149,7 +1148,7 @@ class CLIPSegDecoderLayer(nn.Module): self-attention/MLP, rather than before. """ - # Copied from transformers.models.clip.modeling_clip.CLIPEncoderLayer.__init__ with CLIP->CLIPSeg + # Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoderLayer.__init__ with AltCLIP->CLIPSeg def __init__(self, config: CLIPSegConfig): super().__init__() self.embed_dim = config.hidden_size diff --git a/src/transformers/models/git/modeling_git.py b/src/transformers/models/git/modeling_git.py index 8e14e3a899..27de8c688b 100644 --- a/src/transformers/models/git/modeling_git.py +++ b/src/transformers/models/git/modeling_git.py @@ -632,7 +632,7 @@ class GitVisionMLP(nn.Module): return hidden_states -# Copied from transformers.models.clip.modeling_clip.CLIPAttention +# Copied from transformers.models.clip.modeling_clip.CLIPAttention with CLIP->GitVision class GitVisionAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" @@ -664,7 +664,7 @@ class GitVisionAttention(nn.Module): attention_mask: Optional[torch.Tensor] = None, causal_attention_mask: Optional[torch.Tensor] = None, output_attentions: Optional[bool] = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """Input shape: Batch x Time x Channel""" bsz, tgt_len, embed_dim = hidden_states.size() @@ -737,7 +737,7 @@ class GitVisionAttention(nn.Module): return attn_output, attn_weights_reshaped -# Copied from transformers.models.clip.modeling_clip.CLIPEncoderLayer with CLIP->GitVision +# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoderLayer with AltCLIP->GitVision class GitVisionEncoderLayer(nn.Module): def __init__(self, config: GitVisionConfig): super().__init__() @@ -788,7 +788,7 @@ class GitVisionEncoderLayer(nn.Module): return outputs -# Copied from transformers.models.clip.modeling_clip.CLIPEncoder with CLIP->GitVision, CLIPConfig +# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoder with AltCLIP->GitVision, CLIPConfig class GitVisionEncoder(nn.Module): """ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a @@ -903,7 +903,7 @@ GIT_VISION_INPUTS_DOCSTRING = r""" class GitVisionTransformer(nn.Module): - # Copied from transformers.models.clip.modeling_clip.CLIPVisionTransformer.__init__ with CLIPEncoder->GitVisionEncoder, CLIP->Git + # Copied from transformers.models.altclip.modeling_altclip.AltCLIPVisionTransformer.__init__ with AltCLIPEncoder->GitVisionEncoder, AltCLIP->Git def __init__(self, config: GitVisionConfig): super().__init__() self.config = config diff --git a/src/transformers/models/groupvit/modeling_groupvit.py b/src/transformers/models/groupvit/modeling_groupvit.py index 99be160319..32e1d777cb 100644 --- a/src/transformers/models/groupvit/modeling_groupvit.py +++ b/src/transformers/models/groupvit/modeling_groupvit.py @@ -688,7 +688,7 @@ class GroupViTAttention(nn.Module): return attn_output, attn_weights_reshaped -# Copied from transformers.models.clip.modeling_clip.CLIPEncoderLayer with CLIP->GroupViT +# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoderLayer with AltCLIP->GroupViT class GroupViTEncoderLayer(nn.Module): def __init__(self, config: GroupViTConfig): super().__init__() @@ -1034,7 +1034,6 @@ class GroupViTTextEncoder(nn.Module): ) -# Copied from transformers.models.clip.modeling_clip.CLIPTextTransformer with CLIPText->GroupViTText, CLIPEncoder->GroupViTTextEncoder, CLIP_TEXT->GROUPVIT_TEXT class GroupViTTextTransformer(nn.Module): def __init__(self, config: GroupViTTextConfig): super().__init__() @@ -1081,6 +1080,7 @@ class GroupViTTextTransformer(nn.Module): causal_attention_mask = _create_4d_causal_attention_mask( input_shape, hidden_states.dtype, device=hidden_states.device ) + # expand attention_mask if attention_mask is not None: # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] diff --git a/src/transformers/models/idefics/vision.py b/src/transformers/models/idefics/vision.py index 847e92e89c..5339b70692 100644 --- a/src/transformers/models/idefics/vision.py +++ b/src/transformers/models/idefics/vision.py @@ -192,7 +192,7 @@ class IdeficsVisionAttention(nn.Module): attention_mask: Optional[torch.Tensor] = None, causal_attention_mask: Optional[torch.Tensor] = None, output_attentions: Optional[bool] = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """Input shape: Batch x Time x Channel""" bsz, tgt_len, embed_dim = hidden_states.size() @@ -281,7 +281,7 @@ class IdeficsVisionMLP(nn.Module): return hidden_states -# Copied from transformers.models.clip.modeling_clip.CLIPEncoderLayer with CLIP->IdeficsVision +# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoderLayer with AltCLIP->IdeficsVision class IdeficsVisionEncoderLayer(nn.Module): def __init__(self, config: IdeficsVisionConfig): super().__init__() @@ -332,7 +332,7 @@ class IdeficsVisionEncoderLayer(nn.Module): return outputs -# Copied from transformers.models.clip.modeling_clip.CLIPEncoder with CLIP->IdeficsVision +# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoder with AltCLIP->IdeficsVision class IdeficsVisionEncoder(nn.Module): """ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a diff --git a/src/transformers/models/kosmos2/modeling_kosmos2.py b/src/transformers/models/kosmos2/modeling_kosmos2.py index 9585bd891e..69641790b2 100644 --- a/src/transformers/models/kosmos2/modeling_kosmos2.py +++ b/src/transformers/models/kosmos2/modeling_kosmos2.py @@ -444,7 +444,7 @@ class Kosmos2VisionAttention(nn.Module): attention_mask: Optional[torch.Tensor] = None, causal_attention_mask: Optional[torch.Tensor] = None, output_attentions: Optional[bool] = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """Input shape: Batch x Time x Channel""" bsz, tgt_len, embed_dim = hidden_states.size() @@ -533,7 +533,7 @@ class Kosmos2VisionMLP(nn.Module): return hidden_states -# Copied from transformers.models.clip.modeling_clip.CLIPEncoderLayer with CLIP->Kosmos2Vision +# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoderLayer with AltCLIP->Kosmos2Vision class Kosmos2VisionEncoderLayer(nn.Module): def __init__(self, config: Kosmos2VisionConfig): super().__init__() @@ -584,7 +584,7 @@ class Kosmos2VisionEncoderLayer(nn.Module): return outputs -# Copied from transformers.models.clip.modeling_clip.CLIPEncoder with CLIP->Kosmos2Vision +# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoder with AltCLIP->Kosmos2Vision class Kosmos2VisionEncoder(nn.Module): """ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a @@ -684,7 +684,7 @@ class Kosmos2VisionEncoder(nn.Module): # Similar to `transformers.models.clip.modeling_clip.CLIPVisionTransformer` but without docstring for `forward` class Kosmos2VisionTransformer(nn.Module): - # Copied from transformers.models.clip.modeling_clip.CLIPVisionTransformer.__init__ with CLIPVision->Kosmos2Vision,CLIP_VISION->KOSMOS2_VISION,CLIP->Kosmos2Vision + # Copied from transformers.models.altclip.modeling_altclip.AltCLIPVisionTransformer.__init__ with AltCLIPVision->Kosmos2Vision,ALTCLIP_VISION->KOSMOS2_VISION,AltCLIP->Kosmos2Vision def __init__(self, config: Kosmos2VisionConfig): super().__init__() self.config = config diff --git a/src/transformers/models/owlv2/modeling_owlv2.py b/src/transformers/models/owlv2/modeling_owlv2.py index 638a9d966e..0c4b60a4f5 100644 --- a/src/transformers/models/owlv2/modeling_owlv2.py +++ b/src/transformers/models/owlv2/modeling_owlv2.py @@ -459,7 +459,7 @@ class Owlv2MLP(nn.Module): return hidden_states -# Copied from transformers.models.clip.modeling_clip.CLIPEncoderLayer with CLIP->Owlv2 +# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoderLayer with AltCLIP->Owlv2 class Owlv2EncoderLayer(nn.Module): def __init__(self, config: Owlv2Config): super().__init__() diff --git a/src/transformers/models/owlvit/modeling_owlvit.py b/src/transformers/models/owlvit/modeling_owlvit.py index 32e2012b21..89d92c2209 100644 --- a/src/transformers/models/owlvit/modeling_owlvit.py +++ b/src/transformers/models/owlvit/modeling_owlvit.py @@ -451,7 +451,7 @@ class OwlViTMLP(nn.Module): return hidden_states -# Copied from transformers.models.clip.modeling_clip.CLIPEncoderLayer with CLIP->OwlViT +# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoderLayer with AltCLIP->OwlViT class OwlViTEncoderLayer(nn.Module): def __init__(self, config: OwlViTConfig): super().__init__() diff --git a/src/transformers/models/siglip/modeling_siglip.py b/src/transformers/models/siglip/modeling_siglip.py index cc51dbd776..7c15dea387 100644 --- a/src/transformers/models/siglip/modeling_siglip.py +++ b/src/transformers/models/siglip/modeling_siglip.py @@ -829,7 +829,7 @@ SIGLIP_INPUTS_DOCSTRING = r""" """ -# Copied from transformers.models.clip.modeling_clip.CLIPEncoder with CLIP->Siglip +# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoder with AltCLIP->Siglip class SiglipEncoder(nn.Module): """ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a diff --git a/src/transformers/models/x_clip/modeling_x_clip.py b/src/transformers/models/x_clip/modeling_x_clip.py index 21b1c88aa0..b00b42281b 100644 --- a/src/transformers/models/x_clip/modeling_x_clip.py +++ b/src/transformers/models/x_clip/modeling_x_clip.py @@ -199,7 +199,7 @@ class XCLIPAttention(nn.Module): attention_mask: Optional[torch.Tensor] = None, causal_attention_mask: Optional[torch.Tensor] = None, output_attentions: Optional[bool] = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """Input shape: Batch x Time x Channel""" bsz, tgt_len, embed_dim = hidden_states.size() @@ -288,7 +288,7 @@ class XCLIPMLP(nn.Module): return hidden_states -# Copied from transformers.models.clip.modeling_clip.CLIPEncoderLayer with CLIP->XCLIP +# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoderLayer with AltCLIP->XCLIP class XCLIPEncoderLayer(nn.Module): def __init__(self, config: XCLIPConfig): super().__init__() @@ -609,7 +609,7 @@ X_CLIP_INPUTS_DOCSTRING = r""" """ -# Copied from transformers.models.clip.modeling_clip.CLIPEncoder with CLIP->XCLIP +# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoder with AltCLIP->XCLIP class XCLIPEncoder(nn.Module): """ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a diff --git a/tests/models/clip/test_modeling_clip.py b/tests/models/clip/test_modeling_clip.py index 78a5fb6f9a..3b69944280 100644 --- a/tests/models/clip/test_modeling_clip.py +++ b/tests/models/clip/test_modeling_clip.py @@ -18,21 +18,33 @@ import inspect import os import tempfile import unittest +from typing import Optional, Tuple import numpy as np import requests +from parameterized import parameterized +from pytest import mark import transformers from transformers import CLIPConfig, CLIPTextConfig, CLIPVisionConfig from transformers.testing_utils import ( is_flax_available, is_pt_flax_cross_test, + require_flash_attn, require_torch, + require_torch_gpu, + require_torch_sdpa, require_vision, slow, torch_device, ) -from transformers.utils import is_torch_available, is_vision_available +from transformers.utils import ( + is_torch_available, + is_torch_bf16_available_on_device, + is_torch_fp16_available_on_device, + is_torch_sdpa_available, + is_vision_available, +) from ...test_configuration_common import ConfigTester from ...test_modeling_common import ( @@ -40,6 +52,7 @@ from ...test_modeling_common import ( _config_zero_init, floats_tensor, ids_tensor, + is_flaky, random_attention_mask, ) from ...test_pipeline_mixin import PipelineTesterMixin @@ -59,6 +72,10 @@ if is_torch_available(): ) +if is_torch_sdpa_available(): + from torch.nn.attention import SDPBackend, sdpa_kernel + + if is_vision_available(): from PIL import Image @@ -167,8 +184,180 @@ class CLIPVisionModelTester: return config, inputs_dict +class CLIPModelTesterMixin(ModelTesterMixin): + """ + Subclass of ModelTesterMixin with methods specific to testing CLIP models. + The SDPA equivalence test is overridden here because CLIP models may have test/vision/text+vision inputs, + different output logits, and are not supposed to be used or tested with padding_side="left". + """ + + def test_eager_matches_sdpa_inference( + self, + torch_dtype: str, + use_attention_mask_options: Tuple[Optional[str], ...] = (None, "left", "right"), + logit_keys: Tuple[str, ...] = ("logits_per_image", "logits_per_text", "image_embeds", "text_embeds"), + ): + if not self.all_model_classes[0]._supports_sdpa: + self.skipTest(f"{self.all_model_classes[0].__name__} does not support SDPA") + + if torch_dtype == "float16" and not is_torch_fp16_available_on_device(torch_device): + self.skipTest(f"float16 not supported on {torch_device} (on the specific device currently used)") + + if torch_dtype == "bfloat16" and not is_torch_bf16_available_on_device(torch_device): + self.skipTest( + f"bfloat16 not supported on {torch_device} (on the specific device currently used, e.g. Nvidia T4 GPU)" + ) + + # Convert to torch dtype + dtypes = { + "float16": torch.float16, + "bfloat16": torch.bfloat16, + "float32": torch.float32, + } + torch_dtype = dtypes[torch_dtype] + + atols = { + torch.float32: 1e-5, + torch.bfloat16: 3e-2, + torch.float16: 5e-3, + } + rtols = { + torch.float32: 1e-4, + torch.bfloat16: 3e-2, + torch.float16: 5e-3, + } + + atol = atols[torch_dtype] + rtol = rtols[torch_dtype] + + def get_mean_reldiff(msg, current_case, x, ref, atol, rtol): + return f"{msg} {current_case}: mean relative difference: {((x - ref).abs() / (ref.abs() + 1e-12)).mean():.3e}, torch atol = {atol}, torch rtol = {rtol}" + + for model_class in self.all_model_classes: + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + model = model_class(config) + + with tempfile.TemporaryDirectory() as tmpdirname: + model.save_pretrained(tmpdirname) + + # Load the model with SDPA + model_sdpa = model_class.from_pretrained(tmpdirname, torch_dtype=torch_dtype) + model_sdpa = model_sdpa.eval().to(torch_device) + + # Load model with eager attention + model_eager = model_class.from_pretrained( + tmpdirname, + torch_dtype=torch_dtype, + attn_implementation="eager", + ) + model_eager = model_eager.eval().to(torch_device) + + self.assertTrue(model_sdpa.config._attn_implementation == "sdpa") + self.assertTrue(model_eager.config._attn_implementation == "eager") + + for name, submodule in model_eager.named_modules(): + class_name = submodule.__class__.__name__ + if "SdpaAttention" in class_name or "SdpaSelfAttention" in class_name: + raise ValueError("The eager model should not have SDPA attention layers") + + has_sdpa = False + for name, submodule in model_sdpa.named_modules(): + class_name = submodule.__class__.__name__ + if "SdpaAttention" in class_name or "SdpaSelfAttention" in class_name: + has_sdpa = True + break + + if not has_sdpa: + raise ValueError("The SDPA model should have SDPA attention layers") + + # We use these for loops instead of parameterized.expand just for the interest of avoiding loading/saving the model each time, + # but it would be nicer to have an efficient way to use parameterized.expand + cases = [ + (use_mask, output_attentions, sdpa_backend, batch_size) + for use_mask in use_attention_mask_options + for output_attentions in [True, False] + for sdpa_backend in [ + [SDPBackend.MATH], + [SDPBackend.FLASH_ATTENTION, SDPBackend.MATH], + [SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH], + [SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH], + ] + for batch_size in [1, 5] + ] + fail_cases = [] + + for use_mask, output_attentions, sdpa_backend, batch_size in cases: + processed_inputs = inputs_dict.copy() + + # convert to torch_dtype + if "pixel_values" in processed_inputs: + processed_inputs["pixel_values"] = processed_inputs["pixel_values"].to(torch_dtype) + + # slice for different batch sizes + for key in ["pixel_values", "input_ids", "attention_mask"]: + if key in processed_inputs: + processed_inputs[key] = processed_inputs[key][:batch_size] + + # set attention mask with left padding + if not use_mask: + processed_inputs.pop("attention_mask", None) + elif use_mask == "left": + dummy_attention_mask = processed_inputs["attention_mask"] + dummy_attention_mask[:] = 1 + dummy_attention_mask[:, :1] = 0 + processed_inputs["attention_mask"] = dummy_attention_mask + elif use_mask == "right": + dummy_attention_mask = processed_inputs["attention_mask"] + dummy_attention_mask[:] = 1 + dummy_attention_mask[:, -1:] = 0 + processed_inputs["attention_mask"] = dummy_attention_mask + else: + raise ValueError(f"Invalid value for use_mask={use_mask}") + + processed_inputs["output_attentions"] = output_attentions + processed_inputs["output_hidden_states"] = True + + current_case = f"use_mask={use_mask}, batch_size={batch_size}, sdpa_backend={sdpa_backend}" + + prepared_inputs = self._prepare_for_class(processed_inputs, model_class) + + with torch.no_grad(): + try: + with sdpa_kernel(sdpa_backend): + outputs_eager = model_eager(**prepared_inputs) + outputs_sdpa = model_sdpa(**prepared_inputs) + except Exception as e: + fail_cases.append(f"{current_case}: {e}") + continue + + keys = set(logit_keys) & set(outputs_eager.keys()) + self.assertTrue( + keys, f"Keys {logit_keys} not found in outputs. Available keys: {outputs_eager.keys()}" + ) + + for key in keys: + try: + eager_logits = outputs_eager[key] + sdpa_logits = outputs_sdpa[key] + except KeyError: + raise KeyError(f"Key {key} not found in outputs. Available keys: {outputs_eager.keys()}") + + if "hidden_state" in key and use_mask == "left": + eager_logits = eager_logits[:, 1:] + sdpa_logits = sdpa_logits[:, 1:] + elif "hidden_state" in key and use_mask == "right": + eager_logits = eager_logits[:, :-1] + sdpa_logits = sdpa_logits[:, :-1] + + is_close = torch.allclose(eager_logits, sdpa_logits, atol=atol, rtol=rtol) + if not is_close: + fail_cases.append(get_mean_reldiff(key, current_case, sdpa_logits, eager_logits, atol, rtol)) + + self.assertTrue(len(fail_cases) == 0, "\n".join(fail_cases)) + + @require_torch -class CLIPVisionModelTest(ModelTesterMixin, unittest.TestCase): +class CLIPVisionModelTest(CLIPModelTesterMixin, unittest.TestCase): """ Here we also overwrite some of the tests of test_modeling_common.py, as CLIP does not use input_ids, inputs_embeds, attention_mask and seq_length. @@ -261,6 +450,17 @@ class CLIPVisionModelTest(ModelTesterMixin, unittest.TestCase): self.assertIsNotNone(model) self.assertTrue(hasattr(model, "visual_projection")) + @parameterized.expand([("float16",), ("bfloat16",), ("float32",)]) + @require_torch_sdpa + @slow + @is_flaky() + def test_eager_matches_sdpa_inference(self, torch_dtype: str): + super().test_eager_matches_sdpa_inference( + torch_dtype=torch_dtype, + logit_keys=("last_hidden_state", "pooler_output", "image_embeds"), + use_attention_mask_options=(None,), + ) + class CLIPTextModelTester: def __init__( @@ -361,7 +561,7 @@ class CLIPTextModelTester: @require_torch -class CLIPTextModelTest(ModelTesterMixin, unittest.TestCase): +class CLIPTextModelTest(CLIPModelTesterMixin, unittest.TestCase): all_model_classes = (CLIPTextModel, CLIPTextModelWithProjection) if is_torch_available() else () fx_compatible = True test_pruning = False @@ -428,6 +628,21 @@ class CLIPTextModelTest(ModelTesterMixin, unittest.TestCase): self.assertIsNotNone(model) self.assertTrue(hasattr(model, "text_projection")) + @parameterized.expand([("float16",), ("bfloat16",), ("float32",)]) + @require_torch_sdpa + @slow + @is_flaky() + def test_eager_matches_sdpa_inference(self, torch_dtype: str): + super().test_eager_matches_sdpa_inference( + torch_dtype=torch_dtype, + logit_keys=("last_hidden_state", "pooler_output", "text_embeds"), + use_attention_mask_options=(None, "right"), # "left" is not supported for text model + ) + + @require_torch_sdpa + def test_sdpa_can_dispatch_on_flash(self): + self.skipTest(reason="CLIPTextModel has two attention masks: `causal_attention_mask` and `attention_mask`") + class CLIPModelTester: def __init__(self, parent, text_kwargs=None, vision_kwargs=None, is_training=True): @@ -479,7 +694,7 @@ class CLIPModelTester: @require_torch -class CLIPModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase): +class CLIPModelTest(CLIPModelTesterMixin, PipelineTesterMixin, unittest.TestCase): all_model_classes = (CLIPModel,) if is_torch_available() else () pipeline_model_mapping = ( {"feature-extraction": CLIPModel, "image-feature-extraction": CLIPVisionModel} if is_torch_available() else {} @@ -746,6 +961,115 @@ class CLIPModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase): model = CLIPModel.from_pretrained(model_name) self.assertIsNotNone(model) + @parameterized.expand([("float16",), ("bfloat16",), ("float32",)]) + @require_torch_sdpa + @slow + @is_flaky() + def test_eager_matches_sdpa_inference(self, torch_dtype: str): + super().test_eager_matches_sdpa_inference( + torch_dtype=torch_dtype, + logit_keys=("logits_per_image", "logits_per_text"), + use_attention_mask_options=(None, "right"), # "left" is not supported for text model + ) + + @require_torch_sdpa + def test_sdpa_can_dispatch_on_flash(self): + self.skipTest(reason="CLIP text tower has two attention masks: `causal_attention_mask` and `attention_mask`") + + @require_torch_sdpa + def test_sdpa_can_compile_dynamic(self): + self.skipTest(reason="CLIP model can't be compiled dynamic, error in clip_loss`") + + @require_flash_attn + @require_torch_gpu + @mark.flash_attn_test + @slow + def test_flash_attn_2_inference_equivalence(self): + for model_class in self.all_model_classes: + if not model_class._supports_flash_attn_2: + self.skipTest(f"{model_class.__name__} does not support Flash Attention 2") + + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + model = model_class(config) + + with tempfile.TemporaryDirectory() as tmpdirname: + model.save_pretrained(tmpdirname) + model_fa = model_class.from_pretrained( + tmpdirname, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2" + ) + model_fa.to(torch_device) + + model = model_class.from_pretrained(tmpdirname, torch_dtype=torch.bfloat16) + model.to(torch_device) + + dummy_pixel_values = inputs_dict["pixel_values"].to(torch.bfloat16) + dummy_input_ids = inputs_dict["input_ids"] + + outputs = model(pixel_values=dummy_pixel_values, input_ids=dummy_input_ids, output_hidden_states=True) + outputs_fa = model_fa( + pixel_values=dummy_pixel_values, input_ids=dummy_input_ids, output_hidden_states=True + ) + + self.assertTrue( + torch.allclose(outputs.logits_per_image, outputs_fa.logits_per_image, atol=4e-2, rtol=4e-2), + f"Image logits max diff: {torch.max(torch.abs(outputs.logits_per_image - outputs_fa.logits_per_image))}", + ) + self.assertTrue( + torch.allclose(outputs.logits_per_text, outputs_fa.logits_per_text, atol=4e-2, rtol=4e-2), + f"Text logits max diff: {torch.max(torch.abs(outputs.logits_per_text - outputs_fa.logits_per_text))}", + ) + + @require_flash_attn + @require_torch_gpu + @mark.flash_attn_test + def test_flash_attn_2_inference_equivalence_right_padding(self): + for model_class in self.all_model_classes: + if not model_class._supports_flash_attn_2: + self.skipTest(f"{model_class.__name__} does not support Flash Attention 2") + + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + model = model_class(config) + + with tempfile.TemporaryDirectory() as tmpdirname: + model.save_pretrained(tmpdirname) + model_fa = model_class.from_pretrained( + tmpdirname, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2" + ) + model_fa.to(torch_device) + + model = model_class.from_pretrained( + tmpdirname, torch_dtype=torch.bfloat16, attn_implementation="eager" + ) + model.to(torch_device) + + dummy_pixel_values = inputs_dict["pixel_values"].to(torch.bfloat16) + dummy_input_ids = inputs_dict["input_ids"] + dummy_pixel_mask = inputs_dict["attention_mask"] + + # right padding + dummy_pixel_mask[:] = 1 + dummy_pixel_mask[:, -1:] = 0 + + outputs = model(pixel_values=dummy_pixel_values, input_ids=dummy_input_ids, output_hidden_states=True) + outputs_fa = model_fa( + pixel_values=dummy_pixel_values, input_ids=dummy_input_ids, output_hidden_states=True + ) + + logits_per_image_eager = outputs.logits_per_image[:, :-1] + logits_per_text_eager = outputs.logits_per_text[:, :-1] + + logits_per_image_sdpa = outputs_fa.logits_per_image[:, :-1] + logits_per_text_sdpa = outputs_fa.logits_per_text[:, :-1] + + self.assertTrue( + torch.allclose(logits_per_image_eager, logits_per_image_sdpa, atol=4e-2, rtol=4e-2), + f"Image logits max diff: {torch.max(torch.abs(logits_per_image_eager - logits_per_image_sdpa))}", + ) + self.assertTrue( + torch.allclose(logits_per_text_eager, logits_per_text_sdpa, atol=4e-2, rtol=4e-2), + f"Text logits max diff: {torch.max(torch.abs(logits_per_text_eager - logits_per_text_sdpa))}", + ) + class CLIPForImageClassificationModelTester(CLIPModelTester): def __init__(self, parent): @@ -769,7 +1093,7 @@ class CLIPForImageClassificationModelTester(CLIPModelTester): @require_torch -class CLIPForImageClassificationModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase): +class CLIPForImageClassificationModelTest(CLIPModelTesterMixin, PipelineTesterMixin, unittest.TestCase): all_model_classes = (CLIPForImageClassification,) if is_torch_available() else () pipeline_model_mapping = {"image-classification": CLIPForImageClassification} if is_torch_available() else {} fx_compatible = False @@ -805,6 +1129,17 @@ class CLIPForImageClassificationModelTest(ModelTesterMixin, PipelineTesterMixin, def test_initialization(self): pass + @parameterized.expand([("float16",), ("bfloat16",), ("float32",)]) + @require_torch_sdpa + @slow + @is_flaky() + def test_eager_matches_sdpa_inference(self, torch_dtype: str): + super().test_eager_matches_sdpa_inference( + torch_dtype=torch_dtype, + logit_keys=("logits",), + use_attention_mask_options=(None,), + ) + # We will verify our results on an image of cute cats def prepare_img():