Add `sdpa` and FA2 for CLIP #31940

qubvel · 2024-07-12T21:40:45Z

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:

Before submitting

Did you read the contributor guideline,
Pull Request section?
Was this discussed/approved via a Github issue or the forum? 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, and
here are tips on formatting docstrings.
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.

HuggingFaceDocBuilderDev · 2024-07-12T22:22:43Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

qubvel · 2024-07-15T11:09:00Z

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

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

qubvel · 2024-07-15T15:36:25Z

Note:

CLIPTextModel and CLIPModel as well, are not supposed to be used with left padding:

text model pooled output extracts the last relevant token by the index of the first pad token
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 Fix MusicGen SDPA #31208, SDPA produces NaN with padding mask pytorch/pytorch#103749)

qubvel · 2024-07-15T16:43:59Z

Tests on SSH runner

Flash attention

SDPA

qubvel · 2024-07-15T16:52:30Z

tests/models/clip/test_modeling_clip.py

+    @require_torch_sdpa
+    def test_sdpa_can_compile_dynamic(self):
+        self.skipTest(reason="CLIP model can't be compiled dynamic, error in clip_loss`")


Skip the compile dynamic test, I found that compilation failed due to clip_loss, I don't know how critical this is, but I didn't find a quick fix, so just skipped for now.

Only relevant for CLIPModel

Did that happen because it outputs tuples or was it something else?

No, loss outputs tensor. With dynamic=None it works fine, but with dynamic=True I'm getting the following error in the test:

_call_with_frames_removed: block: [0,0,0], thread: [31,0,0] Assertion 'index out of bounds: tmp4 < ks0' failed.

Oh okay the error message itself doesn't seem to be super helpful. Okay with the current changes.

qubvel · 2024-07-15T17:12:14Z

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

sayakpaul

Thanks, Pavel! My comments are quite minor. Overall, this looks nice.

It's a shame that CLIP doesn't show as much of an improvement SigLIP seems to show.

sayakpaul · 2024-07-16T06:05:55Z

docs/source/en/model_doc/clip.md

+
+### 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)):


Just out of curiosity -- do these trends roughly hold for other checkpoints too? Anyway, you picked the most popular checkpoint I think since it's used in the diffusion community quite heavily.

roughly holds for sdpa, however FA2 is worse for "base" checkpoint for CLIPModel

“openai/clip-vit-base-patch32”

CLIPTextModel

Num text labels Eager (s/iter) FA2 (s/iter) FA2 speedup SDPA (s/iter) SDPA speedup

4 0.007 0.011 0.674 0.006 1.173

16 0.007 0.011 0.648 0.006 1.17

64 0.015 0.017 0.855 0.013 1.112

128 0.03 0.029 1.022 0.027 1.121

256 0.06 0.054 1.106 0.053 1.121

CLIPVisionModel

Image batch size Eager (s/iter) FA2 (s/iter) FA2 speedup SDPA (s/iter) SDPA speedup

1 0.007 0.005 1.224 0.005 1.307

4 0.007 0.005 1.228 0.005 1.287

16 0.007 0.006 1.23 0.005 1.316

32 0.009 0.009 1.023 0.009 1.031

64 0.018 0.017 1.076 0.017 1.082

128 0.035 0.032 1.09 0.032 1.095

CLIPModel

Image batch size Num text labels Eager (s/iter) FA2 (s/iter) FA2 speedup SDPA (s/iter) SDPA speedup

1 4 0.015 0.017 0.866 0.012 1.23

1 16 0.015 0.017 0.86 0.012 1.221

1 64 0.016 0.023 0.705 0.015 1.105

4 4 0.015 0.017 0.871 0.012 1.241

4 16 0.015 0.017 0.866 0.012 1.231

4 64 0.017 0.023 0.741 0.015 1.106

16 4 0.015 0.017 0.869 0.012 1.226

16 16 0.015 0.017 0.878 0.012 1.234

16 64 0.02 0.023 0.868 0.018 1.099

32 4 0.015 0.02 0.746 0.012 1.221

32 16 0.015 0.02 0.756 0.013 1.215

32 64 0.024 0.027 0.911 0.022 1.084

FA2 results are super interesting!

Okay, in general, SDPA seems to be still producing better speedups.

sayakpaul · 2024-07-16T06:14:15Z

docs/source/en/model_doc/clip.md

+#### CLIPTextModel
+
+|   Num text labels |   Eager (s/iter) |   FA2 (s/iter) |   FA2 speedup |   SDPA (s/iter) |   SDPA speedup |
+|------------------:|-----------------:|---------------:|--------------:|----------------:|---------------:|
+|                 4 |            0.007 |          0.011 |         0.677 |           0.006 |          1.175 |
+|                16 |            0.007 |          0.013 |         0.577 |           0.007 |          1.056 |
+|                64 |            0.029 |          0.03  |         0.966 |           0.026 |          1.094 |
+|               128 |            0.052 |          0.049 |         1.069 |           0.047 |          1.108 |
+|               256 |            0.103 |          0.092 |         1.115 |           0.092 |          1.109 |


It is surprising to see that the SigLIP variant (albeit not exactly identical architecturally) shows somewhat consistent speedups across various settings, but CLIP seems to be notorious about this.

The SiglipTextModel also does not gain significant benefits from using sdpa/flash attention. However, the SiglipVisionModel is more resource-intensive, so the SiglipModel performance improves due to the speedup in the vision model.

For example:

CLIPTextModel (64, eager): 0.015 sec
CLIPVisionModel (64, eager): 0.018 sec

SiglipTextModel (64, eager): 0.091 sec
SiglipVisionModel (64, eager): 1.720 sec

Benchmark for Siglip models

"google/siglip-so400m-patch14-384"

SiglipModel

Image batch size Num text labels Eager (s/iter) FA2 (s/iter) FA2 speedup SDPA (s/iter) SDPA speedup

1 4 0.04 0.03 1.359 0.03 1.348

1 16 0.058 0.046 1.248 0.046 1.246

1 64 0.123 0.104 1.182 0.104 1.18

4 4 0.128 0.083 1.549 0.083 1.547

4 16 0.145 0.099 1.469 0.099 1.467

4 64 0.21 0.157 1.343 0.157 1.342

16 4 0.459 0.285 1.61 0.285 1.608

16 16 0.476 0.301 1.58 0.301 1.579

16 64 0.541 0.359 1.508 0.359 1.506

32 4 0.889 0.543 1.638 0.544 1.636

32 16 0.907 0.559 1.622 0.56 1.62

32 64 0.972 0.617 1.575 0.618 1.574

SiglipTextModel

Num text labels Eager (s/iter) FA2 (s/iter) FA2 speedup SDPA (s/iter) SDPA speedup

4 0.018 0.015 1.218 0.014 1.267

16 0.026 0.024 1.071 0.024 1.071

64 0.091 0.082 1.111 0.082 1.111

128 0.18 0.164 1.099 0.164 1.099

256 0.364 0.333 1.093 0.333 1.093

SiglipVisionModel

Image batch size Eager (s/iter) FA2 (s/iter) FA2 speedup SDPA (s/iter) SDPA speedup

1 0.032 0.022 1.449 0.022 1.444

4 0.119 0.075 1.599 0.075 1.597

16 0.45 0.277 1.624 0.277 1.623

32 0.88 0.535 1.645 0.536 1.643

64 1.723 1.03 1.673 1.035 1.664

Hmm yeah. Seems like apart from the vision tower, rest of the speedups are quite identical to SigLIP.

src/transformers/models/clip/modeling_clip.py

tests/models/clip/test_modeling_clip.py

sayakpaul · 2024-07-16T06:29:06Z

tests/models/clip/test_modeling_clip.py

+    @require_torch_sdpa
+    def test_sdpa_can_compile_dynamic(self):
+        self.skipTest(reason="CLIP model can't be compiled dynamic, error in clip_loss`")


Did that happen because it outputs tuples or was it something else?

qubvel · 2024-07-16T21:26:40Z

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 #32012

commit 102842cd477219b9f9bcb23a0bca3a8b92bd732f Author: Pavel Iakubovskii <[email protected]> Date: Fri Jul 12 18:23:52 2024 +0000 Add model-specific sdpa tests commit 60e4c88581abf89ec098da84ed8e92aa904c997d Author: Pavel Iakubovskii <[email protected]> Date: Fri Jul 12 18:20:53 2024 +0000 Add fallback to eager (expensive operation) commit c29033d30e7ffde4327e8a15cbbc6bee37546f80 Author: Pavel Iakubovskii <[email protected]> Date: Thu Jul 11 17:09:55 2024 +0000 Fix attn_implementation propagation commit 783aed05f0f38cb2f99e758f81db6838ac55b9f8 Author: sayakpaul <[email protected]> Date: Sat May 25 09:05:27 2024 +0530 style commit e77e703ca75d00447cda277eca6b886cd32bddc0 Author: sayakpaul <[email protected]> Date: Sat May 25 09:04:57 2024 +0530 add comment to explain why I had to touch forbidden codebase. commit ab9d8849758e7773a31778ccba71588d18552623 Author: sayakpaul <[email protected]> Date: Sat May 25 09:03:02 2024 +0530 fix: flax attribute access. commit c570fc0abf9d1bd58c291aae3c7e384f995996d2 Author: sayakpaul <[email protected]> Date: Sat May 25 08:23:54 2024 +0530 fix tensorflow attribute name. commit 32c812871cfdb268d8a6e3e2c61c5c925c8ed47e Author: sayakpaul <[email protected]> Date: Sat May 25 07:57:10 2024 +0530 fix attribute access. commit 4f41a0138b6c417aed9c9332278f8bcd979cb7c2 Author: sayakpaul <[email protected]> Date: Sat May 25 07:44:02 2024 +0530 _from_config. commit 35aed64ff602422adcf41d7f677a0a24bd9eccae Author: sayakpaul <[email protected]> Date: Fri May 24 18:46:52 2024 +0530 propagation of attn_implementation. commit 4c25c19845438b1dc1d35a5adf9436151c8c5940 Author: sayakpaul <[email protected]> Date: Fri May 24 09:24:36 2024 +0530 style again commit 5f7dc5c5015c0f8116408f737e8c318d1802c80c Author: sayakpaul <[email protected]> Date: Fri May 24 09:19:05 2024 +0530 use from_config. commit b70c409956d0359fa6ae5372275d2a20ba7e3389 Author: sayakpaul <[email protected]> Date: Fri May 24 09:13:43 2024 +0530 quality commit a7b63beff53d0fc754c6564e2a7b51731ddee49d Author: sayakpaul <[email protected]> Date: Fri May 10 14:35:10 2024 +0200 add benchmark numbers commit 455b0eaea50862b8458c8f422b60fe60ae40fdcb Author: sayakpaul <[email protected]> Date: Fri May 10 13:50:16 2024 +0200 Revert "reflect feedback more" This reverts commit dc123e7. commit ca674829d28787349c2a9593a14e0f1d41f04ea4 Author: sayakpaul <[email protected]> Date: Fri May 10 13:50:05 2024 +0200 Revert "fix" This reverts commit 37a1cb3. commit fab2dd8576c099eb1a3464958cb206a664d28247 Author: sayakpaul <[email protected]> Date: Fri May 10 13:47:46 2024 +0200 fix commit fbc6ae50fd6f2d36294d31e191761631b701d696 Author: sayakpaul <[email protected]> Date: Fri May 10 13:38:30 2024 +0200 reflect feedback more commit 87245bb020b2d60a89afe318a951df0159404fc9 Author: sayakpaul <[email protected]> Date: Fri May 3 08:54:34 2024 +0530 fixes commit 1057cc26390ee839251e7f8b3326c4207595fb23 Author: sayakpaul <[email protected]> Date: Fri May 3 07:49:03 2024 +0530 don't explicit set attn_implementation in tests commit e33f75916fc8a99f516b1cf449dbbe9d3aabda81 Author: sayakpaul <[email protected]> Date: Fri May 3 07:43:54 2024 +0530 explicitly override attn_implementation in the towers. commit 4cf41cb1bc885c39df7cb8f2a0694ebf23299235 Author: sayakpaul <[email protected]> Date: Fri May 3 07:38:42 2024 +0530 import in one-line. commit f2cc447ae9e74ccfacb448140cdf88259d4afc8c Author: sayakpaul <[email protected]> Date: Fri May 3 07:34:58 2024 +0530 move sdpa mention to usage tips. commit 92884766c64dbb456926a3a84dd427be1349fa95 Author: sayakpaul <[email protected]> Date: Mon Apr 29 10:58:26 2024 +0530 fix: memory allocation problem. commit d7ffbbfe12f7750b7d0a361420f35c13e0ea787d Author: sayakpaul <[email protected]> Date: Mon Apr 29 09:56:59 2024 +0530 fix-copies commit 8dfc3731cedd02e36acd3fe56bb2e6d61efd25d8 Author: sayakpaul <[email protected]> Date: Fri Apr 26 20:16:12 2024 +0530 address arthur's comments. commit d2ed7b4ce4ff15ae9aa4d3d0500f1544e3dcd9e9 Author: Sayak Paul <[email protected]> Date: Fri Apr 26 20:08:15 2024 +0530 Apply suggestions from code review Co-authored-by: Arthur <[email protected]> commit 46e04361f37ded5c522ff05e9f725b9f82dce40e Author: sayakpaul <[email protected]> Date: Wed Apr 24 09:55:27 2024 +0530 add to docs. commit 831629158ad40d34d8983f209afb2740ba041af2 Author: sayakpaul <[email protected]> Date: Wed Apr 24 09:33:10 2024 +0530 styling.g commit d263a119c77314250f4b4c8469caf42559197f22 Author: sayakpaul <[email protected]> Date: Wed Apr 24 09:15:20 2024 +0530 up commit d44f9d3d7633d4c241a737a1bc317f791f6aedb3 Author: sayakpaul <[email protected]> Date: Tue Apr 23 18:40:42 2024 +0530 handle causal and attention mask commit 122f1d60153df6666b634a94e38d073f3f260926 Author: sayakpaul <[email protected]> Date: Tue Apr 23 15:18:21 2024 +0530 test fixes. commit 4382d8cff6fa1dee5dbcf0d06b3e2841231e36f5 Author: sayakpaul <[email protected]> Date: Tue Apr 23 09:39:25 2024 +0530 fix: scaling inside sdpa. commit 0f629989efc48b7315cf19405a81e02955efe7e5 Author: Sayak Paul <[email protected]> Date: Tue Apr 23 08:14:58 2024 +0530 Update src/transformers/models/clip/modeling_clip.py Co-authored-by: Arthur <[email protected]> commit 14367316877dc27ea40f767ad1aee38bbc97e4ce Author: sayakpaul <[email protected]> Date: Mon Apr 22 16:21:36 2024 +0530 add: sdpa support to clip.

qubvel · 2024-07-17T14:01:45Z

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

amyeroberts

Great piece of work - thanks for adding this!

amyeroberts · 2024-07-17T19:33:54Z

docs/source/en/model_doc/clip.md

+
+### 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)):


FA2 results are super interesting!

amyeroberts · 2024-07-17T19:34:36Z

docs/source/en/model_doc/clip.md

+
+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


These tables are beautiful - thanks for taking the time to get such detailed benchmarks. It's incredible valuable for important checkpoints like CLIP ❤️

sayakpaul · 2024-07-18T05:00:29Z

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

* Squashed commit of the following: commit 102842cd477219b9f9bcb23a0bca3a8b92bd732f Author: Pavel Iakubovskii <[email protected]> Date: Fri Jul 12 18:23:52 2024 +0000 Add model-specific sdpa tests commit 60e4c88581abf89ec098da84ed8e92aa904c997d Author: Pavel Iakubovskii <[email protected]> Date: Fri Jul 12 18:20:53 2024 +0000 Add fallback to eager (expensive operation) commit c29033d30e7ffde4327e8a15cbbc6bee37546f80 Author: Pavel Iakubovskii <[email protected]> Date: Thu Jul 11 17:09:55 2024 +0000 Fix attn_implementation propagation commit 783aed05f0f38cb2f99e758f81db6838ac55b9f8 Author: sayakpaul <[email protected]> Date: Sat May 25 09:05:27 2024 +0530 style commit e77e703ca75d00447cda277eca6b886cd32bddc0 Author: sayakpaul <[email protected]> Date: Sat May 25 09:04:57 2024 +0530 add comment to explain why I had to touch forbidden codebase. commit ab9d8849758e7773a31778ccba71588d18552623 Author: sayakpaul <[email protected]> Date: Sat May 25 09:03:02 2024 +0530 fix: flax attribute access. commit c570fc0abf9d1bd58c291aae3c7e384f995996d2 Author: sayakpaul <[email protected]> Date: Sat May 25 08:23:54 2024 +0530 fix tensorflow attribute name. commit 32c812871cfdb268d8a6e3e2c61c5c925c8ed47e Author: sayakpaul <[email protected]> Date: Sat May 25 07:57:10 2024 +0530 fix attribute access. commit 4f41a0138b6c417aed9c9332278f8bcd979cb7c2 Author: sayakpaul <[email protected]> Date: Sat May 25 07:44:02 2024 +0530 _from_config. commit 35aed64ff602422adcf41d7f677a0a24bd9eccae Author: sayakpaul <[email protected]> Date: Fri May 24 18:46:52 2024 +0530 propagation of attn_implementation. commit 4c25c19845438b1dc1d35a5adf9436151c8c5940 Author: sayakpaul <[email protected]> Date: Fri May 24 09:24:36 2024 +0530 style again commit 5f7dc5c5015c0f8116408f737e8c318d1802c80c Author: sayakpaul <[email protected]> Date: Fri May 24 09:19:05 2024 +0530 use from_config. commit b70c409956d0359fa6ae5372275d2a20ba7e3389 Author: sayakpaul <[email protected]> Date: Fri May 24 09:13:43 2024 +0530 quality commit a7b63beff53d0fc754c6564e2a7b51731ddee49d Author: sayakpaul <[email protected]> Date: Fri May 10 14:35:10 2024 +0200 add benchmark numbers commit 455b0eaea50862b8458c8f422b60fe60ae40fdcb Author: sayakpaul <[email protected]> Date: Fri May 10 13:50:16 2024 +0200 Revert "reflect feedback more" This reverts commit dc123e7. commit ca674829d28787349c2a9593a14e0f1d41f04ea4 Author: sayakpaul <[email protected]> Date: Fri May 10 13:50:05 2024 +0200 Revert "fix" This reverts commit 37a1cb3. commit fab2dd8576c099eb1a3464958cb206a664d28247 Author: sayakpaul <[email protected]> Date: Fri May 10 13:47:46 2024 +0200 fix commit fbc6ae50fd6f2d36294d31e191761631b701d696 Author: sayakpaul <[email protected]> Date: Fri May 10 13:38:30 2024 +0200 reflect feedback more commit 87245bb020b2d60a89afe318a951df0159404fc9 Author: sayakpaul <[email protected]> Date: Fri May 3 08:54:34 2024 +0530 fixes commit 1057cc26390ee839251e7f8b3326c4207595fb23 Author: sayakpaul <[email protected]> Date: Fri May 3 07:49:03 2024 +0530 don't explicit set attn_implementation in tests commit e33f75916fc8a99f516b1cf449dbbe9d3aabda81 Author: sayakpaul <[email protected]> Date: Fri May 3 07:43:54 2024 +0530 explicitly override attn_implementation in the towers. commit 4cf41cb1bc885c39df7cb8f2a0694ebf23299235 Author: sayakpaul <[email protected]> Date: Fri May 3 07:38:42 2024 +0530 import in one-line. commit f2cc447ae9e74ccfacb448140cdf88259d4afc8c Author: sayakpaul <[email protected]> Date: Fri May 3 07:34:58 2024 +0530 move sdpa mention to usage tips. commit 92884766c64dbb456926a3a84dd427be1349fa95 Author: sayakpaul <[email protected]> Date: Mon Apr 29 10:58:26 2024 +0530 fix: memory allocation problem. commit d7ffbbfe12f7750b7d0a361420f35c13e0ea787d Author: sayakpaul <[email protected]> Date: Mon Apr 29 09:56:59 2024 +0530 fix-copies commit 8dfc3731cedd02e36acd3fe56bb2e6d61efd25d8 Author: sayakpaul <[email protected]> Date: Fri Apr 26 20:16:12 2024 +0530 address arthur's comments. commit d2ed7b4ce4ff15ae9aa4d3d0500f1544e3dcd9e9 Author: Sayak Paul <[email protected]> Date: Fri Apr 26 20:08:15 2024 +0530 Apply suggestions from code review Co-authored-by: Arthur <[email protected]> commit 46e04361f37ded5c522ff05e9f725b9f82dce40e Author: sayakpaul <[email protected]> Date: Wed Apr 24 09:55:27 2024 +0530 add to docs. commit 831629158ad40d34d8983f209afb2740ba041af2 Author: sayakpaul <[email protected]> Date: Wed Apr 24 09:33:10 2024 +0530 styling.g commit d263a119c77314250f4b4c8469caf42559197f22 Author: sayakpaul <[email protected]> Date: Wed Apr 24 09:15:20 2024 +0530 up commit d44f9d3d7633d4c241a737a1bc317f791f6aedb3 Author: sayakpaul <[email protected]> Date: Tue Apr 23 18:40:42 2024 +0530 handle causal and attention mask commit 122f1d60153df6666b634a94e38d073f3f260926 Author: sayakpaul <[email protected]> Date: Tue Apr 23 15:18:21 2024 +0530 test fixes. commit 4382d8cff6fa1dee5dbcf0d06b3e2841231e36f5 Author: sayakpaul <[email protected]> Date: Tue Apr 23 09:39:25 2024 +0530 fix: scaling inside sdpa. commit 0f629989efc48b7315cf19405a81e02955efe7e5 Author: Sayak Paul <[email protected]> Date: Tue Apr 23 08:14:58 2024 +0530 Update src/transformers/models/clip/modeling_clip.py Co-authored-by: Arthur <[email protected]> commit 14367316877dc27ea40f767ad1aee38bbc97e4ce Author: sayakpaul <[email protected]> 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

SangbumChoi · 2024-07-26T07:51:58Z

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.

sayakpaul · 2024-07-26T07:53:31Z

Cc @DN6 @yiyixuxu ^

SangbumChoi · 2024-07-26T07:58:27Z

#28610

Fixes missing TE layers resulting from the addition of SDPA/FA2 to CLIP: huggingface/transformers#31940

qubvel added Vision run-slow labels Jul 12, 2024

qubvel commented Jul 15, 2024

View reviewed changes

sayakpaul approved these changes Jul 16, 2024

View reviewed changes

qubvel requested a review from amyeroberts July 16, 2024 21:32

qubvel added 19 commits July 17, 2024 09:15

Remove fallback for empty attention mask (expensive operation)

cca2ec6

Fix typing in copies

f3d8024

Add flash attention

0bbc912

Add flash attention tests

06c1cea

List CLIP in FA docs

650aaa5

Fix embeddings attributes and tf

fa01bb6

[run-slow] clip

60207fe

Update clip documentation

df54b52

Remove commented code, skip compile dynamic for CLIPModel

e5c8aae

Fix doc

8376f43

Fix doc 2

088a4ec

Remove double transpose

9028968

Add torch version check for contiguous()

ec859fc

Add comment to test mixin

778ef65

Fix copies

1528708

