Skip to content

fix(grpo_sync): honor force_on_policy_ratio in TQ trainer - #3087

Merged
terrykong merged 19 commits into
mainfrom
zhiyul/grpo-sync-force-on-policy-parity
Jul 29, 2026
Merged

fix(grpo_sync): honor force_on_policy_ratio in TQ trainer#3087
terrykong merged 19 commits into
mainfrom
zhiyul/grpo-sync-force-on-policy-parity

Conversation

@ZhiyuLi-Nvidia

@ZhiyuLi-Nvidia ZhiyuLi-Nvidia commented Jul 6, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Summary

grpo_sync.py (TQ trainer) never honored loss_fn.force_on_policy_ratio, so it always ran the prev-logprob forward.

Fix

  • Guard policy.get_logprobs_from_meta(meta, …) on not skip_prev_logprobs.
  • Populate prev_logprobs locally with torch.zeros_like(generation_logprobs) when skipping (mirroring legacy train_data["prev_logprobs"] = torch.zeros_like(...)).
  • Skip logprob_inference_prep when both prev and ref are skipped.
  • Emit placeholder seq-level metrics (zero-filled) when the forward is skipped — matches legacy for wandb continuity.

e2e tests

Verified force_on_policy_ratio is honored in grpo-qwen3-235b-16n8g, 16n × 8G e2e run

@ZhiyuLi-Nvidia
ZhiyuLi-Nvidia requested review from a team as code owners July 6, 2026 22:23
@copy-pr-bot

copy-pr-bot Bot commented Jul 6, 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.

@ZhiyuLi-Nvidia ZhiyuLi-Nvidia added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Jul 6, 2026
@ZhiyuLi-Nvidia

ZhiyuLi-Nvidia commented Jul 6, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 8047c26

@ZhiyuLi-Nvidia

ZhiyuLi-Nvidia commented Jul 7, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 4b5ae88

zyzhou5
zyzhou5 reviewed Jul 7, 2026
View reviewed changes

@zyzhou5 zyzhou5 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.

Left some comments. Overall LGTM

Comment thread nemo_rl/algorithms/grpo.py
Comment thread nemo_rl/algorithms/grpo.py
@ZhiyuLi-Nvidia
ZhiyuLi-Nvidia force-pushed the zhiyul/grpo-sync-force-on-policy-parity branch from 4b5ae88 to b6260b8 Compare July 7, 2026 05:26
@ZhiyuLi-Nvidia
ZhiyuLi-Nvidia requested a review from a team as a code owner July 7, 2026 06:27
@ZhiyuLi-Nvidia

ZhiyuLi-Nvidia commented Jul 7, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 4395fde

@ZhiyuLi-Nvidia
ZhiyuLi-Nvidia force-pushed the zhiyul/grpo-sync-force-on-policy-parity branch from 4395fde to 5090ed5 Compare July 24, 2026 07:41
@ZhiyuLi-Nvidia

ZhiyuLi-Nvidia commented Jul 24, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 5090ed5

@ZhiyuLi-Nvidia

ZhiyuLi-Nvidia commented Jul 24, 2026

Copy link
Copy Markdown
Contributor Author

@terrykong could you take a final review at your convenience?

terrykong
terrykong reviewed Jul 25, 2026
View reviewed changes

@terrykong terrykong left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Reviewed with a multi-agent pass (RL/domain, bug scan, tests, design, comment threads, plus an adversarial round). The core fix is correct, and the two parts that were easiest to get wrong both check out:

  • Dropping prev_logprobs from the worker fetch is safe. ClippedPGLossFn never reads data["prev_logprobs"] under force_on_policy_ratio — it substitutes curr_logprobs.detach() — so token_mult_prob_error, gen-kl, policy-kl, JSD and the seq-level IS weights all read curr, never a zero placeholder. Nothing degenerates silently.
  • The grpo.py extraction is behaviour-preserving. The two inlined blocks in grpo_train and async_grpo_train were already identical to each other, and the helper reproduces the predicate, warning text, trigger condition, ordering and all 8 metric keys exactly; the deleted force_on_policy_ratio local has no remaining readers. That was the highest-risk part of the diff and it's clean. Extracting the placeholder-metrics dict is a real win — it stops the three trainers drifting on metric keys, which is invisible in review and painful to debug in wandb.

Also worth saying: the decision not to unify the three trainer loops looks right — grpo_train/async_grpo_train mutate a materialized BatchedDataDict while grpo_sync narrows a TQ column schema, so merging would scatter if data_plane_enabled through a shared path.

Nothing here is blocking; the test note is the only one I'd suggest acting on before merge.

One small ask on the description: the three linked runs are exactly the right shape for a parity claim, but they're linked without numbers — a line quoting the policy_and_reference_logprobs timing pre-fix vs post-fix would show the skip actually took effect rather than merely not breaking anything. Worth a sentence too that _compute_seq_logprob_error_metrics moved outside that timer, so the series step-changes across this commit even on the non-skip path (it now matches legacy, so this looks deliberate).

Separately, and not about this PR: the two data-plane adapters disagree on missing-column semantics — noop.py:213-217 raises KeyError for an unproduced field, while TransferQueue silently omits it (the controller filters unproduced columns out of the schema before readiness is evaluated). That makes noop-backed unit tests stricter than production on a fail-loud-vs-fail-silent axis; it caught several of us during this review. Probably worth a tracking issue for the data-plane owners.

Note: ruff/ruff-format pass on all four files; pyrefly could not run in my environment, so there's no type-check signal from me.

Generated by Claude Code

Comment thread nemo_rl/algorithms/grpo.py Outdated
Comment thread tests/unit/algorithms/test_grpo.py Outdated
Comment thread nemo_rl/models/policy/tq_policy.py Outdated
Comment thread nemo_rl/algorithms/grpo_sync.py Outdated
Comment thread nemo_rl/algorithms/grpo.py Outdated
@ZhiyuLi-Nvidia ZhiyuLi-Nvidia mentioned this pull request Jul 27, 2026
Open
3 tasks
ZhiyuLi-Nvidia and others added 12 commits July 27, 2026 12:58
@ZhiyuLi-Nvidia @claude
fix(grpo_sync): honor force_on_policy_ratio in TQ trainer
8831d62 
grpo_sync.py always ran the prev_logprobs forward, ignoring
loss_fn.force_on_policy_ratio — legacy grpo_train has skipped it since
d1788b2. On grpo-qwen3-235b-16n8g this added ~16s/step (~4% total step
time) of wasted forward compute plus TQ writeback.

Port the skip: call _resolve_logprob_skip_flags, guard
get_logprobs_from_meta on skip_prev_logprobs, write placeholder zeros to
TQ so train_presharded still finds prev_logprobs, and emit placeholder
seq_logprob_error_metrics via the shared factory.

Extract _resolve_logprob_skip_flags and _placeholder_seq_logprob_error_metrics
into private helpers in grpo.py; refactor grpo_train, async_grpo_train, and
grpo_train_sync to share them (prevents future drift across the three trainers).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
@ZhiyuLi-Nvidia @claude
refactor(grpo): drop misleading 0-fill for skipped seq_logprob metrics
efb3081 
Emitting 8 keys pinned to 0.0 whenever prev_logprobs is skipped
(force_on_policy_ratio=True) polluted wandb charts with a permanent
zero line that reads as "seq-level error was 0" — a lie, since no
measurement was taken. metrics.update({}) leaves the keys off for those
steps, which is truthful.

Also restore the pre-existing "Skip prev_logprobs computation when..."
and "todo @jiaqi:" comments that were dropped when the flag-resolution
block was extracted into _resolve_logprob_skip_flags.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
@ZhiyuLi-Nvidia @claude
perf(grpo_sync): drop prev_logprobs TQ round-trip when force_on_polic…
d0efc1c 
…y_ratio

When loss_fn.force_on_policy_ratio=True the loss uses curr_logprobs in
place of prev_logprobs (loss_functions.py:322-324, 337-338) — the fetch
is pure waste. The prior fix wrote a placeholder zero tensor into TQ
just so train_presharded's DP_TRAIN_FIELDS fetch would succeed.

Drop prev_logprobs from the train fetch instead: TQPolicy.train_from_meta
now inspects loss_fn.force_on_policy_ratio and removes prev_logprobs
from meta.fields, and the driver no longer writes zeros to TQ. The
driver-side local `prev_logprobs = torch.zeros_like(generation_logprobs)`
stays (it's still passed to compute_advantage, mirroring legacy).

Also restore _placeholder_seq_logprob_error_metrics() — the all-zero
seq-level metrics dict is intentional (keeps wandb charts continuous
across steps that skip prev_logprobs, matching legacy behavior).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
@ZhiyuLi-Nvidia @claude
refactor(grpo_sync): tighten helper + unify select_fields with dispatch
141afd0 
Follow-up cleanup on the force_on_policy_ratio TQ port:

- _resolve_logprob_skip_flags no longer returns seq_logprob_error_threshold;
  callers grab it from master_config themselves (was only used at the call
  site anyway). Helper is now a 2-tuple.

- grpo_sync.py fuses "should we run the forward?" with "add to
  select_fields" into a single ``if compute_prev`` / ``if compute_ref``
  branch — no more parallel guards on the same flag. Restores the
  positive-named locals (compute_prev, compute_ref) for readability.

- TQPolicy.train_from_meta comment tightened; annotation on train_fields
  dropped (redundant with the tuple-comprehension type).

- Local-zeros comment for the driver-side placeholder shortened to two
  lines pointing at the invariant (workers skip the fetch, see
  train_from_meta).

Empirically identical: run 13498411 already validated the semantics; this
patch only shuffles code inside the same functions.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
@ZhiyuLi-Nvidia @claude
refactor(grpo_sync): thread skip_prev_logprobs through train_from_meta
088c9bb 
Drop the duplicated skip decision. The driver already resolved
skip_prev_logprobs via _resolve_logprob_skip_flags; TQPolicy.train_from_meta
now accepts it as an argument instead of independently inspecting
loss_fn.force_on_policy_ratio. Single config, single source of truth,
threaded from driver to policy.

Also remove a stray meta-comment that described the refactor itself
rather than the code, and tighten the driver-local placeholder comment.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
@ZhiyuLi-Nvidia @claude
refactor(grpo_sync): driver owns the train field list
6cb80f0 
Move the "which fields workers should fetch" decision fully into the
driver. TQPolicy.train_from_meta now takes an optional
``train_fields`` (defaulting to DP_TRAIN_FIELDS), so the policy
doesn't re-derive it from loss config. The driver computes the tuple
once, right after resolving skip flags, and reuses that single source
for both the driver-side read_from_dataplane and the worker-side
train_presharded fetch.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
@ZhiyuLi-Nvidia @claude
refactor(tq_policy): use DP_TRAIN_FIELDS as direct default for train_…
50b1ccf 
…fields

tuple is immutable, so we can put DP_TRAIN_FIELDS in the default arg
directly instead of the None sentinel + runtime fallback.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
@ZhiyuLi-Nvidia @claude
docs(grpo_sync): correct the prev_logprobs placeholder comment
94e5164 
Old comment mentioned "not shipped to workers", which describes the TQ
path (now handled by train_fields), not the reason we allocate a zero
tensor at this line. The real reason: compute_advantage below takes a
logprobs_policy arg that GRPO/Reinforce++ ignore, and OPD is forbidden
in the skip_prev_logprobs=True path (asserted by opd._skip_prev_logprobs).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
@ZhiyuLi-Nvidia @claude
refactor(grpo_sync): pass None to compute_advantage when skipping pre…
b839b14 
…v_logprobs

torch.zeros_like(generation_logprobs) allocated a full [B, S] tensor
just to satisfy compute_advantage's logprobs_policy arg. All estimators
handle logprobs_policy=None: GRPO/Reinforce/GDPO/GAE either swallow it
in **kwargs or gate on ``is not None``, and OPD is asserted out of the
skip_prev_logprobs=True path by opd._skip_prev_logprobs. Drop the
allocation.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
@ZhiyuLi-Nvidia @claude
fix(grpo_sync): guard log_data prev_logprobs.tolist() against None
a82d3ef 
The previous commit switched prev_logprobs to None when force_on_policy_ratio
skips the forward, but missed a downstream debug-log consumer that calls
.tolist() on it. AttributeError: 'NoneType' object has no attribute 'tolist'
on step 1.

Emit None in the log dump (JSON null) when skipped — accurate reflection of
what the driver actually computed.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
@ZhiyuLi-Nvidia @claude
Revert "refactor(grpo_sync): pass None to compute_advantage when skip…
435236d 
…ping prev_logprobs"

Legacy grpo.py::grpo_train and async_grpo_train both set
train_data["prev_logprobs"] = torch.zeros_like(...) when force_on_policy_ratio
skips the forward, and call .tolist() unconditionally in the debug-log path.
Diverging in grpo_sync.py — passing None + guarding .tolist() — was a wash
(the 32 MB allocation is trivial vs. training step time) and broke the
debug-log path with an AttributeError.

Revert to zeros for line-by-line parity with legacy.

This reverts commits a09d39f and e49cb5a.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
@ZhiyuLi-Nvidia @claude
style(grpo_sync): apply ruff-format to train_fields generator
86fbfe6 
CI lint job flagged the train_fields tuple comprehension — ruff-format
prefers each generator clause on its own line.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
ZhiyuLi-Nvidia and others added 6 commits July 27, 2026 12:58
@ZhiyuLi-Nvidia @claude
fix(grpo): assert-out use_kl_in_reward + force_on_policy_ratio
3333116 
The combination is silently broken today (in legacy grpo_train and, as of
this PR's port, grpo_train_sync): when force_on_policy_ratio=True and
seq_logprob_error_threshold is unset, prev_logprobs is a torch.zeros_like
placeholder. Reinforce++/GAE compute_advantage then calls
``calculate_kl(zeros_tensor, reference_policy_logprobs, ...)`` — the
``is not None`` guard at advantage_estimator.py:214 doesn't catch the
zero placeholder, so the KL term corrupts the advantage.

Fail fast in _create_advantage_estimator (shared by all trainers).

Reported in PR review; scope-limited to configuration validation.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
@ZhiyuLi-Nvidia @claude
test(grpo): parametrized coverage for _resolve_logprob_skip_flags + k…
da11484 
…l-reward assertion

Adds a compact matrix over the four flag regimes (default / force_on_policy /
force_on_policy + threshold override / skip_ref) plus a single-case check for
the AssertionError guarding use_kl_in_reward + force_on_policy_ratio in
_create_advantage_estimator.

TQPolicy.train_from_meta's new train_fields arg is exercised by the 235B TQ
e2e run linked in the PR body — unit-testing it would need the full Ray +
worker_group scaffolding.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
@ZhiyuLi-Nvidia @claude
test(grpo): use MasterConfig+ClippedPGLossConfig instead of SimpleNam…
def6ee0 
…espace

Both types are already imported in this file. Using them makes the
fake-config production-realistic and drops the SimpleNamespace shim
entirely.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
@ZhiyuLi-Nvidia @claude
fix(grpo): defensive getattr for use_kl_in_reward assertion
d21dba7 
Existing test_opd.py::test_create_advantage_estimator_opd_branch builds
a SimpleNamespace loss_fn without ``use_kl_in_reward``. Short-circuit on
missing attribute so the assertion doesn't fault on lightweight test
doubles; production ClippedPGLossConfig always has the field.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
@ZhiyuLi-Nvidia @claude
fix(grpo): address PR #3087 review comments
4d04c3d 
Applies terrykong's review feedback on PR #3087:

* Move use_kl_in_reward + skip_prev_logprobs guard from
  _create_advantage_estimator into setup() as a new
  _validate_use_kl_in_reward_compat helper — fails before cluster and
  model weights load rather than after (was grpo_sync.py:455). Guard now
  drops the getattr() and reads loss_config.use_kl_in_reward directly
  (declared field on pydantic-validated ClippedPGLossConfig).

* Allow reference_policy_kl_penalty == 0 through the guard: kl_coef=0
  zeros the KL term regardless, so a zero-placeholder prev_logprobs
  can't corrupt the advantage in that case.

* Assert pytest.warns(UserWarning) on the force_plus_threshold case in
  test_resolve_logprob_skip_flags — pins the helper docstring's "warn
  on incompatible combos" promise (deleting the warnings.warn would
  otherwise leave all four parametrize cases green).

* Correct contradictory docstring in tq_policy.py train_from_meta;
  add train_fields to the Args: block.

* Correct inaccurate comment in grpo_sync.py: train_fields is consumed
  only by train_from_meta (driver read uses select_fields).

* Tighten _resolve_logprob_skip_flags return type from tuple[bool, Any]
  to tuple[bool, bool | None] (skip_reference_policy_logprobs_calculation
  is NotRequired[bool]).

Also adds test_validate_use_kl_in_reward_allows_zero_kl_penalty to pin
the kl_coef=0 exception.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
@ZhiyuLi-Nvidia @claude
refactor(grpo_sync): extract _train_fields_for_step + tighten tq comment
12ed729 
Follow-ups from PR #3087 review:

* Hoist the inline ``prev_logprobs`` filter in ``grpo_train_sync`` into
  a pure ``_train_fields_for_step(skip_prev_logprobs)`` helper (module
  already imports DP_TRAIN_FIELDS). Adds a 2-case parametrize test
  that pins the perf win from silently regressing: TQ ignores an
  unwritten column and ClippedPGLossFn never reads prev_logprobs
  under force_on_policy_ratio, so a broken filter wouldn't crash —
  we would just quietly pay the round-trip again with nothing red.

* Shorten the ``tq_policy.train_from_meta`` inline comment: the Args
  docstring already documents ``train_fields``, so the inline note
  only needs to state the callsite invariant (columns must already
  be in TQ).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
@ZhiyuLi-Nvidia
ZhiyuLi-Nvidia force-pushed the zhiyul/grpo-sync-force-on-policy-parity branch from 5090ed5 to 12ed729 Compare July 27, 2026 20:56
@ZhiyuLi-Nvidia @claude
docs(tq_policy): restore composition + writers in train_from_meta com…
a522962 
…ment

Previous simplification dropped the field composition and delta-write
authorship that terrykong's suggested wording carried. Re-add them in a
tighter 4-line form: composition (rollout + logprob deltas + advantages +
sample_mask; default DP_TRAIN_FIELDS), invariant + writers (must be in
TQ before this call, written by workers + driver delta-writes), and the
narrowing example (drop prev_logprobs under force_on_policy_ratio).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
@ZhiyuLi-Nvidia

ZhiyuLi-Nvidia commented Jul 27, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 12ed729

@copy-pr-bot

copy-pr-bot Bot commented Jul 27, 2026

Copy link
Copy Markdown

/ok to test 12ed729

@ZhiyuLi-Nvidia, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

@ZhiyuLi-Nvidia

ZhiyuLi-Nvidia commented Jul 27, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test a522962

@terrykong
terrykong merged commit e1b2026 into main Jul 29, 2026
80 checks passed
@terrykong
terrykong deleted the zhiyul/grpo-sync-force-on-policy-parity branch July 29, 2026 03:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@terrykong terrykong terrykong approved these changes
@zyzhou5 zyzhou5 zyzhou5 left review comments

Assignees

No one assigned

Labels

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

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@ZhiyuLi-Nvidia @terrykong @zyzhou5