diff --git a/docs/guides/ppo.md b/docs/guides/ppo.md index 8d83e01fca7..1cc21eaad27 100644 --- a/docs/guides/ppo.md +++ b/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_` 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 and checkpointvalue 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) diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index 401b7dc78d8..f5815ec359f 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -54,7 +54,7 @@ uv run examples/run_grpo_single_controller.py --config 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 | diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index e0fea886a4e..043f04e6556 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/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_ 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 diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index 4db2867a7ed..8778cd55401 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/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", ) diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index ce28051b297..1f9259fce0a 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/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 = [ diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 313c3a94ab6..ad577f0729b 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/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 diff --git a/nemo_rl/utils/checkpoint.py b/nemo_rl/utils/checkpoint.py index 2636021d40f..42938fa925f 100644 --- a/nemo_rl/utils/checkpoint.py +++ b/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_ 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_ 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_ directory from a " + "critic-pretrain or PPO run." + ) + + class PretrainedCheckpointConfig(TypedDict): """Configuration for restoring initial weights from a pre-existing Megatron checkpoint. diff --git a/tests/unit/algorithms/test_ppo.py b/tests/unit/algorithms/test_ppo.py index 8f51fa4f9fa..2674d85f5ee 100644 --- a/tests/unit/algorithms/test_ppo.py +++ b/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() diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml index 71a9251a559..71b3bff39a8 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/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_ dir of a critic-pretrain run whose value/ seeds the critic val_period: 20 val_at_start: true val_at_end: false diff --git a/tests/unit/single_controller/test_ppo_setup.py b/tests/unit/single_controller/test_ppo_setup.py index 6fe7853f3bf..6c914abab0e 100644 --- a/tests/unit/single_controller/test_ppo_setup.py +++ b/tests/unit/single_controller/test_ppo_setup.py @@ -408,6 +408,19 @@ def test_rejects_colocated_generation(self, make_config): with pytest.raises(ValueError, match="colocated.enabled=false"): validate_single_controller_config(mc) + @pytest.mark.parametrize( + "make_config", [_ppo_master_config, _make_master_config], ids=["ppo", "grpo"] + ) + @pytest.mark.parametrize("epochs", [-1, 0], ids=["negative", "zero"]) + def test_rejects_a_non_positive_max_num_epochs(self, make_config, epochs): + """docs/guides/ppo.md tells async-PPO users to set -1; carried onto SC it + makes the rollout pump's epoch gate unsatisfiable and the run exits 0.""" + mc = make_config() + algo_config(mc).max_num_epochs = epochs + + with pytest.raises(ValueError, match="trains zero steps"): + validate_single_controller_config(mc) + def test_grpo_is_free_to_use_any_sampler(self): mc = _make_master_config() mc.async_rl.sampler = WindowedSamplerConfig() @@ -570,6 +583,97 @@ def test_grpo_run_builds_no_critic(self, patched_ppo_factories): patched_ppo_factories["policy"].offload_to_cpu.assert_not_called() +class TestValueWarmStart: + def test_fresh_run_builds_the_critic_from_the_warm_start( + self, patched_ppo_factories, tmp_path + ): + seed = tmp_path / "critic_pretrain" / "step_370" + (seed / "value" / "weights").mkdir(parents=True) + (seed / "value" / "optimizer").mkdir() + mc = _ppo_master_config( + ppo=PPOConfig.model_construct( + max_num_steps=100, + warm_start_value_checkpoint=str(seed), + **_STEP_CONFIG, + ) + ) + mc.checkpointing["checkpoint_dir"] = str(tmp_path / "run") + + setup_single_controller(mc, tokenizer=MagicMock(pad_token_id=0)) + + value_kwargs = patched_ppo_factories["_build_value"].call_args.kwargs + assert value_kwargs["weights_path"] == seed / "value" / "weights" + assert value_kwargs["optimizer_path"] == seed / "value" / "optimizer" + # The policy is untouched by a warm start: pi_0 comes from the base model. + trainer_kwargs = patched_ppo_factories["_build_trainer"].call_args.kwargs + assert trainer_kwargs["weights_path"] is None + + def test_unset_leaves_the_critic_cold(self, patched_ppo_factories, tmp_path): + mc = _ppo_master_config() + mc.checkpointing["checkpoint_dir"] = str(tmp_path / "run") + + setup_single_controller(mc, tokenizer=MagicMock(pad_token_id=0)) + + value_kwargs = patched_ppo_factories["_build_value"].call_args.kwargs + assert value_kwargs["weights_path"] is None + assert value_kwargs["optimizer_path"] is None + + @pytest.mark.parametrize( + ("warm_start", "message"), + [ + ("", "warm_start_value_checkpoint is empty"), + ("{tmp}/typo", "would silently start cold"), + ], + ids=["empty_string", "typo"], + ) + def test_rejects_a_warm_start_that_does_not_resolve( + self, patched_ppo_factories, tmp_path, warm_start, message + ): + """get_resume_paths never stats the path, so an unresolvable one would + train a cold critic behind the line claiming a warm start. A Hydra + override written as `key=` with an unset variable arrives as "".""" + mc = _ppo_master_config( + ppo=PPOConfig.model_construct( + max_num_steps=100, + warm_start_value_checkpoint=warm_start.format(tmp=tmp_path), + **_STEP_CONFIG, + ) + ) + mc.checkpointing["checkpoint_dir"] = str(tmp_path / "run") + + with pytest.raises(ValueError, match=message): + setup_single_controller(mc, tokenizer=MagicMock(pad_token_id=0)) + + patched_ppo_factories["_build_value"].assert_not_called() + + def test_resume_ignores_the_warm_start(self, patched_ppo_factories, tmp_path): + """Re-seeding a resumed run would discard the critic's own progress, so + the run's checkpoint has to win over the key left in the config.""" + seed = tmp_path / "critic_pretrain" / "step_370" + (seed / "value" / "weights").mkdir(parents=True) + (seed / "value" / "optimizer").mkdir() + step_5 = tmp_path / "run" / "step_5" + for component in ("policy", "value"): + (step_5 / component / "weights").mkdir(parents=True) + (step_5 / component / "optimizer").mkdir() + (step_5 / "training_info.json").write_text("{}") + mc = _ppo_master_config( + ppo=PPOConfig.model_construct( + max_num_steps=100, + warm_start_value_checkpoint=str(seed), + **_STEP_CONFIG, + ) + ) + mc.checkpointing["checkpoint_dir"] = str(tmp_path / "run") + + with patch.object(sc_setup_mod, "load_dataloader_state"): + setup_single_controller(mc, tokenizer=MagicMock(pad_token_id=0)) + + value_kwargs = patched_ppo_factories["_build_value"].call_args.kwargs + assert value_kwargs["weights_path"] == step_5 / "value" / "weights" + assert value_kwargs["optimizer_path"] == step_5 / "value" / "optimizer" + + def _cluster_config(mc: MasterConfig, *, colocated: bool, backend: str) -> MasterConfig: """Fill in the cluster / generation keys _build_clusters reads.""" mc.cluster = {