Add comment for mask

de681a5

Update docs

cfc5240

[run-slow] clip

d2328d3

qubvel force-pushed the add-clip-sdpa-fa2 branch from 0b4e63b to d2328d3 Compare July 17, 2024 09:16

amyeroberts approved these changes Jul 17, 2024

View reviewed changes

sayakpaul merged commit 1c37e8c into huggingface:main Jul 18, 2024
22 checks passed

sayakpaul mentioned this pull request Jul 18, 2024

[CLIP] add: sdpa support to clip. #30390

Closed

jambayk mentioned this pull request Jul 24, 2024

StableDiffusionSafetyChecker ignores attn_implementation load kwarg huggingface/diffusers#8957

Open

qubvel mentioned this pull request Jul 25, 2024

Fix attention propagation for vision towers of llava-like models #32221

Closed

atiorh mentioned this pull request Aug 4, 2024

ERROR - converting 'scaled_dot_product_attention' op (located at: 'text_encoder/text_model/encoder/0/self_attn'): apple/ml-stable-diffusion#345

Open

yangtian6781 mentioned this pull request Aug 20, 2024

[Feature] 为什么不用pytorch原生的sdpa，反而用flash attention呢？ OpenGVLab/InternVL#519

Closed

LoganBooker added a commit to LoganBooker/LyCORIS that referenced this pull request Sep 8, 2024

Add CLIPSdpaAttention and CLIPFlashAttention2 to target modules.

a8fbbe2

Fixes missing TE layers resulting from the addition of SDPA/FA2 to CLIP: huggingface/transformers#31940

LoganBooker added a commit to LoganBooker/LyCORIS that referenced this pull request Sep 8, 2024

Add CLIPSdpaAttention and CLIPFlashAttention2 to target modules.

962baed

Fixes missing TE layers resulting from the addition of SDPA/FA2 to CLIP: huggingface/transformers#31940

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `sdpa` and FA2 for CLIP #31940

Add `sdpa` and FA2 for CLIP #31940

qubvel commented Jul 12, 2024 •

edited

Loading

HuggingFaceDocBuilderDev commented Jul 12, 2024

qubvel commented Jul 15, 2024 •

edited

Loading

qubvel commented Jul 15, 2024 •

edited

Loading

qubvel commented Jul 15, 2024

qubvel Jul 15, 2024 •

edited

Loading

sayakpaul Jul 16, 2024

qubvel Jul 16, 2024 •

edited

Loading

sayakpaul Jul 18, 2024

qubvel commented Jul 15, 2024

sayakpaul left a comment

sayakpaul Jul 16, 2024

qubvel Jul 16, 2024 •

edited

Loading

amyeroberts Jul 17, 2024

sayakpaul Jul 18, 2024

sayakpaul Jul 16, 2024

qubvel Jul 16, 2024 •

edited

Loading

sayakpaul Jul 18, 2024

sayakpaul Jul 16, 2024

qubvel commented Jul 16, 2024

qubvel commented Jul 17, 2024

amyeroberts left a comment

amyeroberts Jul 17, 2024

amyeroberts Jul 17, 2024

sayakpaul commented Jul 18, 2024

SangbumChoi commented Jul 26, 2024 •

edited

Loading

sayakpaul commented Jul 26, 2024

SangbumChoi commented Jul 26, 2024


		### 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)):

Num text labels	Eager (s/iter)	FA2 (s/iter)	FA2 speedup	SDPA (s/iter)	SDPA speedup
4	0.007	0.011	0.674	0.006	1.173
16	0.007	0.011	0.648	0.006	1.17
64	0.015	0.017	0.855	0.013	1.112
128	0.03	0.029	1.022	0.027	1.121
256	0.06	0.054	1.106	0.053	1.121

Image batch size	Eager (s/iter)	FA2 (s/iter)	FA2 speedup	SDPA (s/iter)	SDPA speedup
1	0.007	0.005	1.224	0.005	1.307
4	0.007	0.005	1.228	0.005	1.287
16	0.007	0.006	1.23	0.005	1.316
32	0.009	0.009	1.023	0.009	1.031
64	0.018	0.017	1.076	0.017	1.082
128	0.035	0.032	1.09	0.032	1.095

Image batch size	Num text labels	Eager (s/iter)	FA2 (s/iter)	FA2 speedup	SDPA (s/iter)	SDPA speedup
1	4	0.015	0.017	0.866	0.012	1.23
1	16	0.015	0.017	0.86	0.012	1.221
1	64	0.016	0.023	0.705	0.015	1.105
4	4	0.015	0.017	0.871	0.012	1.241
4	16	0.015	0.017	0.866	0.012	1.231
4	64	0.017	0.023	0.741	0.015	1.106
16	4	0.015	0.017	0.869	0.012	1.226
16	16	0.015	0.017	0.878	0.012	1.234
16	64	0.02	0.023	0.868	0.018	1.099
32	4	0.015	0.02	0.746	0.012	1.221
32	16	0.015	0.02	0.756	0.013	1.215
32	64	0.024	0.027	0.911	0.022	1.084

Image batch size	Num text labels	Eager (s/iter)	FA2 (s/iter)	FA2 speedup	SDPA (s/iter)	SDPA speedup
1	4	0.04	0.03	1.359	0.03	1.348
1	16	0.058	0.046	1.248	0.046	1.246
1	64	0.123	0.104	1.182	0.104	1.18
4	4	0.128	0.083	1.549	0.083	1.547
4	16	0.145	0.099	1.469	0.099	1.467
4	64	0.21	0.157	1.343	0.157	1.342
16	4	0.459	0.285	1.61	0.285	1.608
16	16	0.476	0.301	1.58	0.301	1.579
16	64	0.541	0.359	1.508	0.359	1.506
32	4	0.889	0.543	1.638	0.544	1.636
32	16	0.907	0.559	1.622	0.56	1.62
32	64	0.972	0.617	1.575	0.618	1.574

Image batch size	Eager (s/iter)	FA2 (s/iter)	FA2 speedup	SDPA (s/iter)	SDPA speedup
1	0.032	0.022	1.449	0.022	1.444
4	0.119	0.075	1.599	0.075	1.597
16	0.45	0.277	1.624	0.277	1.623
32	0.88	0.535	1.645	0.536	1.643
64	1.723	1.03	1.673	1.035	1.664


		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

Add sdpa and FA2 for CLIP #31940

Add sdpa and FA2 for CLIP #31940

Conversation

qubvel commented Jul 12, 2024 • edited Loading

What does this PR do?

Before submitting

Who can review?

HuggingFaceDocBuilderDev commented Jul 12, 2024

qubvel commented Jul 15, 2024 • edited Loading

CLIPModel

CLIPTextModel

CLIPVisionModel

qubvel commented Jul 15, 2024 • edited Loading

qubvel commented Jul 15, 2024

Tests on SSH runner

Flash attention

SDPA

qubvel Jul 15, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

qubvel Jul 16, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

qubvel commented Jul 15, 2024

sayakpaul left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

qubvel Jul 16, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

qubvel Jul 16, 2024 • edited Loading

Choose a reason for hiding this comment

Benchmark for Siglip models

Choose a reason for hiding this comment

Choose a reason for hiding this comment

qubvel commented Jul 16, 2024

qubvel commented Jul 17, 2024

amyeroberts left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sayakpaul commented Jul 18, 2024

SangbumChoi commented Jul 26, 2024 • edited Loading

sayakpaul commented Jul 26, 2024

SangbumChoi commented Jul 26, 2024

Add `sdpa` and FA2 for CLIP #31940

Add `sdpa` and FA2 for CLIP #31940

qubvel commented Jul 12, 2024 •

edited

Loading

qubvel commented Jul 15, 2024 •

edited

Loading

qubvel commented Jul 15, 2024 •

edited

Loading

qubvel Jul 15, 2024 •

edited

Loading

qubvel Jul 16, 2024 •

edited

Loading

qubvel Jul 16, 2024 •

edited

Loading

qubvel Jul 16, 2024 •

edited

Loading

SangbumChoi commented Jul 26, 2024 •

edited

Loading