Skip to content
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.

Sign up for GitHub

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

Jump to bottom

[feat] Multihead dispatch benchmark #273

Merged
merged 3 commits into from
Apr 13, 2022
Merged

[feat] Multihead dispatch benchmark #273

merged 3 commits into from
Apr 13, 2022

Conversation

dianaml0
Copy link
Contributor

What does this PR do?

Benchmarks xFormers MultiHeadDispatch again Pytorch's version.

Before submitting

  • Did you have fun?
    • Make sure you had fun coding 🙃
  • Did you read the contributor guideline?
  • Was this discussed/approved via a Github issue? (no need for typos, doc improvements)
    • N/A
  • Did you make sure to update the docs?
    • N/A
  • Did you write any new necessary tests?
    • N/A
  • Did you update the changelog? (if needed)
    • N/A

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Apr 13, 2022
@dianaml0
Copy link
Contributor Author

--- Type: torch.float16 ---

Units: runtime in ms, lower is better B=8, M=384, K=128, N_HEADS=4 B=8, M=784, K=512, N_HEADS=4 B=4, M=1024, K=768, N_HEADS=4 B=4, M=2048, K=1024, N_HEADS=4 B=2, M=2048, K=2048, N_HEADS=4 B=2, M=2048, K=4096, N_HEADS=4 B=2, M=4096, K=4096, N_HEADS=4 B=1, M=2048, K=12288, N_HEADS=4
torch - fw 0.20 0.97 0.99 3.31 3.77 9.96 24.54 33.71
xf - fw 0.14 0.85 0.90 2.99 3.56 9.91 24.37 33.61

--- Type: torch.float32 ---

Units: runtime in ms, lower is better B=8, M=384, K=128, N_HEADS=4 B=8, M=784, K=512, N_HEADS=4 B=4, M=1024, K=768, N_HEADS=4 B=4, M=2048, K=1024, N_HEADS=4 B=2, M=2048, K=2048, N_HEADS=4 B=2, M=2048, K=4096, N_HEADS=4 B=2, M=4096, K=4096, N_HEADS=4 B=1, M=2048, K=12288, N_HEADS=4
torch - fw 0.36 2.72 3.22 12.96 17.42 51.51 126.85 191.88
xf - fw 0.30 2.53 3.09 12.19 16.42 50.81 125.75 192.74

--- Type: torch.float16 ---

Units: runtime in ms, lower is better B=8, M=384, K=128, N_HEADS=4 B=8, M=784, K=512, N_HEADS=4 B=4, M=1024, K=768, N_HEADS=4 B=4, M=2048, K=1024, N_HEADS=4 B=2, M=2048, K=2048, N_HEADS=4 B=2, M=2048, K=4096, N_HEADS=4 B=2, M=4096, K=4096, N_HEADS=4 B=1, M=2048, K=12288, N_HEADS=4
torch - fw+bw 0.58 2.72 2.82 9.10 10.76 28.95 70.05 103.40
xf - fw+bw 0.81 2.28 2.49 8.15 9.89 28.37 69.78 102.42

--- Type: torch.float32 ---

Units: runtime in ms, lower is better B=8, M=384, K=128, N_HEADS=4 B=8, M=784, K=512, N_HEADS=4 B=4, M=1024, K=768, N_HEADS=4 B=4, M=2048, K=1024, N_HEADS=4 B=2, M=2048, K=2048, N_HEADS=4 B=2, M=2048, K=4096, N_HEADS=4 B=2, M=4096, K=4096, N_HEADS=4 B=1, M=2048, K=12288, N_HEADS=4
torch - fw+bw 1.21 7.71 9.55 37.10 51.18 158.82 381.57 607.77
xf - fw+bw 0.93 7.15 8.96 35.21 49.48 154.06 381.99 605.88

@dianaml0
Copy link
Contributor Author

Screen Shot 2022-04-13 at 2 46 10 PM
Screen Shot 2022-04-13 at 2 46 02 PM

@dianaml0 dianaml0 requested review from blefaudeux and fmassa April 13, 2022 18:47
blefaudeux
blefaudeux reviewed Apr 13, 2022
View reviewed changes
HOWTO.md
@@ -238,7 +238,6 @@ attention = BlockSparseAttention(layout=causal_layout, block_size=BLOCK_SIZE, dr
# "multi_head" will be responsible for the forward
multi_head = (
MultiHeadDispatch(
seq_len=SEQ,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nice catch, thanks

blefaudeux
blefaudeux reviewed Apr 13, 2022
View reviewed changes
xformers/benchmarks/benchmark_multi_head_dispatch.py Outdated
N_HEADS = [4]


def bench_multihead_dispatch(backward: bool):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

super clean, nice

blefaudeux
blefaudeux approved these changes Apr 13, 2022
View reviewed changes
Copy link
Contributor

@blefaudeux blefaudeux left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

perfect, thanks @dianaml0 ! The results are very interesting, I for one thought that we could be a little slow in this particular spot, does not seem to be the case ?

blefaudeux
blefaudeux reviewed Apr 13, 2022
View reviewed changes
xformers/benchmarks/benchmark_multi_head_dispatch.py Outdated
)

def torch_mha():
y, _ = torch_multi_head(query=query, key=query, value=query)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: could be nice to bench with and without self attention, I know that self is prevalent in vision but for NLP that would not always be the case. Could be a boolean here, on top of "backwards", self_attention or not ?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good idea, just added this :)

@blefaudeux blefaudeux changed the title multihead dispatch benchmark [feat] Multihead dispatch benchmark Apr 13, 2022
@codecov-commenter
Copy link

Codecov Report

Merging #273 (235cd83) into main (fb7bbcb) will not change coverage.
The diff coverage is n/a.

@@           Coverage Diff           @@
##             main     #273   +/-   ##
=======================================
  Coverage   92.80%   92.80%           
=======================================
  Files          61       61           
  Lines        3363     3363           
=======================================
  Hits         3121     3121           
  Misses        242      242
Flag Coverage Δ
Python 92.80% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update fb7bbcb...235cd83. Read the comment docs.

@dianaml0
Copy link
Contributor Author

Updated results with and without self attention:

--- Type: torch.float16 ---

Units: runtime in ms, lower is better B=8, M=384, K=128, N_HEADS=4 B=8, M=784, K=512, N_HEADS=4 B=4, M=1024, K=768, N_HEADS=4 B=4, M=2048, K=1024, N_HEADS=4 B=2, M=2048, K=2048, N_HEADS=4 B=2, M=2048, K=4096, N_HEADS=4 B=2, M=4096, K=4096, N_HEADS=4 B=1, M=2048, K=12288, N_HEADS=4
torch - fw 0.20 0.96 0.98 3.30 3.77 9.84 24.30 33.42
xf - fw 0.25 0.83 0.88 2.86 3.56 9.72 23.54 33.20

--- Type: torch.float32 ---

Units: runtime in ms, lower is better B=8, M=384, K=128, N_HEADS=4 B=8, M=784, K=512, N_HEADS=4 B=4, M=1024, K=768, N_HEADS=4 B=4, M=2048, K=1024, N_HEADS=4 B=2, M=2048, K=2048, N_HEADS=4 B=2, M=2048, K=4096, N_HEADS=4 B=2, M=4096, K=4096, N_HEADS=4 B=1, M=2048, K=12288, N_HEADS=4
torch - fw 0.36 2.73 3.20 12.90 17.41 51.31 125.33 189.67
xf - fw 0.32 2.53 3.06 12.27 16.58 50.42 124.79 188.62

--- Type: torch.float16 ---

Units: runtime in ms, lower is better B=8, M=384, K=128, N_HEADS=4 B=8, M=784, K=512, N_HEADS=4 B=4, M=1024, K=768, N_HEADS=4 B=4, M=2048, K=1024, N_HEADS=4 B=2, M=2048, K=2048, N_HEADS=4 B=2, M=2048, K=4096, N_HEADS=4 B=2, M=4096, K=4096, N_HEADS=4 B=1, M=2048, K=12288, N_HEADS=4
torch - fw (self_attn) 0.18 0.96 0.98 3.31 3.78 9.85 24.28 33.43
xf - fw (self_attn) 0.19 0.84 0.90 3.01 3.53 9.80 24.03 32.87

--- Type: torch.float32 ---

Units: runtime in ms, lower is better B=8, M=384, K=128, N_HEADS=4 B=8, M=784, K=512, N_HEADS=4 B=4, M=1024, K=768, N_HEADS=4 B=4, M=2048, K=1024, N_HEADS=4 B=2, M=2048, K=2048, N_HEADS=4 B=2, M=2048, K=4096, N_HEADS=4 B=2, M=4096, K=4096, N_HEADS=4 B=1, M=2048, K=12288, N_HEADS=4
torch - fw (self_attn) 0.39 2.72 3.20 12.95 16.92 51.36 125.19 190.71
xf - fw (self_attn) 0.31 2.54 3.07 12.08 16.34 50.73 124.05 187.88

--- Type: torch.float16 ---

Units: runtime in ms, lower is better B=8, M=384, K=128, N_HEADS=4 B=8, M=784, K=512, N_HEADS=4 B=4, M=1024, K=768, N_HEADS=4 B=4, M=2048, K=1024, N_HEADS=4 B=2, M=2048, K=2048, N_HEADS=4 B=2, M=2048, K=4096, N_HEADS=4 B=2, M=4096, K=4096, N_HEADS=4 B=1, M=2048, K=12288, N_HEADS=4
torch - fw+bw 1.20 2.73 2.83 9.14 10.69 28.95 69.35 101.83
xf - fw+bw 2.02 2.41 2.58 8.04 10.19 28.43 66.89 101.62

--- Type: torch.float32 ---

Units: runtime in ms, lower is better B=8, M=384, K=128, N_HEADS=4 B=8, M=784, K=512, N_HEADS=4 B=4, M=1024, K=768, N_HEADS=4 B=4, M=2048, K=1024, N_HEADS=4 B=2, M=2048, K=2048, N_HEADS=4 B=2, M=2048, K=4096, N_HEADS=4 B=2, M=4096, K=4096, N_HEADS=4 B=1, M=2048, K=12288, N_HEADS=4
torch - fw+bw 1.45 7.73 9.58 36.68 50.36 156.19 369.16 607.29
xf - fw+bw 2.12 7.24 9.19 35.41 49.42 154.14 373.03 601.57

--- Type: torch.float16 ---

Units: runtime in ms, lower is better B=8, M=384, K=128, N_HEADS=4 B=8, M=784, K=512, N_HEADS=4 B=4, M=1024, K=768, N_HEADS=4 B=4, M=2048, K=1024, N_HEADS=4 B=2, M=2048, K=2048, N_HEADS=4 B=2, M=2048, K=4096, N_HEADS=4 B=2, M=4096, K=4096, N_HEADS=4 B=1, M=2048, K=12288, N_HEADS=4
torch - fw+bw (self_attn) 1.59 2.72 2.81 9.15 10.79 29.18 70.02 102.94
xf - fw+bw (self_attn) 1.73 2.27 2.47 8.15 9.93 28.05 69.23 100.98

--- Type: torch.float32 ---

Units: runtime in ms, lower is better B=8, M=384, K=128, N_HEADS=4 B=8, M=784, K=512, N_HEADS=4 B=4, M=1024, K=768, N_HEADS=4 B=4, M=2048, K=1024, N_HEADS=4 B=2, M=2048, K=2048, N_HEADS=4 B=2, M=2048, K=4096, N_HEADS=4 B=2, M=4096, K=4096, N_HEADS=4 B=1, M=2048, K=12288, N_HEADS=4
torch - fw+bw (self_attn) 1.14 7.70 9.57 37.02 50.81 157.03 370.30 603.59
xf - fw+bw (self_attn) 1.31 7.20 8.98 35.13 48.92 153.77 375.83 600.75

@blefaudeux blefaudeux merged commit 0c555e9 into main Apr 13, 2022
@fmassa fmassa deleted the benchmark_MHA branch April 14, 2022 08:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@blefaudeux blefaudeux blefaudeux approved these changes

@fmassa fmassa Awaiting requested review from fmassa

Assignees
No one assigned
Labels
CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@dianaml0 @codecov-commenter @blefaudeux @facebook-github-bot