feat(ppo): warm-start the value model from a critic-pretrain checkpoint

[yuki-97]

What does this PR do ?

Adds ppo.warm_start_value_checkpoint: a fresh PPO run seeds its critic from a separately pretrained value checkpoint, so critic warmup can be paid once offline instead of inside every PPO run.

Today the only way to reach a usable critic is ppo.policy_training_start_step: N — N steps during which generation and environment scoring run at full cost while only the critic trains. Where generation dominates the step cost (SWE, long-horizon agentic runs), that is the most expensive way to buy a value function. Pretraining the critic on already-collected rollouts and starting PPO from it moves that cost off the PPO allocation entirely.

How it works. The value component's resume paths are resolved independently of the policy's (CheckpointManager.get_resume_paths(..., model_component="value")), so the whole feature is a choice of which checkpoint that call reads:

warm_start = ppo_config.warm_start_value_checkpoint
if last_checkpoint_path is None and warm_start is not None:
    validate_warm_start_checkpoint(warm_start)
    print(f"🔥 Warm-starting the value model from {warm_start}")
value_weights_path, value_optimizer_path = checkpointer.get_resume_paths(
    last_checkpoint_path or warm_start,
    model_component="value",
)

validate_warm_start_checkpoint rejects an empty string and a path with no value/weights. get_resume_paths never stats what it is given, so without it a typo prints the 🔥 line over a cold critic.

The seed is a step_<n> directory holding a value/ subtree — what a PPO or critic-pretraining run already checkpoints — and the critic restores its weights, plus optimizer moments and LR-scheduler state whenever the seed carries them. On Megatron the seed and this run must then agree on the value scheduler — docs/guides/ppo.md lists the fields. Nothing else is read: the policy starts from the base model, the dataloader from the beginning, and the step counter at 0. Those are a fresh run's defaults, so no extra handling is needed for them.

Fresh runs only, deliberately. When the run already has a checkpoint of its own, that checkpoint wins and the setting is not read. This is what lets the key stay in the config across resumes: a requeued run restores its own critic rather than re-seeding from the pretrained one, with no config edit between submissions. Making the warm start win instead would silently discard the critic's progress on every resume.

Both paths. The knob lives on PPOConfig, and SC's algo_config() returns that same object on a PPO run, so legacy ppo.setup and setup_single_controller are both covered by one field.

Issues

None closed.

Usage

ppo:
  warm_start_value_checkpoint: /path/to/critic_pretrain_run/step_370
  policy_training_start_step: 0  # the critic is already warm

policy_training_start_step: 0 is the natural pairing, but keeping a short online warmup is equally valid — it calibrates the seeded critic on this run's own rollout distribution before the policy starts moving. A nonzero value is also what async_ppo.warmup_generation_lead_steps requires (async_rl.sampler.warmup_lookahead_versions on the SingleController path); both must be null when it is 0.

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? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• Byte-identical on a PPO run when unset. The default is null, so last_checkpoint_path or warm_start reduces to last_checkpoint_path.

• One shape change on the SC path. Value resume-path resolution now sits inside if is_ppo_run(master_config), so a GRPO run no longer resolves value paths it never uses. Side effect on the GRPO path: a GRPO resume previously warned Optimizer state not found about a value model it does not have. Warning text only.

• Tests. test_ppo_setup.py::TestValueWarmStart covers the near end on the SC path — the seeded paths reach _build_value while the trainer's weights_path stays None, so the policy is not seeded, plus test_ppo.py::test_ppo_setup_rejects_a_warm_start_that_does_not_resolve for the legacy ppo.setup copy. test_{megatron,dtensor}_value_worker.py::test_value_worker_checkpoint_save_and_load covers the far end: a restored value worker reproduces the saved model's predictions exactly, so a wrong checkpoint or a partial restore fails.

  uv run --group test pytest tests/unit/single_controller/test_ppo_setup.py tests/unit/algorithms/test_ppo.py tests/unit/utils/test_checkpoint.py
  # 2 GPUs + HF_TOKEN
  uv run --group test --extra mcore pytest tests/unit/models/value/test_megatron_value_worker.py::test_value_worker_checkpoint_save_and_load

• One unrelated fix to resolve #3773 (comment). SC has no -1 convention — the rollout pump gates on _current_epoch < max_num_epochs, so an async-PPO user following docs/guides/ppo.md and carrying max_num_epochs: -1 onto SC got zero optimizer steps and exit status 0. _validate_algo_settings now rejects a non-positive value and names the remedy; _clamp_max_num_steps no longer needs its own <= 0 early return.

• No functional test. A seeded run and a cold control do not share rollouts (async vLLM is not reproducible even at a fixed seed), so a cross-run critic-loss comparison is a statistical claim, not an identity check. The one property e2e would add — a real SC PPO run writing a loadable step_<n>/value/ subtree — is already asserted by tests/functional/ppo_async_single_controller.sh.

[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 6adf1a4

[yuki-97]

Team review of the critic warm-start path, run across five agents (RL/guidelines, bug-finder, test, design, devil's advocate) with an adversarial pass over every finding.

15 raw findings deduped to 11 clusters; 4 survived confidence + impact filtering. Seven were dropped as confirmed-but-not-worth-a-round-trip, and one (a claim that the new test fixture models an unreachable checkpoint layout) was disputed outright — get_resume_paths returns the identical tuple on both branches, so the fixture is fine.

No lint/type results: this host has no uv or pre-commit, and nothing was executed except two standalone stdlib repros. Test suggestions are unverified locally and flagged as such.

Generated by Claude Code

[yuki-97]

/ok to test c14539c

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 7107664 (PR #3821 from yukih/sc-warm-start-critic)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of yukih/sc-drop-incomplete-target branch (fast-forward)

All submodule changes look good! ✨

[terrykong]

Team review of the critic warm-start path: seven agents (RL/guidelines, Megatron-Bridge, bug-finder, test, design, comment-audit, devil's advocate), with an adversarial pass over every finding including my own.

Nice feature — moving critic warmup off the PPO allocation is a real saving where generation dominates, and the resume-wins precedence is the right call.

Three comments below. Eleven candidate findings were cut to three: four were disputed outright by the adversarial pass, and the rest folded into the surviving comments. Two design calls are worth affirming rather than questioning:

• Routing the seed through the resume path instead of checkpointing.pretrained_checkpoint is correct, and non-obviously so. That path auto-sets finetune=True in megatron-bridge, which suppresses optimizer/RNG state and triggers the value worker's hide_loss_modules hook that drops output_layer.* — so it would restore the trunk and randomly re-initialize the value head, discarding exactly what critic pretraining produces. megatron_value_worker.py:97 already states the rule.
• Gating value-path resolution on is_ppo_run incidentally fixes a pre-existing GRPO wart — a GRPO resume previously stat'd step_N/value/optimizer, missed, and emitted a spurious "Optimizer state not found" warning about a value model the run does not have. Worth a line in the PR body, since it is a behavior change on the GRPO path in a PR titled feat(ppo).

Not verified — no runnable environment. Everything here is static reading plus upstream source at the pinned SHAs (Megatron-LM 14346b65, Megatron-Bridge 8c46dc42). No test and no linter was executed: this tree's uv.lock is linux-only and uv run / pre-commit both fail on macOS with Failed to parse entry: nemo-gym. I make no claim about lint status, and the suggested test is unrun.

FYI, no action: the PR description's "How it works" snippet is stale — it still shows the call site before validate_warm_start_checkpoint was added in c14539c.

Generated by Claude Code

[yuki-97]

/ok to test 95c03dc

[yuki-97]

/ok to test c39b179