feat(ppo): decouple actor and critic training epochs by bg51717 · Pull Request #3829 · NVIDIA-NeMo/RL

bg51717 · 2026-08-25T19:18:58Z

What does this PR do ?

Adds an independent ppo.critic_ppo_epochs setting while preserving the existing behavior by defaulting it to ppo_epochs.

Critic updates are grouped before policy updates so each colocated model stays resident across its epochs instead of being repeatedly moved between CPU and GPU. The policy and critic Megatron scheduler budgets are sized independently.

Issues

None closed.

Usage

ppo:
  ppo_epochs: 2
  critic_ppo_epochs: 3

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you run the unit tests and functional tests locally?
Did you add or update any necessary documentation?

Additional Information

Tests: 177 targeted unit tests passed with 7 GPU-only skips.
A two-step GB300 functional smoke completed with ppo_epochs=2 and critic_ppo_epochs=3.

copy-pr-bot · 2026-08-25T19:19:03Z

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.

github-actions · 2026-08-28T08:14:22Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 2199947 (PR #3829 from biguo/sc-critic-ppo-epochs)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of yukih/sc-warm-start-critic branch (fast-forward)

All submodule changes look good! ✨

bg51717 · 2026-08-28T17:29:27Z

/ok to test b008666

bg51717

Reviewed with a 4-agent pass (rl-expert, test-agent, bug-finder, devil's-advocate) covering config validation fail-fast paths, GPU residency correctness across all 3 restructured PPO epoch loops (sync/async/single-controller), 3-way consistency between them, and the two independent Megatron train_iters formulas (cross-checked against Megatron-Bridge's scheduler consumption at the pinned submodule SHA).

No bugs found. 4 low/medium-severity suggestions below (2 test-coverage gaps, 1 docstring gap, 1 process ask around memory-fragmentation evidence) survived an adversarial devil's-advocate pass that disputed and discarded a 5th (a doc-page omission that turned out to match that page's own established convention).

Generated by Claude Code

Address review feedback on #3829. Add test_ppo_train_critic_keeps_extra_epochs_during_policy_warmup, which pins the interaction between critic_ppo_epochs and policy_training_start_step at the mock-call-order level. Both knobs were covered individually, but nothing held the critic loop unconditional while the actor loop stays gated behind warmup. Add test_ppo_config_rejects_zero_ppo_epochs. The ppo_epochs < 1 guard is now centralized in PPOConfig.validate_settings, but was only exercised through model_construct, which bypasses pydantic validation entirely. Restore the Returns: section on _value_train_epochs, dropped when it replaced _value_train. Its sibling helpers (_value_stage, _advantage_stage, _sync_weights) all document their return value. Signed-off-by: bg51717 <biguo@nvidia.com>

bg51717 · 2026-08-28T18:33:46Z

/ok to test a39596a

yuki-97

Reviewed the epoch-decoupling change across ppo.py, single_controller.py, the SC setup/config validators, the docs and the tests, along four lenses: correctness of the regrouped train loop, the scheduler-budget contract, config/doc parity, and test coverage. The regrouping itself looks numerically equivalent to the old interleaving (advantages and returns are computed once before the block and the policy loss never reads the critic), and the offload lifecycle is preserved.

Address review feedback on #3829. Add test_ppo_train_critic_keeps_extra_epochs_during_policy_warmup, which pins the interaction between critic_ppo_epochs and policy_training_start_step at the mock-call-order level. Both knobs were covered individually, but nothing held the critic loop unconditional while the actor loop stays gated behind warmup. Add test_ppo_config_rejects_zero_ppo_epochs. The ppo_epochs < 1 guard is now centralized in PPOConfig.validate_settings, but was only exercised through model_construct, which bypasses pydantic validation entirely. Restore the Returns: section on _value_train_epochs, dropped when it replaced _value_train. Its sibling helpers (_value_stage, _advantage_stage, _sync_weights) all document their return value. Signed-off-by: bg51717 <biguo@nvidia.com>

bg51717 · 2026-08-29T14:49:29Z

/ok to test bf66a15

Group critic updates before actor updates so colocated models stay resident across epochs. Add critic_ppo_epochs and separate Megatron scheduler budgets for legacy and single-controller PPO. Signed-off-by: bg51717 <biguo@nvidia.com>

Signed-off-by: bg51717 <biguo@nvidia.com>

Address review feedback on #3829. Add test_ppo_train_critic_keeps_extra_epochs_during_policy_warmup, which pins the interaction between critic_ppo_epochs and policy_training_start_step at the mock-call-order level. Both knobs were covered individually, but nothing held the critic loop unconditional while the actor loop stays gated behind warmup. Add test_ppo_config_rejects_zero_ppo_epochs. The ppo_epochs < 1 guard is now centralized in PPOConfig.validate_settings, but was only exercised through model_construct, which bypasses pydantic validation entirely. Restore the Returns: section on _value_train_epochs, dropped when it replaced _value_train. Its sibling helpers (_value_stage, _advantage_stage, _sync_weights) all document their return value. Signed-off-by: bg51717 <biguo@nvidia.com>

