feat(sc): support PPO in single controller

[yuki-97]

What does this PR do ?

Ports PPO — including the critic/value model — to the SingleController (SC) path, so an async PPO run can be driven by examples/run_grpo_single_controller.py instead of run_ppo.py. Megatron critic only, async only.

What lands here

Config. MasterConfig.grpo becomes optional and a ppo: block is added beside it; the two are mutually exclusive and every read goes through algo_config(). The advantage estimator config moves from dict to a BaseModel with name pinned to a Literal. _validate_algo_settings rejects the misconfigurations listed below at setup time, each message naming the fix.

Value model on the data plane. New TQValue (driver) and get_values_presharded (worker), with TQDriverMixin extracted from tq_policy.py so the policy and the critic share one copy of the packing / data-plane helpers. PPO_VALUE_FIELDS is deliberately kept out of DP_TRAIN_FIELDS so a GRPO run never fetches a column nobody wrote.

Train pump. _value_stage (critic forward) and _value_train (critic optimizer step) around the existing advantage stage, plus ppo.ppo_epochs (N optimizer steps per RL step, the policy offloaded between epochs — closes the "multiple mini-steps inside a single RL step" item in #2625), ppo.policy_training_start_step (critic warmup: weight sync skipped, policy optimizer not checkpointed, step excluded from best-checkpoint tracking), and async_rl.sampler.warmup_lookahead_versions (the SC equivalent of legacy warmup_generation_lead_steps; buffer capacity is sized from the peak window). The critic forward runs after the policy and reference logprobs rather than before as in ppo.py, the policy is parked on CPU across both critic stages, and the two checkpoint saves are serialized critic-first.

Setup. The train cluster is sized for two worker groups on a PPO run, the critic is built and resumed alongside the policy, and the Megatron tick budget is scaled by ppo_epochs.

Constraints enforced at config time

• ppo.async_ppo: null — SC drives the async loop itself (rejected in the entrypoint).
• ppo and grpo are mutually exclusive, and exactly one must be set.
• value and value_loss_fn are present iff the ppo block is, and value.megatron_cfg.enabled: true.
• ppo.ppo_epochs >= 1.
• async_rl.sampler.name: in_order.
• min_groups_for_streaming_train == ppo.num_prompts_per_step.
• value.train_global_batch_size == num_prompts_per_step * num_generations_per_prompt, so one RL step is one critic optimizer.step.
• async_rl.rollout_failure.max_skipped_prompts and max_consecutive_dropped_prompts are both 0.
• warmup_lookahead_versions requires policy_training_start_step > 0.
• ckpt_assume_constant_structure is rejected together with critic warmup + optimizer checkpointing.

policy.offload_optimizer_for_logprob: true is required but not validated — nothing else on the SC path takes the policy optimizer off the GPU.

Not covered (intentional, tracked)

Rejected at setup rather than silently skipped, except where noted.

• ppo.overlong_filtering, reward_shaping, reward_scaling, use_dynamic_sampling — implemented on neither algorithm block on the SC path ([Single Controller / Async RL] cleanup tracking issue #2625).
• Samplers other than in_order — windowed, weight_fifo and custom are not supported under PPO.
• Streaming critic training — the value workers have no split begin/microbatch/finish train API yet ([Single Controller / Async RL] cleanup tracking issue #2625), so one RL step must be one chunk.
• Rollout drop budgets — a drop shortens the step, and the critic shards it against the configured value.train_global_batch_size rather than its actual size, so the first short step fails a divisibility assert inside the value workers. The resiliency layer (feat(sc): tolerate and replace dropped rollouts in the SingleController #3665) stays available on GRPO.
• DTensor critic — out of scope, Megatron only.
• adv_estimator.name: raw_reward under PPO — yields no returns column; ppo.py behaves the same way. Not rejected.
• total_flops is reported from the last epoch only, so it under-reports N-fold under ppo_epochs > 1; throughput here is not comparable to the v1 path.
• The exemplar's default sizing OOMs (PPO default yaml OOM #3793) — pre-existing across the PPO exemplar family, not introduced here. The functional and nightly runs both pass with their own sizing.
• Renaming run_grpo_single_controller.py to run_single_controller.py — separate PR, to keep this diff reviewable.

Issues

Advances #2625 (does not close it). The SC exemplar added here inherits the exemplar-sizing OOM tracked in #3793.

Usage

uv run examples/run_grpo_single_controller.py --config examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml

examples/configs/ppo_math_1B_megatron_single_controller.yaml is the exemplar (schema documentation / starting point). Running it as-is OOMs — see #3793, which covers all three PPO exemplars. The functional test drives it with the usual smoke-test overrides (Qwen/Qwen2.5-0.5B, num_prompts_per_step=2, train_global_batch_size=8, 2 GPUs) and passes, as does the 2-node nightly recipe.

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 added. Unit: test_ppo_setup.py, tests/unit/models/value/test_tq_value.py and test_sampler_interface.py are new, plus PPO coverage in test_single_controller_actor.py and test_checkpointing.py — the epoch loop, the warmup path, the value stage, pad-target isolation, and save/resume across the warmup boundary. Functional: tests/functional/ppo_async_single_controller.sh, registered in the L1 list. Nightly: the 2n8g recipe above, added to nightly.txt; its thresholds come from this recipe's own 40-step run, not the legacy recipe's — at 256 prompts/step it has consumed a quarter of the legacy run's data by step 40.

Docs. docs/guides/single-controller.md now covers the PPO path: the ppo: / value: / value_loss_fn: blocks, the constraints above, warmup_lookahead_versions, and the legacy ppo.async_ppo.* → async_rl.* migration table.

Test Results. yellow: legacy with same settings; red: v2, this PR; blue: v2, this PR. resume from ckpt 20.

[copy-pr-bot]

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.

[yuki-97]

/ok to test a61654e

[yuki-97]

/ok to test aa5bfb3

[yuki-97]

/ok to test 5026969

[tianyi-zhang-02]

Four places this touches what I have open. Flagging them so the ordering is yours, not something we hit in a conflict.

1. scale_rewards does not reach SC. For PPO that is a regression — ppo.py calls it in ppo_train, and grep -c scale_rewards nemo_rl/algorithms/single_controller.py on this branch is 0. So grpo.reward_scaling silently stops working when a PPO run moves to SC. #3787 applies it at the top of _advantage_stage, after _value_stage and before compute_advantage, so GAE sees scaled rewards. (advantage_clip_low/high is in that same PR but is GRPO-only, so not a PPO regression.)

2. payload.py:83 still reads sample_mask = torch.ones(n, ...). grpo.py builds that from DatumSpec.loss_multiplier and zeroes it for env-flagged and, under overlong_filtering, truncated rows. Both signals are on the Completion already; nothing reads them. So the PPO path inherits it — rows every v1 driver drops get a value target and a GAE advantage. #3786 fixes it.

3. advantage_estimator.py — #3512 is the other half of this. You type the config side; #3512 types the return and the contract (AdvantageResult, plus a Protocol both _create_advantage_estimator functions are annotated with). Same direction, and they will conflict — both rewrite the estimator __init__s and compute_advantage signatures. Yours is the feature and mine has sat since 2026-08-06 with no reviewer, so I would rather adapt: rebase #3512 on top of this once it settles, or close it and hand you the AdvantageResult change to fold in. Your call.

4. The test renames. This moves test_single_controller.py → test_single_controller_actor.py, where #3787 adds its tests. Nothing for you to do — I will rebase after this lands.

1 and 2 are pre-existing and open separately; no action needed for this PR to be correct. Mentioning them because "SC never applied it" is the same shape as your #3770, and the PPO path picks both up :)

[yuki-97]

hi @tianyi-zhang-02 , thanks for your interest. this is the init implementation of single controller + PPO and is aim to get into main in recent days.
it's intended that there will be many missing features. critical features will be updated in this PR, others may record in #2625 and may or may not need implementing.

many things is still subject to change, feel free to follow this PR to watch.

[tianyi-zhang-02]

Went through the whole diff rather than just my overlaps. One finding, one note, and four things that hold up — recording those too, since "I looked and it is fine" is worth as much to you as a list.

Finding, low severity: OPDAdvantageEstimator is the only estimator left out of the typed-config migration.

Five of six are converted — advantage_estimator.py:80, :124, :218 to AdvEstimatorConfig, :292, :345 to GAEConfig. :589 is still def __init__(self, estimator_config: dict, loss_config: dict), and the only one whose loss_config is still dict.

It matters a little more than an annotation. This PR narrows the name to Literal["grpo", "gdpo", "opd", "reinforce_plus_plus"], making opd first-class — while grpo.py:2275 still constructs it as OPDAdvantageEstimator({"name": "opd"}, loss_config), a literal dict, and the __init__ body ignores both arguments. So setting adv_estimator.name: opd alongside any other adv_estimator.* field drops those fields silently. Typing the signature and passing the real config through would make that visible.

Note: this closes a #2625 item. "Support multiple mini-steps inside a single RL step" is still unchecked there, and self._ppo_epochs / for epoch in range(self._ppo_epochs) in single_controller.py is that. Worth ticking off when this lands.

Checked and clean — listing these so you know where I did not find anything:

• ppo.py parity. Enumerated every ppo.* key v1 reads. policy_training_start_step is honoured (11 sites), ppo_epochs is honoured. batch_multiplier is absent but gated on use_dynamic_sampling, already tracked in [Single Controller / Async RL] cleanup tracking issue #2625.
• The dict → model migration. estimator_config[ and estimator_config.get( both return nothing tree-wide, so no subscript survived.
• The fix(sc): reduce and finalize gradients once per step, not per chunk #3660 shape on the value model. Went looking for the critic's gradient being reduced per chunk instead of per step. You avoided it deliberately and said why in the _train_pump docstring.
• And that invariant is enforced, not assumed — the thing I came to check. _validate_algo_settings in config.py hard-errors when min_groups_for_streaming_train != num_prompts_per_step under PPO, with a message naming the critic's per-chunk optimizer step. Right place for it.

Three lenses I did not get to: the GAE recursion itself (mask handling at sequence ends, VAPO decoupled λ, normalize_advantages's axis), driver_mixin.py versus the 77 lines leaving tq_policy.py, and whether the six-file rename dropped a test. Happy to take any — say which and I will do it against the branch rather than guess :)

[tianyi-zhang-02]

Closing two of the three lenses. Both clean, no action needed.

The six-file rename drops no test. Collected tests/unit/single_controller/ on both trees and diffed the ID sets with the filename stripped, so a test moving between renamed files does not register as a change:

clean main (ffbf33f3): 357 collected
this branch:           424 collected
main-only:             (empty)

67 net new, nothing lost.

driver_mixin.py is a real extraction, not a parallel implementation. tq_policy.py loses exactly _stamp_pad_seqlen, read_from_dataplane, write_to_dataplane, _packing_args; the mixin provides exactly those four plus a new _isolated_meta. Two consumers — TQPolicy(TQDriverMixin, Policy) and TQValue(TQDriverMixin, Value) — so it earns being shared.

That leaves the GAE recursion: mask handling at sequence ends, the decoupled gae_lambda_value / gae_lambda_policy with length_adaptive_alpha, and which axis normalize_advantages reduces over. That one I would want to run rather than read — it is where reading is least reliable, and I would rather say so than give you a confident-sounding opinion about a recursion I only eyeballed :)

[yuki-97]

/ok to test 9d61ad0

[yuki-97]

Self-review pass over the whole diff, done with a team of specialised agents plus an adversarial verification round. 14 inline comments below: 3 correctness bugs, 9 suggestions, 1 batched nit.

Rebase needed -- the PR is currently CONFLICTING against main. Worth doing before the comments below, since a few of them touch files that will move.

What held up well, so it does not get lost among the findings: test coverage is genuinely strong -- ppo_epochs > 1, the critic warmup path, the value stage, the pad-target isolation and save/resume across the warmup boundary all have real unit coverage, and three reviewers independently said so. _validate_algo_settings is the right pattern, with seven config-time rejections whose messages each name the fix. Keeping PPO_VALUE_FIELDS out of DP_TRAIN_FIELDS so a missing column errors rather than reading zeros is the correct call. driver_mixin.py is a real two-consumer extraction. The Megatron loss-averaging counteraction from 6e103d0 is fully preserved, VPP is asserted off, PP is handled correctly, and every Megatron API call checks out against the pinned SHA.

Checked and cleared, listed so nobody re-derives them: the GAE recursion, masking and normalize_advantages axis are unchanged at the merge base; the logprobs -> logprobs_policy rename is a fix rather than a break (SC already used the new names, so its KL kwargs were being swallowed by **kwargs); train_iters = max_num_steps * ppo_epochs matches ppo.py verbatim; old logprobs, values and advantages are computed once before the epoch loop, which is correct PPO; the six-file rename drops no test and leaves no stale reference anywhere in the tree; both new test files genuinely collect in their CI shards.

Two candidate findings died in verification and are recorded here rather than posted. The nightly GPU-hour cap going to 4190 when the measured total is 4155 looked like slack that should be tightened -- but git log -L on that assertion shows the cap has never been an exact fit, so 35 hours of headroom is normal practice. And adv_estimator.name: raw_reward reaching the critic with no returns column is real, but ppo.py:1580 does the identical thing, so it is ported behaviour, not something this PR introduces.

One scope note: roughly 60 lines of the diff are unrelated churn -- pyrefly.toml re-sorted end to end, the L1 driver reordered (which is what carried away the 42 comments flagged below), and two of the six renames are pure no-ops. Splitting the rename and re-sort commits out would make this easier to review.

Generated by Claude Code

[yuki-97]

/ok to test b5b1bb8

[yuki-97]

/ok to test 90c48c9

[tianyi-zhang-02]

Looked at the five reds. They are two separate things, and I could only pin one down.

The three L0_Unit_Tests_Megatron_Policy_* shards are one missing line in a test fixture.

All three fail the same way:

if self._first_train_step_forward_pre_hook_disabled and update_successful:
E   AttributeError: 'MegatronPolicyWorkerImpl' object has no attribute '_first_train_step_forward_pre_hook_disabled'
megatron_policy_worker.py:1657

bea00ff added that read in _finish_train_step_body. The attribute is only set in __init__, at megatron_policy_worker.py:599. And test_megatron_split_state.py:113 builds the worker with object.__new__(MegatronPolicyWorkerImpl) to skip __init__, then hand-sets what it needs — there is already a comment there saying exactly that about another attribute:

# Normally set from get_rank_safe() in __init__, which object.__new__ skips.

So the fixture needs one more line next to it, w._first_train_step_forward_pre_hook_disabled = False, matching the __init__ default. Nothing wrong with the production change.

fast_L1_Functional_Tests_Megatron_4 is unrelated, and I could not attribute it.

A threshold miss in grpo_megatron_generation_topp_topk:

FAIL  max(data["train/token_mult_prob_error"]) < 1.06   ->  1.0856791734695435

2.4% over. With no main baseline I cannot tell whether this PR moved it or it was already close. One thing that might help you decide: git log -L on that assertion shows it has been raised once already, 1.05 → 1.06 in #3392, so it has a history of being tight. If you have a recent green main run of that job, its max value would settle it.

CI quality check is just the roll-up of the four :)

[yuki-97]

/ok to test 8fda593

[yuki-97]

/ok to test 6d303d3

[yuki-97]

/ok to test b60aee0