[FSDP] Optimize context parallel #1383
Closed
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
…#1362) Co-authored-by: fy1214 <[email protected]> Co-authored-by: Gao016 <[email protected]> Co-authored-by: yefei12 <[email protected]> Co-authored-by: yzlnew <[email protected]>
Collaborator
|
(this might not be compatible with new FSDP impl...) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 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.
Motivation
Try to optimize context parallel for FSDP backend #1062 by eliminating the
all_gatherfor entropy and log_probs. Instead of gathering log_probs and entropy across all CP ranks and then distributing them, each rank now computes and retains only its local portian, with handling during unpacking and loss computation.Changes
get_logprob_and_entropy_with_cpchange from all_gather full log_probs/entropy to return local shards onlyunpack_sequences_with_cphandle tensor slicing, extracting only the portion of each sequence's response that overlaps with the current rankall_reducethe final loss across cp ranksImpact
During training,
get_logprob_and_entropy_with_cpis called 3 times to obtain ref logprobs, actor logprobs and current logprobs. Threeall_gatheroperations on full-lengths response log_probs+entropy replaced by 1 scalarall_reduce.Concerns
sum_of_sample_meanis now computed as a sum of local contributions.