Skip to content

fix(sc): apply advantage clipping and reward scaling in the advantage stage - #3787

Open
tianyi-zhang-02 wants to merge 13 commits into
NVIDIA-NeMo:mainfrom
tianyi-zhang-02:fix/sc-grpo-knob-parity
Open

fix(sc): apply advantage clipping and reward scaling in the advantage stage#3787
tianyi-zhang-02 wants to merge 13 commits into
NVIDIA-NeMo:mainfrom
tianyi-zhang-02:fix/sc-grpo-knob-parity

Conversation

@tianyi-zhang-02

@tianyi-zhang-02 tianyi-zhang-02 commented Aug 24, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Makes SingleController honor the existing GRPO advantage-clipping and reward-scaling settings.

  • grpo.reward_scaling uses the same scale_rewards helper and runs before advantage estimation.
  • grpo.advantage_clip_low/high uses _clip_grpo_advantages after pre-clip diagnostics are logged and before the data-plane write, matching the legacy and synchronous GRPO paths.
  • clipping is GRPO-only; a similarly named extra field on PPO cannot activate it.

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 upstream main at ccbcd4cc5.

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:

Run Result
default knobs, 2 steps train_steps=2, trainer_version=2; two refits succeeded; no stale/drop/abort events
scaling + clipping enabled, 1 step reward scaling 0..1 → -1..1; pre-clip advantages [-0.8096, 1.2351]; clip bounds [-0.5, 0.5]; loss 0.1040; grad norm 2.6518

The enabled run used Reinforce++ without baseline subtraction so a uniform zero reward could not make the clipping path vacuous.

Unit

  • PPO setup and SingleController actor suites: 108 passed
  • covers enabled/disabled scaling, clip ordering, data-plane values, and the PPO guard

… 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>
@tianyi-zhang-02
tianyi-zhang-02 requested review from a team as code owners August 24, 2026 04:40
@copy-pr-bot

copy-pr-bot Bot commented Aug 24, 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.

@tianyi-zhang-02
tianyi-zhang-02 requested review from a team as code owners August 26, 2026 07:05
# Conflicts:
#	nemo_rl/algorithms/single_controller.py
#	tests/unit/single_controller/test_single_controller_actor.py
@tianyi-zhang-02
tianyi-zhang-02 force-pushed the fix/sc-grpo-knob-parity branch from fbed880 to 93c2551 Compare August 26, 2026 07:44
@svcnvidia-nemo-ci svcnvidia-nemo-ci added the waiting-on-maintainers Waiting on maintainers to respond label Aug 26, 2026
…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>
tianyi-zhang-02 added a commit to tianyi-zhang-02/RL that referenced this pull request Aug 26, 2026
…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>
@tianyi-zhang-02

Copy link
Copy Markdown
Contributor Author

Heads-up on merge order, with the resolution written out.

Four of my open PRs edit the same twelve-line block — the unsupported knob list in _validate_algo_settings:

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.

tianyi-zhang-02 added a commit to tianyi-zhang-02/RL that referenced this pull request Aug 26, 2026
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>
tianyi-zhang-02 added a commit to tianyi-zhang-02/RL that referenced this pull request Aug 26, 2026
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>
tianyi-zhang-02 added a commit to tianyi-zhang-02/RL that referenced this pull request Aug 26, 2026
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>
tianyi-zhang-02 added a commit to tianyi-zhang-02/RL that referenced this pull request Aug 26, 2026
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>
@svcnvidia-nemo-ci svcnvidia-nemo-ci removed the waiting-on-maintainers Waiting on maintainers to respond label Aug 27, 2026
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>
@tianyi-zhang-02
tianyi-zhang-02 force-pushed the fix/sc-grpo-knob-parity branch from 687fcba to 9e3cfc4 Compare August 27, 2026 02:52
…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>
@tianyi-zhang-02
tianyi-zhang-02 force-pushed the fix/sc-grpo-knob-parity branch from 9e3cfc4 to 6535295 Compare August 27, 2026 02:56
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>
@tianyi-zhang-02
tianyi-zhang-02 requested a review from a team as a code owner August 28, 2026 17:29
@github-actions github-actions Bot added the Documentation Improvements or additions to documentation label Aug 28, 2026
# Conflicts:
#	nemo_rl/algorithms/single_controller_utils/config.py
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

community-request Documentation Improvements or additions to documentation waiting-on-maintainers Waiting on maintainers to respond

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants