Skip to content

decouple force_on_policy_ratio from step0 reforward. - #2061

Merged
copybara-service[bot] merged 1 commit into
google:mainfrom
entrpn:jfacevedo_add_on_policy_ratio
Sep 4, 2026
Merged

decouple force_on_policy_ratio from step0 reforward.#2061
copybara-service[bot] merged 1 commit into
google:mainfrom
entrpn:jfacevedo_add_on_policy_ratio

Conversation

@entrpn

@entrpn entrpn commented Sep 1, 2026

Copy link
Copy Markdown
Collaborator

Measured on qwen3-4b (OpenMathInstruct-2, agentic rollout):

  • clipfrac 0.0006 -> 0.0000
  • log_ratio_abs 0.0157 -> 0.0000
  • step time 984s -> 742s (~24%)

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:

  • log_sampler_trainer_agreement: opt-in flag that spends one trainer
    forward pass to keep the sampler/trainer agreement metrics alive.
    Without it, enabling force_on_policy_ratio silently drops them.
  • force_on_policy_ratio + off_policy_steps > 0 now warns instead of
    raising. The MLPerf reference pairs the two deliberately.
  • Compare completion_tokens[-1] against EOS when counting clipped
    completions; the mask was being compared against a vocab id, so
    clip_ratio over-counted.
  • Buffer metrics once per group instead of once per env/reward-time
    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

  • I have added all the necessary unit tests for my change.
  • I have verified that my change does not break existing code and all unit tests pass.
  • I have added all appropriate doc-strings/documentation.
  • My PR is based on the latest changes of the main branch (if unsure, rebase the code).
  • I have signed the Contributor License Agreement.
  • I have followed Contribution Guidelines.

Comment thread tests/rl/agentic/agentic_grpo_learner_test.py Outdated
Comment thread tunix/rl/agentic/agentic_grpo_learner.py Outdated
Comment thread tunix/rl/agentic/agentic_grpo_learner.py Outdated
Comment thread tunix/rl/agentic/agentic_grpo_learner.py Outdated
Comment thread tunix/rl/algo_core.py Outdated
Comment thread tunix/rl/agentic/agentic_grpo_learner.py Outdated
@entrpn
entrpn force-pushed the jfacevedo_add_on_policy_ratio branch from 6049f84 to d7207f3 Compare September 2, 2026 16:16
@entrpn

entrpn commented Sep 2, 2026

Copy link
Copy Markdown
Collaborator Author

@lc5211 @tianshub resolved your comments, PTAL.

Comment thread tunix/rl/agentic/agentic_grpo_learner.py Outdated
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

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.

@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

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.

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.

Comment thread tunix/rl/agentic/agentic_grpo_learner.py
Comment thread tunix/rl/agentic/agentic_grpo_learner.py
@entrpn
entrpn force-pushed the jfacevedo_add_on_policy_ratio branch from d7207f3 to 5f8597d Compare September 3, 2026 17:00
Comment thread tunix/rl/agentic/agentic_grpo_learner.py
@copybara-service
copybara-service Bot merged commit f45a945 into google:main Sep 4, 2026
35 of 37 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants