Include token as part of the input/output tuple in all-gather and reduce-scatter #7338
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
…tter Imported from GitHub PR openxla#5740 This PR adds tuple input support to all-gather and reduce-scatter. This is a revival of part of tensorflow/tensorflow#58377 and to be used in conjunction with pytorch/xla#5624 . In FSDP, different layers' weights need to be all-gathered/reduced-scatter during training. If some layers are small, multiple layers' weights can be aggregated for more efficient data transfer (same concept as bucket_cap_mb in DDP). With existing all-gather and reduce-scatter in PyTorch-XLA, you would have to do the bucketing and decomposing outside of the operation. This PR enables multiple different tensors to be all-gathered/reduce-scatter, keeping the original tensor shapes to enable bucketing and decomposing optimizations inside the operation. Original PR has token support like the token used for allreduce to ensure order between CCops. That will be separate PR if needed. Copybara import of the project: -- 7ea1159 by Junmin Hao <[email protected]>: Add Tuple input and token support to all-gather and reduce-scatter. Committer: Junmin Hao <[email protected]> -- cdb873e by Junmin Hao <[email protected]>: lint fix -- aad3521 by Jeffrey Huynh <[email protected]>: Fix hlo_verifier_test failure due to changed expectation -- 32e8145 by Jeffrey Huynh <[email protected]>: Separate the token change out into a separate PR with RFC. -- b301c2a by Jeffrey Huynh <[email protected]>: Change *WithToken tests to *WithTuple -- 5890278 by Jeffrey Huynh <[email protected]>: Fix missing parenthesis Merging this change closes openxla#5740 COPYBARA_INTEGRATE_REVIEW=openxla#5740 from jeffhataws:ag_rs_coalesce_revived 14e09f0 PiperOrigin-RevId: 573976449
This reverts commit 5890278.
This reverts commit b301c2a.
This reverts commit 32e8145.
ddunl
added
kokoro:force-run
|
I think if you rebase the CI will run as expected, sorry about that |
|
Hi @jeffhataws, Can you please check @ddunl comments? Thanks. |
Rebased in another branch, so switching to #7677 . |
|
Switching to #7677 |
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.
(Replaced by #7677)
This is a follow up change to #5740 to include token as part of the input/output tuple in all-gather and reduce-scatter. The change helps clean up the interface to be similar to all-reduce and also enable to use to XLA token type. Per discussions in the previous PR we will open RFC discussion on openxla-discuss.