decouple force_on_policy_ratio from step0 reforward. - #2061
Merged
copybara-service[bot] merged 1 commit intoSep 4, 2026
Conversation
entrpn
requested review from
abheesht17,
hgao327,
jiangyangmu,
lc5211,
s-noghabi,
sizhit2,
tianshub and
wang2yn84
as code owners
September 1, 2026 00:10
tianshub
requested changes
Sep 1, 2026
tianshub
reviewed
Sep 1, 2026
lc5211
reviewed
Sep 1, 2026
entrpn
force-pushed
the
jfacevedo_add_on_policy_ratio
branch
from
September 2, 2026 16:16
6049f84 to
d7207f3
Compare
Collaborator
Author
4 tasks
tianshub
reviewed
Sep 2, 2026
| force_on_policy_ratio: bool = False | ||
| # Costs one trainer forward pass; keeps the sampler/trainer agreement metrics | ||
| # alive when force_on_policy_ratio would otherwise leave nothing to compare. | ||
| log_sampler_trainer_agreement: bool = False |
Collaborator
There was a problem hiding this comment.
@lc5211 are there any concerns if we disable this? probably no, but just to be aware this is another knob if we want to see the logged metrics
Collaborator
There was a problem hiding this comment.
It's used in combination with force_on_policy_ratio.
or self.algo_config.log_sampler_trainer_agreement``` in this implementation so we are good. maybe let's add comment in line 486 just for the context purpose on when we enforce training logps recomputation.
tianshub
reviewed
Sep 2, 2026
entrpn
force-pushed
the
jfacevedo_add_on_policy_ratio
branch
from
September 3, 2026 17:00
d7207f3 to
5f8597d
Compare
tianshub
approved these changes
Sep 3, 2026
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
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.
Measured on qwen3-4b (OpenMathInstruct-2, agentic rollout):
Gradient magnitude is essentially unchanged; clipping was only firing on
~0.06% of tokens, all of it kernel noise rather than policy drift.
Also in this change:
forward pass to keep the sampler/trainer agreement metrics alive.
Without it, enabling force_on_policy_ratio silently drops them.
raising. The MLPerf reference pairs the two deliberately.
completions; the mask was being compared against a vocab id, so
clip_ratio over-counted.
sub-key. Previously duplicated every metric, or dropped all of them
when no timing keys were present.
Matches the MLPerf Qwen3.5-397B GRPO reference recipe.
Reference
Same implementation as nemo_rl for force_on_policy_ratio, found here: https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/algorithms/loss/loss_functions.py#L520C9-L525C36
Checklist