feat(grpo): make env-flagged sample masking optional - #3402
Merged
yuki-97 merged 6 commits intoJul 31, 2026
Conversation
Adds grpo.mask_env_flagged_samples (NotRequired). Absent or true keeps today's behavior byte-for-byte: NeMo-Gym rollout batches carry the env-driven mask_sample flags and _apply_mask_sample_filter masks the flagged samples from the loss. Setting false keeps the mask_sample key out of the rollout batch entirely (flags are still read from the env results; they are just not attached to the training batch). Why an off switch: on a large sequence-packed async NeMo-Gym run (Qwen3.5-397B, 256 rollouts/step, Megatron sequence packing), the mere presence of the [B]-bool mask_sample key in the batch reproducibly collapsed training even when every flag was False (6-run A/B: 4/4 collapses with the key present, 2/2 healthy with it absent; num_mask_sample_filtered reported 0 in the collapsing runs, so the filter itself was not masking anything). We have not root-caused the interaction between the extra per-sample key and the packed-batch slicing yet; the gate is containment that lets affected runs opt out without removing or altering the feature for everyone else. _apply_mask_sample_filter and its tests are unchanged; the default path keeps mask_sample in the batch as before. Signed-off-by: Michal Futrega <mfutrega@nvidia.com>
michal2409
force-pushed
the
split/mask-env-flagged-gate
branch
from
July 29, 2026 11:35
5f18a04 to
aae1e80
Compare
yuki-97
reviewed
Jul 29, 2026
yuki-97
left a comment
yuki-97
left a comment
Contributor
There was a problem hiding this comment.
overall LGTM, left some minor comments.
cc @ananthsub for this WAR for gym issue, details in #3201 (comment).
yuki-97
reviewed
Jul 29, 2026
25 tasks
- Move the config key from grpo.mask_env_flagged_samples to env.should_mask_flagged_samples, following the should_log_nemo_gym_responses precedent - Read it in one helper in rollouts.py with a docstring on when to turn it off; document the key in the gym exemplar YAMLs - Apply the same gate in AsyncNemoGymRolloutImpl so the single-controller flow honors it too Signed-off-by: Michal Futrega <mfutrega@nvidia.com>
Gym flags rollouts that hit max iterations even when they solve the task, and those are samples worth training on (see NVIDIA-NeMo#3201 discussion). Signed-off-by: Michal Futrega <mfutrega@nvidia.com>
yuki-97
reviewed
Jul 30, 2026
yuki-97
left a comment
yuki-97
left a comment
Contributor
There was a problem hiding this comment.
@michal2409 thanks for the update!
the new added one in nemo_rl/experience/rollout_manager.py seems not pass the config correctly, and could you help to add some unit tests to guard? flagged inline.
…nager RolloutManager never forwarded mask_env_flagged_samples to AsyncNemoGymRolloutImpl, so env.should_mask_flagged_samples had no effect on the single-controller path. Forward it and pass should_mask_flagged_samples(master_config.env) from setup_single_controller. Add unit coverage: the config helper, the batched postprocess gate, the streaming _result_to_completion gate, and a RolloutManager forwarding test that fails if the wiring is dropped again. Signed-off-by: Michal Futrega <mfutrega@nvidia.com>
yuki-97
reviewed
Jul 30, 2026
Signed-off-by: Michal Futrega <mfutrega@nvidia.com>
yuki-97
approved these changes
Jul 31, 2026
yuki-97
added
the
CI:Lfast
Contributor
|
/ok to test 79514c2 |
@yuki-97, there was an error processing your request: See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/ |
Contributor
|
/ok to test d305d56 |
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
2 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.
NeMo-Gym environments and agents can flag individual rollout samples as "don't train on this" (the mask_sample flag, added in #3163). Flagged samples are dropped from the policy-gradient loss while still counting toward advantage computation. Currently this behavior is always on: every gym rollout batch carries the flags, and any sample the environment flags is silently masked out.
This PR makes that behavior configurable via one optional boolean,
env.should_mask_flagged_samples(documented in the gym exemplar YAMLs, next toshould_log_nemo_gym_responses):The key is read in a single helper (
should_mask_flagged_samplesinrollouts.py) used by the sync trainer, the async trajectory collector, andAsyncNemoGymRolloutImpl(so the future single-controller flow honors the gate too — there, turning it off drops the flag from the result before it lands inCompletion.env_extras).Why you'd want to turn it off: the flags can be too coarse to honor — Gym flags rollouts that hit max iterations even when they solve the task, and those are samples worth training on (#3201 (comment)). Also, how many samples get flagged is decided by the environment/agent and varies run to run, so masking changes the effective batch composition non-deterministically. For controlled experiments and benchmark runs we need the loss to see every sample; A/B runs showed the masked-sample ramp changing training dynamics mid-run. The gate makes that choice explicit while leaving the default exactly as it is.