Use reduce-scatter coalescing for FSDP #6024
Merged
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
| I guess we need to rebase to the master once the dependent PR is landed? |
jeffhataws force-pushed the jeffhataws_fsdp_coaelesce branch from dd3bdac to 1bfae0e Compare 
I had rebased over the dependent PR #5956 |
jeffhataws force-pushed the jeffhataws_fsdp_coaelesce branch from 1bfae0e to 66636c8 Compare 
alanwaketan reviewed Dec 14, 2023
There was a problem hiding this comment.
The change makes a lot of sense to support coalescing reduce-scatter. Just one question, what if I don't need this feature and want to preserve the initial behavior where the reduce-scatter is fired immediately?
I wish I have the resources to perform through-out performance tests in TPU but...
Therefore, will it be possible to add this as an optional feature?
jeffhataws force-pushed the jeffhataws_fsdp_coaelesce branch from f06545d to 3dce325 Compare | Let me know when it's ready for review? |
jeffhataws force-pushed the jeffhataws_fsdp_coaelesce branch from 3dce325 to aac286b Compare 
| @alanwaketan can you take a look at this one? |
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
3 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.
This PR uses reduce-scatter coalescence in FSDP in addition to reduce-scatter's scale param. This PR is companion to #5950 and #5956 and to be used in conjunction with openxla openxla/xla#5740 .
This is a revival of #4145.