Add Packing Support for Context Parallelism (Ring Attention) #2906
+129
−24
Conversation
This file contains hidden or 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
kocchop
requested review from
A9isha,
NicoGrande,
NuojCheng,
RissyRan,
SurbhiJainUSC,
aireenmei,
bvandermoon,
gagika,
gobbleturk,
hengtaoguo,
jiangjy1982,
khatwanimohit,
parambole,
richjames0,
shralex,
shuningjin,
suexu1025 and
vipannalla
as code owners
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
richjames0
approved these changes
Dec 31, 2025
Collaborator
richjames0
left a comment
richjames0
left a comment
There was a problem hiding this comment.
a couple of nits but lgtm
| # Handle packing configurations | ||
| if self.config.packing and self.config.dataset_type != "synthetic": | ||
| if using_context_parallelism and not using_load_balanced_ring_cp: | ||
| raise AssertionError("Packing is only supported for load balanced ring attention with context parallelism.") |
Collaborator
There was a problem hiding this comment.
Nit: AssertionError feels weird here to me. Maybe an argumenterror?
|
|
||
|
|
||
| def get_reorder_callable(cp_size, shard_mode): | ||
| def get_reorder_callable(cp_size, shard_mode, reorder_strategy=0): # 0=DualChunkSwap, 1=Striped |
Collaborator
There was a problem hiding this comment.
As I read this late at night I imagine you're using an integer here so it's comprehensible by JAX but could this be made into an enum without breaking things (at worse using .value?
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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
Enables sequence packing for context parallelism with
ringstrategy using TransformerEngine's DotProductAttention. Includes comprehensive GPU tests for ring attention with packing for sm90+.reorder_causal_load_balancingapiTests
Added a GPU integration test that works for sm90+.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.