[JAX] Context Parallel Attention with All-Gather #1106
Merged
+942
−262
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
mgoldfarb-nvidia
changed the title
|
How is this different from #1059? |
|
The other method is implemented as ring attention with point to point comms. Both will ultimately provide context parallel attention but some forms of attention e.g window attention are more easily supported with ag approach. We also may get better scaling of the ag comms on multi node setups.
Michael Goldfarb
Get Outlook for iOS<https://aka.ms/o0ukef>
…________________________________
From: Przemyslaw Tredak ***@***.***>
Sent: Wednesday, August 14, 2024 7:32:03 PM
To: NVIDIA/TransformerEngine ***@***.***>
Cc: Michael Goldfarb ***@***.***>; Author ***@***.***>
Subject: Re: [NVIDIA/TransformerEngine] [JAX] DRAFT: Context Parallel Attention with All-Gather (PR #1106)
How is this different from #1059<#1059>?
—
Reply to this email directly, view it on GitHub<#1106 (comment)>, or unsubscribe<https://github.com/notifications/unsubscribe-auth/BI2EORL7SERQM3OUS62GNU3ZRPZIHAVCNFSM6AAAAABMRF7F5CVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDEOJQGE2TEMBTGM>.
You are receiving this because you authored the thread.Message ID: ***@***.***>
|
|
Hmmm, sure, but it still feels like duplicate of the functionality. Could you maybe collaborate with @mingxu1067 to merge your work with his? |
|
Sure thing. Ming and I have already been in discussion on how to merge the PRs. Likely there should be an initial CP attention PR and will need to follow with updated to other components on the jax side as we implement the rest of the CP features for jax. |
mgoldfarb-nvidia
force-pushed
the
mgoldfarb-nvidia/context_parallel_attention_with_all_gather
branch
2 times, most recently
from
4cd496a
to
050bce8
Compare
mgoldfarb-nvidia
changed the title
mgoldfarb-nvidia
force-pushed
the
mgoldfarb-nvidia/context_parallel_attention_with_all_gather
branch
2 times, most recently
from
43a8d78
to
988e3f6
Compare
|
|
||
| register_primitive(FusedAttnCPWithAllGatherBwdPrimitive) | ||
|
|
||
|
|
||
| def fused_attn_fwd( |
There was a problem hiding this comment.
This is an important discussion point: Do we keep CP hidden behind a common function or add a separate flavor of fused_attn_cp_allgather_fwd. The though here was it makes sense to keep a common interface that naturally support CP. Other implementation strategies e.g. ring can be exposed via argument.
mgoldfarb-nvidia
force-pushed
the
mgoldfarb-nvidia/context_parallel_attention_with_all_gather
branch
from
988e3f6
to
e912b64
Compare
mgoldfarb-nvidia
force-pushed
the
mgoldfarb-nvidia/context_parallel_attention_with_all_gather
branch
from
e912b64
to
03e34c0
Compare
zlsh80826
reviewed
Aug 28, 2024
mgoldfarb-nvidia
force-pushed
the
mgoldfarb-nvidia/context_parallel_attention_with_all_gather
branch
5 times, most recently
from
a83875c
to
df5267a
Compare
zlsh80826
reviewed
Sep 5, 2024
zlsh80826
reviewed
Sep 5, 2024
mgoldfarb-nvidia
force-pushed
the
mgoldfarb-nvidia/context_parallel_attention_with_all_gather
branch
2 times, most recently
from
ba88ede
to
d3c9d06
Compare
zlsh80826
approved these changes
Sep 6, 2024
|
Please help check the pre-commit failure. You can run |
mgoldfarb-nvidia
force-pushed
the
mgoldfarb-nvidia/context_parallel_attention_with_all_gather
branch
from
5c6bc6c
to
83a62b4
Compare
|
/te-ci jax |
mgoldfarb-nvidia
force-pushed
the
mgoldfarb-nvidia/context_parallel_attention_with_all_gather
branch
from
83a62b4
to
0691181
Compare
|
/te-ci jax |
1 similar comment
|
/te-ci jax |
Signed-off-by: Michael Goldfarb <[email protected]>
mgoldfarb-nvidia
force-pushed
the
mgoldfarb-nvidia/context_parallel_attention_with_all_gather
branch
from
0691181
to
5194eb4
Compare
|
/te-ci jax |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
Adds support for context parallel fused attention based on an all-gather/reduce-scatter approach. This implementation exposes the collective communication between CP ranks.
The first implementation of CP only support causal and no masking without bias. Additional QKV formats and configurations will be added to subsequent PRs.
Fixes # (issue)
Type of change
Changes
Adds context parallel axis resource and new primitives to fused attention.
Checklist: