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

docs/guides/ppo.md

@@ -67,7 +67,7 @@ Async PPO requires non-colocated vLLM generation with `vllm_cfg.async_engine: tr
 
 `max_trajectory_age_steps` is the normal policy-training age limit. The recommended value is `1`; larger values improve overlap but increase off-policy bias in GAE. When `policy_training_start_step > 0`, set `warmup_generation_lead_steps` to a larger value to bank additional rollout batches while the policy is frozen for critic warmup. The collector caps frozen-policy targets at `policy_training_start_step + max_trajectory_age_steps`, so their actual policy-update age remains within the normal limit. The buffer keeps these batches valid through that frontier and then restores the normal age limit. `null` uses `max_trajectory_age_steps` as the generation lead throughout.
 
-Async training stops at `max_num_steps`; the collector cycles the training dataloader as needed. `max_num_epochs` is not supported yet and must be set to `-1`; use `max_num_steps` to control training length. Async checkpoints save the collector dataloader and replay-buffer state together with policy and value state. By default, incomplete restored targets are retained and gap-filled. Setting `drop_incomplete_targets_on_restore: true` discards their restored rows and fills the target from subsequent dataloader prompts; it does not regenerate the original prompts.
+Async training stops at `max_num_steps`; the collector cycles the training dataloader as needed. `max_num_epochs` is not supported yet and must be set to `-1`; use `max_num_steps` to control training length. This is a v1-only convention — the SingleController path rejects any non-positive `max_num_epochs`, so set a positive value there (see [single-controller.md](single-controller.md#migrating-a-legacy-async-config)). Async checkpoints save the collector dataloader and replay-buffer state together with policy and value state. By default, incomplete restored targets are retained and gap-filled. Setting `drop_incomplete_targets_on_restore: true` discards their restored rows and fills the target from subsequent dataloader prompts; it does not regenerate the original prompts.
 
 ### Value Model Configuration
 
@@ -202,6 +202,39 @@ ppo:
 
 During warmup, generation and environment scoring still run normally — only policy weight updates are skipped.
 
+### Warm-Starting the Critic
+
+Warmup can be paid once, offline, instead of inside every PPO run: pretrain the value model separately and point a fresh PPO run at that checkpoint.
+
+```yaml
+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.
+
+The path is a `step_<n>` directory holding a `value/` subtree — the layout a PPO or critic-pretraining run checkpoints. The critic restores its weights from it, plus optimizer moments and LR-scheduler state whenever the seed carries them; a seed written with `save_optimizer: false` restores weights only and warns. Setup rejects a path with no `value/weights` subtree rather than letting the critic start cold behind the message above. Nothing else is read: the policy starts from the base model, the dataloader from the beginning, and the step counter at 0.
+
+**The seed and this run must agree on the value scheduler.** The seed is restored through the ordinary resume path, so Megatron compares nine scheduler fields against the ones this run builds and raises on the first mismatch — `use_checkpoint_opt_param_scheduler` is off, so `OptimizerParamScheduler._check_and_set` asserts equality. Match all of these across the two runs:
+
+- `value.megatron_cfg.optimizer.lr` and `.min_lr` — they feed `max_lr`/`min_lr` and are the *first* two fields checked. They live in the optimizer block, not the scheduler block.
+- `value.megatron_cfg.scheduler`.
+- `value.train_global_batch_size` — it multiplies `lr_decay_steps`, `wd_incr_steps` and `lr_warmup_steps`.
+- the tick budget `train_iters`. A synchronous run sets it to `min(max_num_steps, max_num_epochs × len(dataloader)) × ppo_epochs`; an async run sets it to `max_num_steps × ppo_epochs`, since async requires `max_num_epochs: -1`. `len(dataloader)` is prompt batches per epoch, so on a synchronous run the dataset size and `num_prompts_per_step` are part of the budget whenever the epoch term is the smaller one — as it is for the shipped recipes that set `max_num_epochs: 15`. Matching `max_num_steps` and `ppo_epochs` alone is not enough there.
+
+A mismatch fails during critic init. Which field is named depends on which input differs: a batch-size difference reports `warmup iterations`, a learning-rate difference reports `learning rate`.
+
+```
+AssertionError: OptimizerParamScheduler: class input value <X> and checkpointvalue <Y> for total number of weight decay iterations do not match
+```
+
+(`checkpointvalue` runs together in the upstream message; search for it as written.)
+
+Carrying the seed's schedule over is deliberate — it is what lets the post-warmup LR continue instead of restarting. Set `value.megatron_cfg.scheduler.override_opt_param_scheduler: true` if you would rather this run's settings win; the seed's step position is still restored, so the critic resumes at the seed's tick count rather than at step 0. All of this is Megatron-only — a DTensor critic (the default in `ppo_math_1B.yaml`) loads the seed's scheduler state with no comparison and has no override knob.
+
+The warm start applies to a fresh run only: once the run has written a checkpoint of its own, that checkpoint wins. That is what lets the setting stay in the config across resumes — a resubmitted run restores its own critic instead of re-seeding from the pretrained one, with no config edit in between.
+
 ## Loss
 
 ### Policy Loss
@@ -238,6 +271,7 @@ ppo:
   max_num_steps: 100000
   ppo_epochs: 4
   policy_training_start_step: 0
+  warm_start_value_checkpoint: null
   val_period: 20
   val_at_start: true
   val_at_end: false
@@ -294,6 +328,7 @@ value_loss_fn:
 **PPO-specific parameters:**
 - **`ppo.ppo_epochs`**: Number of training updates per rollout batch
 - **`ppo.policy_training_start_step`**: Number of critic-only warmup steps before policy training begins
+- **`ppo.warm_start_value_checkpoint`**: Checkpoint step directory whose `value/` seeds the critic on a fresh run. See [Warm-Starting the Critic](#warm-starting-the-critic)
 - **`ppo.seq_logprob_error_threshold`**: Nullable sequence-level multiplicative probability-error threshold. PPO always logs sequence-level train/generation mismatch metrics; when this is set, sequences above the threshold are excluded from advantage and loss computation.
 - **`ppo.async_ppo`**: Enables replay-buffer-based asynchronous PPO. See [Asynchronous PPO](#asynchronous-ppo) for requirements and staleness controls.
 - **`ppo.adv_estimator.name`**: Set to `"gae"` for GAE advantage estimation (PPO default)

docs/guides/single-controller.md

@@ -54,7 +54,7 @@ uv run examples/run_grpo_single_controller.py --config <your-sc.yaml>
       use_importance_sampling_correction: true
     ```
 
-5. **(PPO) Set `ppo:` instead of `grpo:`** — the two algorithm blocks are mutually exclusive, and SC reads every step setting from whichever one is present. A PPO run also needs `value:`, `value_loss_fn:` and `ppo.adv_estimator.name: gae` (same schemas as legacy PPO), a Megatron critic, and `policy.offload_optimizer_for_logprob: true`, which is what keeps the policy optimizer off the GPU while the critic runs. `ppo.policy_training_start_step: N` gives the usual critic warmup: for the first N steps the policy is neither trained nor refit, while the critic trains every step.
+5. **(PPO) Set `ppo:` instead of `grpo:`** — the two algorithm blocks are mutually exclusive, and SC reads every step setting from whichever one is present. A PPO run also needs `value:`, `value_loss_fn:` and `ppo.adv_estimator.name: gae` (same schemas as legacy PPO), a Megatron critic, and `policy.offload_optimizer_for_logprob: true`, which is what keeps the policy optimizer off the GPU while the critic runs. `ppo.policy_training_start_step: N` gives the usual critic warmup: for the first N steps the policy is neither trained nor refit, while the critic trains every step. `ppo.warm_start_value_checkpoint` seeds that critic from another run's checkpoint instead, so a fresh run can skip the online warmup entirely — see [Warm-Starting the Critic](./ppo.md#warm-starting-the-critic).
 
 ## Async-RL Knobs and Sampler Modes
 
@@ -160,6 +160,8 @@ The [legacy async GRPO](./async-grpo.md) (`grpo.async_grpo.enabled: true` under
 
 SC reads its async knobs from `async_rl:` and **requires `grpo.async_grpo: null`** (or `ppo.async_ppo: null` on a PPO run) — `run_grpo_single_controller.py` raises if a legacy block is still present, so null it out when porting rather than leaving it in place.
 
+Do not carry `max_num_epochs: -1` across either. [ppo.md](./ppo.md#asynchronous-ppo) requires that value for legacy async PPO, but SC has no `-1` convention: the rollout pump gates on `_current_epoch < max_num_epochs`, so any non-positive value trains zero steps and exits successfully. Setup rejects it — set a positive `max_num_epochs` and bound the run with `max_num_steps`.
+
 | Legacy `grpo.async_grpo.*` / `ppo.async_ppo.*` | SC equivalent `async_rl.*` |
 | -------------------------- | -------------------------- |
 | `enabled: true` | Implicit — SC is always async; use `sampler.max_lookahead_versions: 0` for sync semantics, `>= 1` for async |

examples/configs/ppo_math_1B.yaml

@@ -10,6 +10,10 @@ ppo:
   max_num_steps: 100000
   ppo_epochs: 4
   policy_training_start_step: 0  # number of PPO steps of critic-only warmup before policy training begins
+  # step_<n> dir of a critic-pretrain run whose value/ seeds the critic.
+  # Only a fresh run reads it; a resume ignores it and restores the critic from
+  # its own checkpoint, so it can stay set.
+  warm_start_value_checkpoint: null
   val_period: 20
   val_at_start: true
   val_at_end: false

nemo_rl/algorithms/ppo.py

@@ -96,7 +96,11 @@
 from nemo_rl.models.policy.lm_policy import Policy
 from nemo_rl.models.value import Value, ValueConfig
 from nemo_rl.models.value.interfaces import ValueInterface
-from nemo_rl.utils.checkpoint import CheckpointingConfig, CheckpointManager
+from nemo_rl.utils.checkpoint import (
+    CheckpointingConfig,
+    CheckpointManager,
+    validate_warm_start_checkpoint,
+)
 from nemo_rl.utils.logger import (
     Logger,
     LoggerConfig,
@@ -185,6 +189,10 @@ class PPOConfig(BaseModel, extra="allow"):
     # Value model trains from step 0; policy training is skipped for
     # total_steps < this value. Default 0 (train from start).
     policy_training_start_step: int = 0
+    # Step directory of a critic-pretrain run whose value/ seeds this run's critic.
+    # Only a fresh run reads it; a resume ignores it and restores the critic from
+    # its own checkpoint, so it can stay set.
+    warm_start_value_checkpoint: str | None = None
     # Nullable sequence-level multiplicative probability-error threshold.
     # None logs metrics without masking; values above the threshold are excluded.
     seq_logprob_error_threshold: float | None = None
@@ -705,8 +713,14 @@ def setup(
     worker_init_timing_metrics = {}
 
     weights_path, optimizer_path = checkpointer.get_resume_paths(last_checkpoint_path)
+    # Only a fresh run reads this; a resume ignores it and restores the critic from
+    # its own checkpoint, so the key can stay in the config.
+    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,
+        last_checkpoint_path or warm_start,
         model_component="value",
     )
 

nemo_rl/algorithms/single_controller_utils/config.py

@@ -731,6 +731,16 @@ def _validate_algo_settings(master_config: MasterConfig) -> None:
     """
     algo_cfg = algo_config(master_config)
 
+    # None means no epoch bound. SC has no -1 convention though: the rollout pump
+    # gates on _current_epoch < max_num_epochs, so <= 0 trains nothing and exits 0.
+    if algo_cfg.max_num_epochs is not None and algo_cfg.max_num_epochs <= 0:
+        raise ValueError(
+            f"max_num_epochs={algo_cfg.max_num_epochs} trains zero steps on the "
+            "SingleController path, which does not use the -1 convention that v1 "
+            "async PPO requires. Set a positive max_num_epochs and bound the run "
+            "with max_num_steps."
+        )
+
     # SC reads none of these on either path, so an enabled one describes shaping
     # this run does not do. Async GRPO rejects three of them the same way.
     unsupported = [

nemo_rl/algorithms/single_controller_utils/setup.py

@@ -100,7 +100,10 @@
 )
 from nemo_rl.models.policy.tq_policy import TQPolicy
 from nemo_rl.models.value.tq_value import TQValue
-from nemo_rl.utils.checkpoint import CheckpointManager
+from nemo_rl.utils.checkpoint import (
+    CheckpointManager,
+    validate_warm_start_checkpoint,
+)
 from nemo_rl.weight_sync import WeightSynchronizer, create_weight_synchronizer
 
 
@@ -560,7 +563,7 @@ def _clamp_max_num_steps(
     """Clamp max_num_steps to max_num_epochs * len(dataloader)."""
     algo_cfg = algo_config(master_config)
     max_num_epochs = algo_cfg.max_num_epochs
-    if max_num_epochs is None or max_num_epochs <= 0:
+    if max_num_epochs is None:
         return
     algo_cfg.max_num_steps = min(
         algo_cfg.max_num_steps,
@@ -789,10 +792,17 @@ def setup_single_controller(
     )
     save_state = _get_grpo_save_state(loaded_state)
     weights_path, optimizer_path = checkpointer.get_resume_paths(last_checkpoint_path)
-    value_weights_path, value_optimizer_path = checkpointer.get_resume_paths(
-        last_checkpoint_path,
-        model_component="value",
-    )
+    if is_ppo_run(master_config):
+        # Only a fresh run reads this; a resume ignores it and restores the critic
+        # from its own checkpoint, so the key can stay in the config.
+        warm_start = master_config.ppo.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",
+        )
 
     # ==========================
     # Setup Dataset & Environments

nemo_rl/utils/checkpoint.py

@@ -61,6 +61,36 @@ def _load_megatron_common_state_dict(iteration_dir: Path) -> dict[str, Any]:
     return load_common_state_dict(str(iteration_dir))
 
 
+def validate_warm_start_checkpoint(
+    warm_start: PathLike, model_component: str = "value"
+) -> None:
+    """Reject a warm-start checkpoint whose model subtree does not exist.
+
+    This is the one checkpoint path that comes straight from user config, and
+    get_resume_paths never stats it -- an unresolvable path would train a cold
+    model behind a line claiming a warm start. Callers must gate this on a fresh
+    run, since a resume ignores the warm start entirely.
+
+    Args:
+        warm_start: step_<n> directory the user pointed the warm start at. A Hydra
+            override written as ``key=`` with an unset variable arrives as "".
+        model_component: Subtree the seed must carry, matching get_resume_paths.
+    """
+    if not str(warm_start).strip():
+        raise ValueError(
+            "warm_start_value_checkpoint is empty. A Hydra override written as "
+            "`ppo.warm_start_value_checkpoint=` with an unset variable produces "
+            "this; point it at a step_<n> directory or drop the override."
+        )
+    if not (Path(warm_start) / model_component / "weights").exists():
+        raise ValueError(
+            f"warm_start_value_checkpoint={str(warm_start)!r} has no "
+            f"{model_component}/weights subtree, so the {model_component} model "
+            "would silently start cold. Point it at a step_<n> directory from a "
+            "critic-pretrain or PPO run."
+        )
+
+
 class PretrainedCheckpointConfig(TypedDict):
     """Configuration for restoring initial weights from a pre-existing Megatron checkpoint.
 

tests/unit/algorithms/test_ppo.py

@@ -1970,6 +1970,15 @@ def test_megatron_train_iters_matches_ppo_training_limit(
     assert config.value["megatron_cfg"]["train_iters"] == expected_train_iters
 
 
+def test_ppo_setup_rejects_a_warm_start_that_does_not_resolve(monkeypatch, tmp_path):
+    """A fresh run's warm-start checkpoint must resolve, or the critic would silently start cold."""
+    config = _make_noncolocated_setup_config()
+    config.ppo.warm_start_value_checkpoint = str(tmp_path / "typo")
+
+    with pytest.raises(ValueError, match="would silently start cold"):
+        _run_noncolocated_setup(monkeypatch, config)
+
+
 def test_colocated_setup_keeps_single_cluster_and_skips_collective(monkeypatch):
     """The default colocated setup remains unchanged by the cluster split."""
     config = _make_noncolocated_setup_config()

tests/unit/reference_configs/ppo_math_1B_megatron.yaml

@@ -10,6 +10,7 @@ ppo:
   max_num_steps: 100000
   ppo_epochs: 4
   policy_training_start_step: 0  # number of PPO steps of critic-only warmup before policy training begins
+  warm_start_value_checkpoint: null  # step_<n> dir of a critic-pretrain run whose value/ seeds the critic
   val_period: 20
   val_at_start: true
   val_at_end: false