feat: optimize get logprobs when cp enabled.#528
Merged
Conversation
github-actions
bot
added
the
documentation
github-actions
bot
removed
the
documentation
joyang-nv
force-pushed
the
joyang/cp_opt
branch
3 times, most recently
from
34f829d to
a068b07
Compare
joyang-nv
changed the title
joyang-nv
requested review from
SahilJain314,
abukharin-nv and
gshennvm
and removed request for
SahilJain314
joyang-nv
added
CI:L1
terrykong
reviewed
Jun 26, 2025
terrykong
reviewed
Jun 26, 2025
joyang-nv
added
CI:L1
SahilJain314
reviewed
Jul 1, 2025
Signed-off-by: Jonas yang <joyang@nvidia.com>
Signed-off-by: Jonas yang <joyang@nvidia.com>
Signed-off-by: Jonas yang <joyang@nvidia.com>
Signed-off-by: Jonas yang <joyang@nvidia.com>
Signed-off-by: Jonas yang <joyang@nvidia.com>
joyang-nv
added
CI:L1
terrykong
approved these changes
Jul 3, 2025
Collaborator
terrykong
left a comment
terrykong
left a comment
There was a problem hiding this comment.
change lgtm, but @SahilJain314 should also give approval since it affects CP in mcore
SahilJain314
approved these changes
Jul 3, 2025
joyang-nv
added
CI:L1
github-merge-queue
bot
removed this pull request from the merge queue due to failed status checks
therealnaveenkamal
pushed a commit
to therealnaveenkamal/RL
that referenced
this pull request
Jul 7, 2025
Signed-off-by: Jonas yang <joyang@nvidia.com>
jialei777
pushed a commit
to jialei777/nemo-rl
that referenced
this pull request
Jul 23, 2025
Signed-off-by: Jonas yang <joyang@nvidia.com> Signed-off-by: Jialei Chen <jialeic@google.com>
KiddoZhu
pushed a commit
that referenced
this pull request
Jul 28, 2025
Signed-off-by: Jonas yang <joyang@nvidia.com>
FannYYW
pushed a commit
to xxman-google/NeMo-RL
that referenced
this pull request
Aug 5, 2025
Signed-off-by: Jonas yang <joyang@nvidia.com>
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
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.
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.