feat: optimize get logprobs when cp enabled. #528
Draft
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
34f829d to
a068b07
Compare
Signed-off-by: Jonas yang <joyang@nvidia.com>
Signed-off-by: Jonas yang <joyang@nvidia.com>
terrykong
reviewed
Jun 26, 2025
| logits = DTensor.from_local( | ||
| local_logits, | ||
| device_mesh=self.device_mesh["cp", "tp"], | ||
| placements=[Shard(sequence_dim), Shard(-1)], |
There was a problem hiding this comment.
is there any benefit to doing the redistribute like this vs. logits.redistribute(device_mesh=..., placements=....)?
Also, can this be set to async_op=True?
There was a problem hiding this comment.
Accepted. :) Just want to unify full tensor/dtensor format.
terrykong
reviewed
Jun 26, 2025
| assert isinstance(target, DTensor), ( | ||
| "target must be a DTensor if seq_index is provided" | ||
| ) | ||
| cp_mesh = target.device_mesh |
There was a problem hiding this comment.
@SahilJain314 to comment on CP making its appearance in the model agnostic utilities
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.
What does this PR do ?
This PR optimize get logprobs when CP is enabled for FSDP2. Issue #549
Issues
In previous PR, the logits were retrieved from sharded one (local tensor shape [b, s / cp_size, v / tp_size]) into full tensor with shape [b, s, v] and passed to loss function.
The key reason was we had to ensure sequence order was correct when get log probs.
This PR allows permuted sequenced to pass to loss function with additional full tensor
seq_indexwhich indicates the order of the permuted sequence and allow parallel logprobs computation even mixed with tp enabled.Test Result
cp8
tp4cp2
Average step has saved 3.+ seconds.