-
Notifications
You must be signed in to change notification settings - Fork 31.2k
Add sdpa and FA2 for CLIP
#31940
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add sdpa and FA2 for CLIP
#31940
Conversation
|
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. |
|
Benchmarking results for the CLIP models ( Code: CLIPModel
CLIPTextModel
CLIPVisionModel
|
|
Note: CLIPTextModel and CLIPModel as well, are not supposed to be used with left padding:
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Did that happen because it outputs tuples or was it something else?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Oh okay the error message itself doesn't seem to be super helpful. Okay with the current changes.
|
@sayakpaul in case you have time, can you please review this PR? |
sayakpaul
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
docs/source/en/model_doc/clip.md
Outdated
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
FA2 results are super interesting!
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Okay, in general, SDPA seems to be still producing better speedups.
docs/source/en/model_doc/clip.md
Outdated
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hmm yeah. Seems like apart from the vision tower, rest of the speedups are quite identical to SigLIP.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Did that happen because it outputs tuples or was it something else?
|
Slow test error: addressed in #32012 |
commit 102842cd477219b9f9bcb23a0bca3a8b92bd732f
Author: Pavel Iakubovskii <qubvel@gmail.com>
Date: Fri Jul 12 18:23:52 2024 +0000
Add model-specific sdpa tests
commit 60e4c88581abf89ec098da84ed8e92aa904c997d
Author: Pavel Iakubovskii <qubvel@gmail.com>
Date: Fri Jul 12 18:20:53 2024 +0000
Add fallback to eager (expensive operation)
commit c29033d30e7ffde4327e8a15cbbc6bee37546f80
Author: Pavel Iakubovskii <qubvel@gmail.com>
Date: Thu Jul 11 17:09:55 2024 +0000
Fix attn_implementation propagation
commit 783aed05f0f38cb2f99e758f81db6838ac55b9f8
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Sat May 25 09:05:27 2024 +0530
style
commit e77e703ca75d00447cda277eca6b886cd32bddc0
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Sat May 25 09:04:57 2024 +0530
add comment to explain why I had to touch forbidden codebase.
commit ab9d8849758e7773a31778ccba71588d18552623
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Sat May 25 09:03:02 2024 +0530
fix: flax attribute access.
commit c570fc0abf9d1bd58c291aae3c7e384f995996d2
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Sat May 25 08:23:54 2024 +0530
fix tensorflow attribute name.
commit 32c812871cfdb268d8a6e3e2c61c5c925c8ed47e
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Sat May 25 07:57:10 2024 +0530
fix attribute access.
commit 4f41a0138b6c417aed9c9332278f8bcd979cb7c2
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Sat May 25 07:44:02 2024 +0530
_from_config.
commit 35aed64ff602422adcf41d7f677a0a24bd9eccae
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Fri May 24 18:46:52 2024 +0530
propagation of attn_implementation.
commit 4c25c19845438b1dc1d35a5adf9436151c8c5940
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Fri May 24 09:24:36 2024 +0530
style again
commit 5f7dc5c5015c0f8116408f737e8c318d1802c80c
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Fri May 24 09:19:05 2024 +0530
use from_config.
commit b70c409956d0359fa6ae5372275d2a20ba7e3389
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Fri May 24 09:13:43 2024 +0530
quality
commit a7b63beff53d0fc754c6564e2a7b51731ddee49d
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Fri May 10 14:35:10 2024 +0200
add benchmark numbers
commit 455b0eaea50862b8458c8f422b60fe60ae40fdcb
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Fri May 10 13:50:16 2024 +0200
Revert "reflect feedback more"
This reverts commit dc123e7.
commit ca674829d28787349c2a9593a14e0f1d41f04ea4
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Fri May 10 13:50:05 2024 +0200
Revert "fix"
This reverts commit 37a1cb3.
commit fab2dd8576c099eb1a3464958cb206a664d28247
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Fri May 10 13:47:46 2024 +0200
fix
commit fbc6ae50fd6f2d36294d31e191761631b701d696
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Fri May 10 13:38:30 2024 +0200
reflect feedback more
commit 87245bb020b2d60a89afe318a951df0159404fc9
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Fri May 3 08:54:34 2024 +0530
fixes
commit 1057cc26390ee839251e7f8b3326c4207595fb23
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Fri May 3 07:49:03 2024 +0530
don't explicit set attn_implementation in tests
commit e33f75916fc8a99f516b1cf449dbbe9d3aabda81
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Fri May 3 07:43:54 2024 +0530
explicitly override attn_implementation in the towers.
commit 4cf41cb1bc885c39df7cb8f2a0694ebf23299235
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Fri May 3 07:38:42 2024 +0530
import in one-line.
commit f2cc447ae9e74ccfacb448140cdf88259d4afc8c
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Fri May 3 07:34:58 2024 +0530
move sdpa mention to usage tips.
commit 92884766c64dbb456926a3a84dd427be1349fa95
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Mon Apr 29 10:58:26 2024 +0530
fix: memory allocation problem.
commit d7ffbbfe12f7750b7d0a361420f35c13e0ea787d
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Mon Apr 29 09:56:59 2024 +0530
fix-copies
commit 8dfc3731cedd02e36acd3fe56bb2e6d61efd25d8
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Fri Apr 26 20:16:12 2024 +0530
address arthur's comments.
commit d2ed7b4ce4ff15ae9aa4d3d0500f1544e3dcd9e9
Author: Sayak Paul <spsayakpaul@gmail.com>
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 <spsayakpaul@gmail.com>
Date: Wed Apr 24 09:55:27 2024 +0530
add to docs.
commit 831629158ad40d34d8983f209afb2740ba041af2
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Wed Apr 24 09:33:10 2024 +0530
styling.g
commit d263a119c77314250f4b4c8469caf42559197f22
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Wed Apr 24 09:15:20 2024 +0530
up
commit d44f9d3d7633d4c241a737a1bc317f791f6aedb3
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Tue Apr 23 18:40:42 2024 +0530
handle causal and attention mask
commit 122f1d60153df6666b634a94e38d073f3f260926
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Tue Apr 23 15:18:21 2024 +0530
test fixes.
commit 4382d8cff6fa1dee5dbcf0d06b3e2841231e36f5
Author: sayakpaul <spsayakpaul@gmail.com>
Date: Tue Apr 23 09:39:25 2024 +0530
fix: scaling inside sdpa.
commit 0f629989efc48b7315cf19405a81e02955efe7e5
Author: Sayak Paul <spsayakpaul@gmail.com>
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 <spsayakpaul@gmail.com>
Date: Mon Apr 22 16:21:36 2024 +0530
add: sdpa support to clip.
0b4e63b to
d2328d3
Compare
|
@amyeroberts whenever you have time can you please review this PR 🙂 |
amyeroberts
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Great piece of work - thanks for adding this!
docs/source/en/model_doc/clip.md
Outdated
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
FA2 results are super interesting!
|
|
||
| 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 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
These tables are beautiful - thanks for taking the time to get such detailed benchmarks. It's incredible valuable for important checkpoints like CLIP ❤️
|
Great work, Pavel. I think this is ready to be shipped. |
|
FYI) @sayakpaul I think you are the managing |
Fixes missing TE layers resulting from the addition of SDPA/FA2 to CLIP: huggingface/transformers#31940
Fixes missing TE layers resulting from the addition of SDPA/FA2 to CLIP: huggingface/transformers#31940




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:
sdpasupport for SigLIP #31499Before submitting
Pull Request section?
to it if that's the case.
documentation guidelines, and
here are tips on formatting docstrings.
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.