fix(sc): apply advantage clipping and reward scaling in the advantage stage - #3787
fix(sc): apply advantage clipping and reward scaling in the advantage stage#3787tianyi-zhang-02 wants to merge 13 commits into
Conversation
… stage grpo.advantage_clip_low/high and grpo.reward_scaling parse on the SingleController path and then do nothing: grep for either in single_controller.py or single_controller_utils/ returns no hits. Every other driver applies them. _clip_grpo_advantages is called from grpo.py:3321, grpo.py:4883 and grpo_sync.py:915; scale_rewards from grpo.py:3024. grpo_sync.py:915 is the closest analog to SC's _advantage_stage -- it clips immediately before its data-plane write, which is the same point in the pipeline -- so this calls the same two helpers rather than reimplementing either, and the source-range clamp and its warning stay identical to the batched path. Both are no-ops at their defaults: advantage_clip_* default to None and reward_scaling.enabled to False, so a run that does not set them is unchanged. Two knobs are deliberately left out because SC does not fetch what they need. grpo.reward_shaping reads batch["message_log"] (reward_functions.py:145) and optionally response_token_lengths; the message-level advantage penalties behind grpo.invalid_tool_call_advantage / malformed_thinking_advantage need the message logs too. _advantage_stage pulls only _advantage_input_fields() from the data plane, so both want plumbing beyond this change. Five tests. The three that assert a knob does something fail on main; the two that assert the disabled default changes nothing pass either way, which is the point of having them. The grpo config namespace those tests build by hand moved into a shared _grpo_stub so the next field the stage reads does not break four unrelated tests again. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
# Conflicts: # nemo_rl/algorithms/single_controller.py # tests/unit/single_controller/test_single_controller_actor.py
fbed880 to
93c2551
Compare
…ity PRs NVIDIA-NeMo#3786 and NVIDIA-NeMo#3787 each rewrote this comment to name the knob they removed, so they conflicted on nothing but the prose -- the tuple entries merged fine. Stating the rule once, without naming a knob, lets the two auto-merge in either order. Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com> Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
…ity PRs NVIDIA-NeMo#3786 and NVIDIA-NeMo#3787 each rewrote this comment to name the knob they removed, so they conflicted on nothing but the prose -- the tuple entries merged fine. Stating the rule once, without naming a knob, lets the two auto-merge in either order. Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com> Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
…on point Both parity PRs added a test right before the same anchor, so they conflicted on nothing but placement. Moving this one to the end of the file lets the two auto-merge in either order. Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com> Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
|
Heads-up on merge order, with the resolution written out. Four of my open PRs edit the same twelve-line block — the
I checked every pair by actually merging them. #3786 × #3787 is clean now — they only ever conflicted on the comment wording and on both adding a test before the same anchor, and both are fixed. The remaining four pairs are a genuine textual conflict on that block, and no restructuring avoids it: extracting or reordering the list conflicts just as hard. So rather than force a stack, here is the resolution. Whichever order they land in, this is the merged form: # An enabled one here describes shaping this run does not do. An entry
# leaves this list when the SC path starts implementing it -- rejecting a
# knob is only right while nobody honours it.
#
# DistillationConfig defines none of them: there is no reward to shape or
# filter on, so the list cannot even be evaluated on that path.
unsupported = (
[]
if is_distillation_run(master_config)
else [
name
for name, enabled in (
("use_dynamic_sampling", algo_cfg.use_dynamic_sampling),
("reward_scaling", algo_cfg.reward_scaling.enabled),
("reward_shaping", algo_cfg.reward_shaping.enabled),
)
if enabled
]
)I built that merge locally and ran the SC and config suites against it — 1084 passed. Happy to rebase whichever ones are left once the first lands; just say which order you want. |
The guard used to wrap that list in a conditional, which put this PR on the same twelve lines as NVIDIA-NeMo#3786 and NVIDIA-NeMo#3787 -- each of those removes an entry from it. Four pairwise conflicts on nothing but placement. Returns early instead. DistillationConfig defines none of the knobs the list names, so the check does not apply and the comprehension could not be evaluated anyway; everything after it is PPO-specific and returns early on this path already. The colocated requirement is the one thing that does apply, so it moves into a helper and the guard calls it directly rather than falling through. The list itself is now untouched by this PR, so all four conflicts go away and none of the three needs to land before the others. Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com> Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
The guard used to wrap that list in a conditional, which put this PR on the same twelve lines as NVIDIA-NeMo#3786 and NVIDIA-NeMo#3787 -- each of those removes an entry from it. Four pairwise conflicts on nothing but placement. Returns early instead. DistillationConfig defines none of the knobs the list names, so the check does not apply and the comprehension could not be evaluated anyway; everything after it is PPO-specific and returns early on this path already. The colocated requirement is the one thing that does apply, so it moves into a helper and the guard calls it directly rather than falling through. The list itself is now untouched by this PR, so all four conflicts go away and none of the three needs to land before the others. Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com> Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
The guard used to wrap that list in a conditional, which put this PR on the same twelve lines as NVIDIA-NeMo#3786 and NVIDIA-NeMo#3787 -- each of those removes an entry from it. Four pairwise conflicts on nothing but placement. Returns early instead. DistillationConfig defines none of the knobs the list names, so the check does not apply and the comprehension could not be evaluated anyway; everything after it is PPO-specific and returns early on this path already. The colocated requirement is the one thing that does apply, so it moves into a helper and the guard calls it directly rather than falling through. The list itself is now untouched by this PR, so all four conflicts go away and none of the three needs to land before the others. Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com> Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
The guard used to wrap that list in a conditional, which put this PR on the same twelve lines as NVIDIA-NeMo#3786 and NVIDIA-NeMo#3787 -- each of those removes an entry from it. Four pairwise conflicts on nothing but placement. Returns early instead. DistillationConfig defines none of the knobs the list names, so the check does not apply and the comprehension could not be evaluated anyway; everything after it is PPO-specific and returns early on this path already. The colocated requirement is the one thing that does apply, so it moves into a helper and the guard calls it directly rather than falling through. The list itself is now untouched by this PR, so all four conflicts go away and none of the three needs to land before the others. Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com> Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Not every algorithm block defines it. The OPD test doubles main added in NVIDIA-NeMo#3768 build a bare namespace, and DistillationConfig has no reward to scale at all, so a hard attribute read turns this into an AttributeError on paths that never asked for scaling. A missing block is a disabled one, which is the shipped default anyway. Found by merging all twenty of my open PRs together: this passed alone and failed against main's MOPD tests. Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com> Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
…t-PPO Same class as the reward_scaling read one commit up. Only GRPOConfig carries advantage_clip_low/high; DistillationConfig does not, and the OPD test doubles main added in NVIDIA-NeMo#3768 build a bare namespace. "Not PPO" is not the same as "has the knob", so the not-is_ppo guard alone turns this into an AttributeError on paths that never configured clipping. Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com> Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
…ge reads Nine controller stubs in this file predate that attribute, so they pass here and fail once merged with current main -- NVIDIA-NeMo#3768 made _advantage_stage read it unconditionally. Same shape as the _rollout_manager stub gap in NVIDIA-NeMo#3783. Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com> Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
687fcba to
9e3cfc4
Compare
…t has NVIDIA-NeMo#3512 replaces the bare-tensor return with AdvantageResult. The double now resolves that class through the module instead of importing it, so it returns a tensor on this branch and an AdvantageResult once NVIDIA-NeMo#3512 lands, and this file does not depend on a name that exists only there. An earlier version of this commit also added _teacher_logprobs_required to nine controller stubs. main already sets it in seven places and the suite is green without the other two, so that was duplicate work that collided with NVIDIA-NeMo#3512 on the same lines -- dropped. Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com> Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
9e3cfc4 to
6535295
Compare
Two gaps in this PR's own parity claim. grpo.py logs advantages at :3471 and clips at :3487; grpo_sync.py at :903 then :915. Both report pre-clip. This PR clipped first, so SC's advantages/mean|max|min would have meant something no other driver's does. Clip after the log; the data-plane write is unchanged. The clip guard's 'not self._is_ppo' half was untested, because the PPO test handed the ppo block a GRPO-shaped stub carrying advantage_clip_low -- a field real PPOConfig does not declare. Deleting the clause left every test green. PPOConfig is extra="allow", so the clause is exactly what stops a user-set ppo.advantage_clip_low from running GRPO clipping on a PPO run; there is a test for that now, and the PPO stub matches the real config. Also drop the importlib lookup for AdvantageResult: it does not exist on main, NVIDIA-NeMo#3512 is unmerged, and if that branch ever returned the wrapper the production code at single_controller.py:2028 would break anyway. The guide said setup rejects reward_scaling. It does not any more. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
# Conflicts: # nemo_rl/algorithms/single_controller_utils/config.py
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
What does this PR do?
Makes SingleController honor the existing GRPO advantage-clipping and reward-scaling settings.
grpo.reward_scalinguses the samescale_rewardshelper and runs before advantage estimation.grpo.advantage_clip_low/highuses_clip_grpo_advantagesafter pre-clip diagnostics are logged and before the data-plane write, matching the legacy and synchronous GRPO paths.Both changes are no-ops at their shipped defaults. Reward shaping and message-level advantage penalties remain out of scope because SC does not yet carry the response-length/message fields they require.
Validation
Final SHA:
e8686a6f5149b48898a3abe6b322d38516f9546c, based on upstreammainatccbcd4cc5.2×H100 end to end
Runpod Secure Cloud, Qwen3-0.6B, Megatron policy/reference, real vLLM rollout, TransferQueue data plane, and
NCCL_NVLS_ENABLE=0:train_steps=2,trainer_version=2; two refits succeeded; no stale/drop/abort events0..1 → -1..1; pre-clip advantages[-0.8096, 1.2351]; clip bounds[-0.5, 0.5]; loss0.1040; grad norm2.6518The enabled run used Reinforce++ without baseline subtraction so a uniform zero reward could not make the clipping path vacuous.
Unit