Signed-off-by: bg51717 <biguo@nvidia.com>

Split PPOConfig's combined validator into validate_epoch and validate_async_warmup, and let critic_ppo_epochs implicitly follow ppo_epochs unless explicitly set. Restore the single-controller warmup test's ability to catch a critic/actor epoch mix-up by giving it a critic_ppo_epochs distinct from ppo_epochs, and drop the now-redundant epoch guard in single_controller_utils/config.py along with the two tests that only exercised it. Signed-off-by: ruit <ruit@nvidia.com>

RayenTian · 2026-08-31T23:00:33Z

/ok to test a440755

github-actions Bot added the Documentation Improvements or additions to documentation label Aug 25, 2026

bg51717 added CI:L1 Run doctests, unit tests, and functional tests and removed CI:L1 Run doctests, unit tests, and functional tests labels Aug 25, 2026

yuki-97 force-pushed the yukih/sc-warm-start-critic branch from c14539c to 6e8790c Compare August 26, 2026 02:11

yuki-97 force-pushed the biguo/sc-critic-ppo-epochs branch 2 times, most recently from 73b6590 to 2199947 Compare August 28, 2026 08:13

github-actions Bot added the CI Relating to CI label Aug 28, 2026

yuki-97 changed the base branch from yukih/sc-warm-start-critic to main August 28, 2026 08:13

bg51717 force-pushed the biguo/sc-critic-ppo-epochs branch from 2199947 to 4918f54 Compare August 28, 2026 15:10

github-actions Bot removed the CI Relating to CI label Aug 28, 2026

bg51717 force-pushed the biguo/sc-critic-ppo-epochs branch from 4918f54 to b008666 Compare August 28, 2026 17:23

bg51717 added the CI:L0 Run doctests and unit tests label Aug 28, 2026

bg51717 commented Aug 28, 2026

View reviewed changes

Comment thread tests/unit/algorithms/test_ppo.py

Comment thread tests/unit/algorithms/test_ppo.py Outdated

Comment thread nemo_rl/algorithms/single_controller.py Outdated

bg51717 marked this pull request as ready for review August 28, 2026 18:40

bg51717 requested review from a team as code owners August 28, 2026 18:40

bg51717 requested a review from yuki-97 August 28, 2026 18:41

yuki-97 reviewed Aug 28, 2026

View reviewed changes

bg51717 force-pushed the biguo/sc-critic-ppo-epochs branch from a39596a to bf66a15 Compare August 29, 2026 14:46

yuki-97 reviewed Aug 31, 2026

View reviewed changes

bg51717 added 2 commits August 31, 2026 15:58

feat(ppo): decouple actor and critic training epochs

20d5d2a

Group critic updates before actor updates so colocated models stay resident across epochs. Add critic_ppo_epochs and separate Megatron scheduler budgets for legacy and single-controller PPO. Signed-off-by: bg51717 <biguo@nvidia.com>

fix(ppo): account for critic warmup in scheduler budget

8644dac

Signed-off-by: bg51717 <biguo@nvidia.com>

bg51717 and others added 3 commits August 31, 2026 15:58

fix(ppo): simplify critic epoch configuration

fe2b9c6

Signed-off-by: bg51717 <biguo@nvidia.com>

RayenTian force-pushed the biguo/sc-critic-ppo-epochs branch from bf66a15 to a440755 Compare August 31, 2026 23:00

RayenTian requested a review from yuki-97 August 31, 2026 23:22

yuki-97 approved these changes Sep 1, 2026

View reviewed changes

yuki-97 enabled auto-merge (squash) September 1, 2026 00:13

yuki-97 merged commit 3a0a625 into main Sep 1, 2026
68 checks passed

yuki-97 deleted the biguo/sc-critic-ppo-epochs branch September 1, 2026 01:19

tdene mentioned this pull request Sep 1, 2026

fix(ci): Fix L1 on main #3949

Merged

4 tasks

Uh oh!

Conversation

bg51717 commented Aug 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Issues

Usage

Before your PR is "Ready for review"

Additional Information

Uh oh!

copy-pr-bot Bot commented Aug 25, 2026

Uh oh!

github-actions Bot commented Aug 28, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

bg51717 commented Aug 28, 2026

Uh oh!

bg51717 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

bg51717 commented Aug 28, 2026

Uh oh!

yuki-97 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

bg51717 commented Aug 29, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

RayenTian commented Aug 31, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

bg51717 commented Aug 25, 2026 •

edited

Loading