[ContextParallel] add process-time based Round-Robin load-balance to CP#163617
Closed
XilunWu wants to merge 17 commits intogh/XilunWu/172/basefrom
Closed
[ContextParallel] add process-time based Round-Robin load-balance to CP#163617XilunWu wants to merge 17 commits intogh/XilunWu/172/basefrom
XilunWu wants to merge 17 commits intogh/XilunWu/172/basefrom
Conversation
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/163617
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit ebb2dcd with merge base b54e466 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
ciflow/inductor
oncall: distributed
…balance to CP" cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
…balance to CP" cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
…balance to CP" cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
…balance to CP" cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
…balance to CP" cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
…balance to CP" cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
…balance to CP" cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
…balance to CP" cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
…balance to CP" cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
…balance to CP" cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
XilunWu
added
module: context parallel
drisspg
reviewed
Oct 7, 2025
drisspg
reviewed
Oct 7, 2025
drisspg
reviewed
Oct 7, 2025
| tasks_in_group, _ = torch.sort(tasks_in_group, dim=1) | ||
| return tasks_in_group | ||
|
|
||
| def _generate_indices(self, restore: bool = False) -> Tensor: |
Contributor
There was a problem hiding this comment.
Very pretty :)
we dont have a DtoH sync here right?
Contributor
Author
There was a problem hiding this comment.
no, all operations are tensor ops and happen on block_mask's device.
…balance to CP" **Summary** The load-balancing problem can be modeled as [identical-machines scheduling](https://en.wikipedia.org/wiki/Identical-machines_scheduling) problem. We already provided an easy-to-extend interface in #161062 for implementing load-balancing and in this PR we start with adding a Round-Robin solution as an example and also a verification. This can be easily adapted to other solutions like Shortest-processing-time-first/ Longest-processing-time-first with extra padding added for collectives. - Added a new type of `_LoadBalancer` implementation `_PTRRLoadBalancer` which is designed for `flex_attention()`. This load-balance strategy analyzes the `BlockMask` sparsity info and perform Round-Robin (unlike traditional Round-Robin doing it in circular order, we do in zig-zag order). - Make `_context_parallel_buffers` and `context_parallel_unshard` handle batched load-balance index (previously it can only handle non-batched load-balance index), like in `create_cp_block_mask`. **Test** `pytest test/distributed/tensor/test_attention.py` cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
…balance to CP" **Summary** The load-balancing problem can be modeled as [identical-machines scheduling](https://en.wikipedia.org/wiki/Identical-machines_scheduling) problem. We already provided an easy-to-extend interface in #161062 for implementing load-balancing and in this PR we start with adding a Round-Robin solution as an example and also a verification. This can be easily adapted to other solutions like Shortest-processing-time-first/ Longest-processing-time-first with extra padding added for collectives. - Added a new type of `_LoadBalancer` implementation `_PTRRLoadBalancer` which is designed for `flex_attention()`. This load-balance strategy analyzes the `BlockMask` sparsity info and perform Round-Robin (unlike traditional Round-Robin doing it in circular order, we do in zig-zag order). - Make `_context_parallel_buffers` and `context_parallel_unshard` handle batched load-balance index (previously it can only handle non-batched load-balance index), like in `create_cp_block_mask`. **Test** `pytest test/distributed/tensor/test_attention.py` cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
…balance to CP" **Summary** The load-balancing problem can be modeled as [identical-machines scheduling](https://en.wikipedia.org/wiki/Identical-machines_scheduling) problem. We already provided an easy-to-extend interface in #161062 for implementing load-balancing and in this PR we start with adding a Round-Robin solution as an example and also a verification. This can be easily adapted to other solutions like Shortest-processing-time-first/ Longest-processing-time-first with extra padding added for collectives. - Added a new type of `_LoadBalancer` implementation `_PTRRLoadBalancer` which is designed for `flex_attention()`. This load-balance strategy analyzes the `BlockMask` sparsity info and perform Round-Robin (unlike traditional Round-Robin doing it in circular order, we do in zig-zag order). - Make `_context_parallel_buffers` and `context_parallel_unshard` handle batched load-balance index (previously it can only handle non-batched load-balance index), like in `create_cp_block_mask`. **Test** `pytest test/distributed/tensor/test_attention.py` cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
…balance to CP" **Summary** The load-balancing problem can be modeled as [identical-machines scheduling](https://en.wikipedia.org/wiki/Identical-machines_scheduling) problem. We already provided an easy-to-extend interface in #161062 for implementing load-balancing and in this PR we start with adding a Round-Robin solution as an example and also a verification. This can be easily adapted to other solutions like Shortest-processing-time-first/ Longest-processing-time-first with extra padding added for collectives. - Added a new type of `_LoadBalancer` implementation `_PTRRLoadBalancer` which is designed for `flex_attention()`. This load-balance strategy analyzes the `BlockMask` sparsity info and perform Round-Robin (unlike traditional Round-Robin doing it in circular order, we do in zig-zag order). - Make `_context_parallel_buffers` and `context_parallel_unshard` handle batched load-balance index (previously it can only handle non-batched load-balance index), like in `create_cp_block_mask`. **Test** `pytest test/distributed/tensor/test_attention.py` cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
fegin
approved these changes
Oct 14, 2025
Contributor
fegin
left a comment
fegin
left a comment
There was a problem hiding this comment.
LGTM, please address the comments before merging the PR.
|
|
||
| Warning: | ||
| For Multi-Head Attention, we require the masks over the head dimension are identical | ||
| (i.e. `self.block_mask` must have shape (B, 1, seq_len, seq_len) or (1, 1, seq_len, seq_len)). |
Contributor
There was a problem hiding this comment.
We should add the check in __init__().
Comment on lines
+442
to
+474
| non_sparse_kv_num_blocks = ( | ||
| kv_num_blocks + full_kv_num_blocks | ||
| if full_kv_num_blocks is not None | ||
| else kv_num_blocks | ||
| ) | ||
| B, H, Q = non_sparse_kv_num_blocks.shape | ||
| # requirement: the masking is identical across heads (i.e. H == 1 in BlockMask) | ||
| non_sparse_kv_num_blocks = non_sparse_kv_num_blocks.view(-1, Q) # (B, Q_BLK) | ||
|
|
||
| batch_ptrr = torch.vmap( | ||
| functools.partial( | ||
| _PTRRLoadBalancer.ptrr_scheduling, | ||
| group_size=self.world_size, | ||
| ) | ||
| ) | ||
| ptrr_indices = batch_ptrr( | ||
| non_sparse_kv_num_blocks | ||
| ) # (B, group_size, num_blks_in_group) | ||
| ptrr_indices = ptrr_indices.reshape(B, -1) # (B, num_blocks) | ||
|
|
||
| # NOTE: only support the case where the qkv block size are equal | ||
| q_blk_size, kv_blk_size = block_mask.BLOCK_SIZE | ||
| assert q_blk_size == kv_blk_size, ( | ||
| "for now only support q_blk_size == kv_blk_size" | ||
| ) | ||
|
|
||
| indices = torch.arange( | ||
| q_blk_size * ptrr_indices.size(1), device=ptrr_indices.device | ||
| ).view(-1, q_blk_size) # (NUM_BLOCKS, BLOCK_SIZE) | ||
| indices = indices[ptrr_indices].view(B, -1) # (B, qkv_size) | ||
|
|
||
| if restore: | ||
| indices = torch.vmap(torch.argsort)(indices) |
Contributor
There was a problem hiding this comment.
I'm thinking that should we put the logic to a separate function? The main reason is that I am worried about the performance indication and am thinking if we should compile the code.
Contributor
Author
There was a problem hiding this comment.
Yes, simply add @torch.compile(fullgraph=True) to the function definition. But we'll need to remove "raise" and "assert" since they would break graph when hit.
…balance to CP" **Summary** The load-balancing problem can be modeled as [identical-machines scheduling](https://en.wikipedia.org/wiki/Identical-machines_scheduling) problem. We already provided an easy-to-extend interface in #161062 for implementing load-balancing and in this PR we start with adding a Round-Robin solution as an example and also a verification. This can be easily adapted to other solutions like Shortest-processing-time-first/ Longest-processing-time-first with extra padding added for collectives. - Added a new type of `_LoadBalancer` implementation `_PTRRLoadBalancer` which is designed for `flex_attention()`. This load-balance strategy analyzes the `BlockMask` sparsity info and perform Round-Robin (unlike traditional Round-Robin doing it in circular order, we do in zig-zag order). - Make `_context_parallel_buffers` and `context_parallel_unshard` handle batched load-balance index (previously it can only handle non-batched load-balance index), like in `create_cp_block_mask`. **Test** `pytest test/distributed/tensor/test_attention.py` cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
…balance to CP" **Summary** The load-balancing problem can be modeled as [identical-machines scheduling](https://en.wikipedia.org/wiki/Identical-machines_scheduling) problem. We already provided an easy-to-extend interface in #161062 for implementing load-balancing and in this PR we start with adding a Round-Robin solution as an example and also a verification. This can be easily adapted to other solutions like Shortest-processing-time-first/ Longest-processing-time-first with extra padding added for collectives. - Added a new type of `_LoadBalancer` implementation `_PTRRLoadBalancer` which is designed for `flex_attention()`. This load-balance strategy analyzes the `BlockMask` sparsity info and perform Round-Robin (unlike traditional Round-Robin doing it in circular order, we do in zig-zag order). - Make `_context_parallel_buffers` and `context_parallel_unshard` handle batched load-balance index (previously it can only handle non-batched load-balance index), like in `create_cp_block_mask`. **Test** `pytest test/distributed/tensor/test_attention.py` cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta ezyang msaroufim dcci [ghstack-poisoned]
Contributor
Author
|
@pytorchbot merge |
Collaborator
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
fegin
added a commit
that referenced
this pull request
Oct 17, 2025
#163617 removes the if/else statement to check if the input buffers have the batch dimension. This PR fixes the issue and also adds a test. In the future, we should explicitly ask users to unsqueeze the batch dimension. This is a BC of the existing contract but implicitly infers the batch dimension existence is not safe. ghstack-source-id: e441157 Pull-Request: #165792
fegin
added a commit
that referenced
this pull request
Oct 18, 2025
#163617 removes the if/else statement to check if the input buffers have the batch dimension. This PR fixes the issue and also adds a test. In the future, we should explicitly ask users to unsqueeze the batch dimension. This is a BC of the existing contract but implicitly infers the batch dimension existence is not safe. ghstack-source-id: 321161d Pull-Request: #165792
pytorchmergebot
pushed a commit
that referenced
this pull request
Oct 18, 2025
…165792) #163617 removes the if/else statement to check if the input buffers have the batch dimension. This PR fixes the issue and also adds a test. In the future, we should explicitly ask users to unsqueeze the batch dimension. This is a BC of the existing contract but implicitly infers the batch dimension existence is not safe. Pull Request resolved: #165792 Approved by: https://github.com/XilunWu
Chao1Han
pushed a commit
to Chao1Han/pytorch
that referenced
this pull request
Oct 21, 2025
…CP (pytorch#163617) **Summary** The load-balancing problem can be modeled as [identical-machines scheduling](https://en.wikipedia.org/wiki/Identical-machines_scheduling) problem. We already provided an easy-to-extend interface in pytorch#161062 for implementing load-balancing and in this PR we start with adding a Round-Robin solution as an example and also a verification. This can be easily adapted to other solutions like Shortest-processing-time-first/ Longest-processing-time-first with extra padding added for collectives. - Added a new type of `_LoadBalancer` implementation `_PTRRLoadBalancer` which is designed for `flex_attention()`. This load-balance strategy analyzes the `BlockMask` sparsity info and perform Round-Robin (unlike traditional Round-Robin doing it in circular order, we do in zig-zag order). - Make `_context_parallel_buffers` and `context_parallel_unshard` handle batched load-balance index (previously it can only handle non-batched load-balance index), like in `create_cp_block_mask`. **Test** `pytest test/distributed/tensor/test_attention.py` Pull Request resolved: pytorch#163617 Approved by: https://github.com/fegin
Chao1Han
pushed a commit
to Chao1Han/pytorch
that referenced
this pull request
Oct 21, 2025
…ytorch#165792) pytorch#163617 removes the if/else statement to check if the input buffers have the batch dimension. This PR fixes the issue and also adds a test. In the future, we should explicitly ask users to unsqueeze the batch dimension. This is a BC of the existing contract but implicitly infers the batch dimension existence is not safe. Pull Request resolved: pytorch#165792 Approved by: https://github.com/XilunWu
zhudada0120
pushed a commit
to zhudada0120/pytorch
that referenced
this pull request
Oct 22, 2025
…ytorch#165792) pytorch#163617 removes the if/else statement to check if the input buffers have the batch dimension. This PR fixes the issue and also adds a test. In the future, we should explicitly ask users to unsqueeze the batch dimension. This is a BC of the existing contract but implicitly infers the batch dimension existence is not safe. Pull Request resolved: pytorch#165792 Approved by: https://github.com/XilunWu
Khanaksahu
pushed a commit
to Khanaksahu/pytorch-fork
that referenced
this pull request
Nov 17, 2025
ghstack-source-id: 970f2cc Pull Request resolved: pytorch/pytorch#163617
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 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.
Stack from ghstack (oldest at bottom):
Summary
The load-balancing problem can be modeled as identical-machines scheduling problem. We already provided an easy-to-extend interface in #161062 for
implementing load-balancing and in this PR we start with adding a Round-Robin solution as an example
and also a verification. This can be easily adapted to other solutions like Shortest-processing-time-first/
Longest-processing-time-first with extra padding added for collectives.
_LoadBalancerimplementation_PTRRLoadBalancerwhich is designed forflex_attention(). This load-balance strategy analyzes theBlockMasksparsity info and performRound-Robin (unlike traditional Round-Robin doing it in circular order, we do in zig-zag order).
_context_parallel_buffersandcontext_parallel_unshardhandle batched load-balanceindex (previously it can only handle non-batched load-balance index), like in
create_cp_block_mask.Test
pytest test/distributed/tensor/test_attention.pycc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @ezyang @msaroufim @dcci