Skip to content

feat(grpo): mask env-flagged samples from policy loss - #3163

Merged
yfw merged 2 commits into
mainfrom
amahishi/add-mask-sample-mr355
Jul 14, 2026
Merged

feat(grpo): mask env-flagged samples from policy loss#3163
yfw merged 2 commits into
mainfrom
amahishi/add-mask-sample-mr355

Conversation

@macandro96

Copy link
Copy Markdown
Contributor

What does this PR do ?

Adds a per-sample gradient mask driven by the rollout environment. When instance_config.mask_sample is True on a rollout result, the sample's loss_multiplier is zeroed out before the policy update — but the sample is kept for advantage computation (baseline / std).

Cherry-pick commit here

Issues

List issues that this PR closes (syntax):

Usage

  • You can potentially add a usage example below
# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
  • Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

  • ...

@macandro96
macandro96 requested review from a team as code owners July 11, 2026 03:25
@macandro96 macandro96 added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Jul 11, 2026
@copy-pr-bot

copy-pr-bot Bot commented Jul 11, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@macandro96

Copy link
Copy Markdown
Contributor Author

/ok to test c112b82

@macandro96
macandro96 requested a review from sdevare-nv July 11, 2026 03:25
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96
macandro96 force-pushed the amahishi/add-mask-sample-mr355 branch from c112b82 to 8112e34 Compare July 13, 2026 15:21
@macandro96

Copy link
Copy Markdown
Contributor Author

/ok to test 8112e34

sdevare-nv
sdevare-nv previously approved these changes Jul 13, 2026

@sdevare-nv sdevare-nv left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

looks good to me!

Comment thread nemo_rl/experience/rollouts.py
Comment thread nemo_rl/algorithms/grpo.py
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96

Copy link
Copy Markdown
Contributor Author

/ok to test 27a0f20

@macandro96
macandro96 requested a review from yfw July 13, 2026 21:28
@yfw
yfw enabled auto-merge (squash) July 13, 2026 21:41
@yfw
yfw merged commit 4372db1 into main Jul 14, 2026
129 of 139 checks passed
@yfw
yfw deleted the amahishi/add-mask-sample-mr355 branch July 14, 2026 01:08
binhu-nv pushed a commit that referenced this pull request Jul 22, 2026
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Co-authored-by: Sugam Devare <sdevare@nvidia.com>
(cherry picked from commit 4372db1)
binhu-nv pushed a commit that referenced this pull request Jul 31, 2026
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Co-authored-by: Sugam Devare <sdevare@nvidia.com>
(cherry picked from commit 4372db1)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants