feat(grpo): support message penalties in single controller - #3734
Conversation
❌ Submodule Fast-Forward Check FailedCheck based on commit: d7318d1 (PR #3734 from ❌ Submodules that need attention:Gym: ❌ PR branch is BEHIND main branch Megatron-Bridge: ❌ PR branch is BEHIND main branch Please ensure all submodule commits are fast-forwards of the main branch before merging. |
d7318d1 to
458f25b
Compare
macandro96
left a comment
macandro96
left a comment
There was a problem hiding this comment.
Team review of PR #3734 by rl-expert, bug-finder, test-agent, and design-reviewer, with devil-advocate stress-testing findings. Semantics faithfully mirror the legacy dtensor GRPO path (PR #2656) via precomputed masks + torch.where; all 42 relevant unit tests pass locally and pre-commit is clean. Design surface is bounded; no over-abstraction warranted for the first violation flag.
4 low-severity suggestions inline (1 consistency, 1 efficiency, 1 test coverage, 1 docstring). 6 additional findings were downgraded or disputed by devil-advocate (hypothetical KeyError path already blocked by the SC validator; unconditional _step_log_dict init is load-bearing for the reduce_advantage_pump_metrics kwargs unpack; ordering already locked by test_invalid_tool_call_takes_precedence_on_overlap).
Generated by Claude Code
macandro96
added
the
CI:Lfast
|
/ok to test 542ce26 |
|
/ok to test 53094ea |
yuki-97
left a comment
yuki-97
left a comment
There was a problem hiding this comment.
Reviewed the SingleController port of message-level advantage penalties (adapting #2656): the mask/count plumbing through payload.py → data-plane schema/column_io.py → single_controller.py's advantage stage → metrics reduction, the NeMo-Gym validation gate, and invalid-tool-call-vs-malformed-thinking precedence on mask overlap. All consistent, tests cover the added paths, and I don't see anything blocking — LGTM.
Also cross-checked against the review comments on #2656 (the legacy path this ports from); none of them recur here.
One thing to handle before merge: this branch currently conflicts with main (mergeable_state: dirty) and the submodule fast-forward check is failing — please rebase.
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
53094ea to
971332f
Compare
|
/ok to test 971332f |
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
|
/ok to test dfe2d2d |
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
|
/ok to test 9cd5e2e |
Main's #3734 message penalties require invalid_tool_call_mask / malformed_thinking_mask train-batch columns; the native rollout path stamps them per message, but the token-capture finalizer publishes a fixed column set that predates the feature, so the first streamed group kills the train pump with a KeyError at the advantage stage (job 6810181). Fail loudly at config validation instead; emitting the masks from capture receipts is a follow-up. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@nvidia.com>
Main's #3734 message penalties require invalid_tool_call_mask / malformed_thinking_mask train-batch columns; the native rollout path stamps them per message, but the token-capture finalizer publishes a fixed column set that predates the feature, so the first streamed group kills the train pump with a KeyError at the advantage stage (job 6810181). Fail loudly at config validation instead; emitting the masks from capture receipts is a follow-up. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@nvidia.com>
2 participants
What does this PR do ?
Adds SingleController support for NeMo-Gym message-level advantage penalties for invalid tool calls and malformed thinking. (Adapts #2656 for single controller path)
The rollout payload now converts NeMo-Gym message flags into token-aligned masks and per-sample counters, carries them through the data plane, and applies the configured advantage values after normal advantage estimation. Invalid-tool-call penalties take precedence when both masks overlap. The SingleController path also validates that these penalties are only enabled with NeMo-Gym and reports global per-step violation counts and rates.
Issues
None.
Usage
Before your PR is "Ready for review"
Pre checks:
Additional Information
Unit tests were added for payload conversion, mask application and precedence, metric reduction, configuration validation, and SingleController advantage-stage integration. Tests and commit hooks have not been run locally.