diff --git a/docs/about/algorithms/ppo.md b/docs/about/algorithms/ppo.md index 21b152c294..a993a261b2 100644 --- a/docs/about/algorithms/ppo.md +++ b/docs/about/algorithms/ppo.md @@ -69,6 +69,7 @@ ppo: gae_lambda: 0.95 gae_gamma: 1.0 ppo_epochs: 4 + critic_ppo_epochs: ${ppo.ppo_epochs} policy_training_start_step: 0 value_loss_fn: @@ -79,6 +80,10 @@ value: model_name: "Qwen/Qwen2.5-1.5B" ``` +`ppo_epochs` and `critic_ppo_epochs` are independent positive integers. The +example uses interpolation so the critic follows the actor epoch count unless +you explicitly override it. + ## Additional Resources - [PPO Paper](https://arxiv.org/abs/1707.06347) diff --git a/docs/guides/ppo.md b/docs/guides/ppo.md index 1cc21eaad2..d1b3f0a8ac 100644 --- a/docs/guides/ppo.md +++ b/docs/guides/ppo.md @@ -59,7 +59,7 @@ When only one node remains for policy and generation after other resources are r ### Asynchronous PPO -Set `ppo.async_ppo.enabled: true` to overlap rollout generation with training. A background collector fills a replay buffer on the non-colocated vLLM GPUs while the policy and value model train on their shared cluster. Values and policy/reference log probabilities are recomputed when a trajectory is sampled, then PPO runs GAE and its normal `ppo_epochs` updates before publishing one new policy version to vLLM. +Set `ppo.async_ppo.enabled: true` to overlap rollout generation with training. A background collector fills a replay buffer on the non-colocated vLLM GPUs while the policy and value model train on their shared cluster. Values and policy/reference log probabilities are recomputed when a trajectory is sampled, then PPO runs GAE, all `critic_ppo_epochs` critic updates, and all `ppo_epochs` policy updates before publishing one new policy version to vLLM. Async PPO reuses the trajectory collector, replay buffer, and weight-versioning infrastructure described in the [Async GRPO guide](async-grpo.md); this section focuses on PPO-specific behavior and constraints. @@ -175,10 +175,10 @@ The PPO training loop, [ppo_train](../../nemo_rl/algorithms/ppo.py), follows thi 3. **Value inference**: the value model predicts per-token state values 4. **Logprob computation**: the policy computes log probabilities for advantage estimation 5. **Advantage estimation**: GAE computes advantages using value predictions and rewards -6. **Value training**: the critic is updated first (critic-before-actor, following [veRL](https://arxiv.org/abs/2412.09613)) -7. **Policy training**: the actor is updated with the clipped surrogate objective +6. **Value training**: the critic completes all of its updates first +7. **Policy training**: the actor completes all of its updates with the clipped surrogate objective -Steps 6–7 repeat `ppo_epochs` times per rollout before generating new responses. +The critic stays resident for all `critic_ppo_epochs` updates, then the policy stays resident for all `ppo_epochs` updates. This avoids moving the colocated models between CPU and GPU after every epoch. ### Multiple Training Steps per Rollout @@ -186,10 +186,14 @@ Unlike GRPO, which performs one training update per rollout, PPO can perform mul ```yaml ppo: - ppo_epochs: 4 # Train 4 times on each rollout batch + ppo_epochs: 4 # actor passes over each rollout batch + critic_ppo_epochs: ${ppo.ppo_epochs} # critic passes; follows actor by default ``` -Each step trains both the critic and the actor on the same advantage estimates computed from the initial rollout. +Each pass uses the same returns and advantage estimates computed from the initial +rollout. Both epoch counts must be at least 1 and can be configured independently; +the exemplar uses interpolation so the critic follows the actor unless explicitly +overridden. ### Critic Warmup @@ -221,7 +225,7 @@ The path is a `step_` directory holding a `value/` subtree — the layout a P - `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. +- the tick budget `train_iters`. A synchronous run sets it to `min(max_num_steps, max_num_epochs × len(dataloader)) × critic_ppo_epochs`; an async run sets it to `max_num_steps × critic_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 `critic_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`. @@ -270,6 +274,7 @@ ppo: max_num_epochs: 100000 max_num_steps: 100000 ppo_epochs: 4 + critic_ppo_epochs: ${ppo.ppo_epochs} policy_training_start_step: 0 warm_start_value_checkpoint: null val_period: 20 @@ -326,7 +331,8 @@ value_loss_fn: ``` **PPO-specific parameters:** -- **`ppo.ppo_epochs`**: Number of training updates per rollout batch +- **`ppo.ppo_epochs`**: Number of actor training updates per rollout batch +- **`ppo.critic_ppo_epochs`**: Number of critic training updates per rollout batch. It can differ from `ppo_epochs`; the exemplar defaults it to `${ppo.ppo_epochs}`. - **`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. diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index edce5196e9..6e6fd69322 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -57,7 +57,7 @@ uv run examples/run_grpo_single_controller.py --config gpus_per_node: 1 # inference GPUs; remainder go to training ``` -3. **One RL step = one training batch.** The batch a step trains on is the whole step (see `validate_single_controller_config` in [nemo_rl/algorithms/single_controller_utils/config.py](../../nemo_rl/algorithms/single_controller_utils/config.py)). A GRPO step is also one optimizer step; a PPO step is `ppo.ppo_epochs` of them over that same batch. +3. **One RL step = one training batch.** The batch a step trains on is the whole step (see `validate_single_controller_config` in [nemo_rl/algorithms/single_controller_utils/config.py](../../nemo_rl/algorithms/single_controller_utils/config.py)). A GRPO step is also one optimizer step. A PPO step applies `ppo.ppo_epochs` actor updates and `ppo.critic_ppo_epochs` critic updates over that same batch. Both counts must be at least 1 and can be configured independently; the exemplar defaults the critic count to `${ppo.ppo_epochs}`. ```python num_prompts_per_step * num_generations_per_prompt == policy.train_global_batch_size @@ -150,7 +150,7 @@ The shipped exemplars cover three of the five modes: Field definitions: - `max_buffered_rollouts` — hard cap on unconsumed rollout groups buffered in the data plane. Validated at setup against the gated sampler's required capacity; a value too small deadlocks the rollout pump, so setup raises instead of silently blocking. Sized from the widest window the run ever uses, so `warmup_lookahead_versions` rather than `max_lookahead_versions` when it is set. -- `min_groups_for_streaming_train` — minimum ready groups the trainer waits for before dispatching a batch. Set to `num_prompts_per_step` for sync/legacy semantics; lower for streaming. (PPO) Must equal `num_prompts_per_step` — the critic has no split train API, so one `train_from_meta` call is one optimizer step, and streaming a step across chunks would step the critic once per chunk. +- `min_groups_for_streaming_train` — minimum ready groups the trainer waits for before dispatching a batch. Set to `num_prompts_per_step` for sync/legacy semantics; lower for streaming. (PPO) Must equal `num_prompts_per_step` — the critic has no split train API, so each critic epoch calls the full-step `train_from_meta` once per chunk. Splitting an RL step across chunks would multiply both models' configured optimizer updates by the number of chunks. - `sampler.warmup_lookahead_versions` (PPO) — lookahead used while `ppo.policy_training_start_step` critic warmup is in progress, shrinking back to `max_lookahead_versions` afterwards. The SC equivalent of `ppo.async_ppo.warmup_generation_lead_steps`. ## Implementation Structure @@ -189,14 +189,14 @@ The SC path splits the async-GRPO loop across a rollout pump and a train pump th #### 5. `_rollout_pump` and `_train_pump` - `_rollout_pump`: pulls prompts from the dataloader, calls `sampler.admit`, dispatches `RolloutManager.generate_and_push`, and honours `max_inflight_prompts` as a backpressure cap. -- `_train_pump`: `sampler.evict → sampler.select → _value_stage (PPO only) → _advantage_stage → _value_train (PPO only) → TQPolicy split API (begin_train_step / train_microbatches_from_meta / finish_train_step) → dp_client.clear_samples`. +- `_train_pump`: `sampler.evict → sampler.select → _value_stage (PPO only) → _advantage_stage → _value_train_epochs (PPO only) → TQPolicy split API (begin_train_step / train_microbatches_from_meta / finish_train_step) → dp_client.clear_samples`. ### Coordination Flow 1. **Driver setup**: `setup_single_controller` builds the worker groups, virtual cluster, dp client, dataloader, `TQReplayBuffer`, `RolloutManager`, and weight synchronizer, and packs them into a `SingleControllerActorArgs` that the entrypoint cloudpickles into the actor. 2. **Actor startup**: `SingleControllerActor` launches `_rollout_pump` and `_train_pump` concurrently as asyncio tasks; both share the same `TQReplayBuffer` and `StalenessSampler`. 3. **Rollout pump loop**: `sampler.admit` gates dispatch against the current trainer version (returning a `target_step` for `in_order`); the pump then reserves a buffer slot, drives `RolloutManager.generate_and_push`, and commits with the observed `start_weight` / `end_weight`. -4. **Train pump loop**: `sampler.evict` drops out-of-window groups, `sampler.select` picks the next batch, `_value_stage` and `_value_train` run the critic forward and its optimizer step on a PPO run, `_advantage_stage` computes advantages, and the TQPolicy split API runs one optimizer step per RL step on GRPO, or `ppo.ppo_epochs` of them on PPO. +4. **Train pump loop**: `sampler.evict` drops out-of-window groups and `sampler.select` picks the next batch. On PPO, `_value_stage` runs the critic forward, `_advantage_stage` computes advantages, and `_value_train_epochs` runs `ppo.critic_ppo_epochs` critic updates. The TQPolicy split API then runs one optimizer step per RL step on GRPO, or `ppo.ppo_epochs` policy updates on PPO. 5. **Weight sync**: after each optimizer step the pump bumps the trainer version, clears rollout permission, calls the weight synchronizer, and re-opens the rollout pump for the next version. ## Relation to Legacy Async GRPO diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index 043f04e655..a58d37c701 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -8,7 +8,9 @@ ppo: max_rollout_turns: 1 max_num_epochs: 100000 max_num_steps: 100000 - ppo_epochs: 4 + ppo_epochs: 4 # actor passes over each rollout batch + # Critic passes over each rollout batch; follows ppo_epochs unless overridden. + critic_ppo_epochs: ${ppo.ppo_epochs} 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 diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index a7902e208b..25c4b60659 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -190,7 +190,11 @@ class PPOConfig(BaseModel, extra="allow"): # When using dynamic sampling, generation prompt batch size will equal # num_prompts_per_step * batch_multiplier batch_multiplier: float = 1.0 + # Number of actor (policy) passes over each rollout batch. ppo_epochs: int = 4 + # Number of critic (value) passes over each rollout batch. Defaults to + # ppo_epochs (see validate_epoch) unless explicitly set. + critic_ppo_epochs: int = 4 reward_shaping: RewardShapingConfig = Field(default_factory=RewardShapingConfig) reward_scaling: RewardScalingConfig = Field(default_factory=RewardScalingConfig) adv_estimator: GAEConfig = Field(default_factory=GAEConfig) @@ -215,7 +219,17 @@ class PPOConfig(BaseModel, extra="allow"): async_ppo: AsyncPPOConfig | None = Field(default_factory=AsyncPPOConfig) @model_validator(mode="after") - def validate_async_warmup_settings(self) -> "PPOConfig": + def validate_epoch(self) -> "PPOConfig": + if "critic_ppo_epochs" not in self.model_fields_set: + self.critic_ppo_epochs = self.ppo_epochs + if self.ppo_epochs < 1: + raise ValueError("ppo.ppo_epochs must be at least 1") + if self.critic_ppo_epochs < 1: + raise ValueError("ppo.critic_ppo_epochs must be at least 1") + return self + + @model_validator(mode="after") + def validate_async_warmup(self) -> "PPOConfig": if ( self.async_ppo is not None and self.async_ppo.enabled @@ -740,12 +754,10 @@ def setup( ) # train_iters is the total scheduler-tick budget. Each Megatron worker - # ticks once per train() call (matching upstream main's per-rollout - # convention), and PPO calls each worker's train() `ppo_epochs` times - # per outer step. So total ticks = (outer steps) * ppo_epochs. - # Scale train_iters accordingly so the configured warmup/decay horizon - # matches the actual scheduler-step count. + # ticks once per train() call, so policy and value need separate budgets + # when their epoch counts or training start steps differ. ppo_epochs = ppo_config.ppo_epochs + critic_ppo_epochs = ppo_config.critic_ppo_epochs async_config = ppo_config.async_ppo if async_config.enabled: outer_training_steps = ppo_config.max_num_steps @@ -754,13 +766,22 @@ def setup( ppo_config.max_num_steps, ppo_config.max_num_epochs * len(dataloader), ) - total_train_iters = outer_training_steps * ppo_epochs - if policy_config.get("megatron_cfg", {}).get("enabled", False): - policy_config["megatron_cfg"]["train_iters"] = total_train_iters + policy_training_steps = max( + outer_training_steps - ppo_config.policy_training_start_step, + 0, + ) + # Megatron-Bridge requires a positive scheduler horizon at setup. The + # scheduler is never advanced when critic warmup spans the whole run. + policy_config["megatron_cfg"]["train_iters"] = max( + policy_training_steps * ppo_epochs, + 1, + ) if value_config.get("megatron_cfg", {}).get("enabled", False): - value_config["megatron_cfg"]["train_iters"] = total_train_iters + value_config["megatron_cfg"]["train_iters"] = ( + outer_training_steps * critic_ppo_epochs + ) # Define initialization functions that will be used in all paths def init_policy(): @@ -1234,7 +1255,7 @@ def ppo_train( Based on the grpo_train loop with PPO-specific modifications: - Value model inference and training (actor-critic) - GAE advantage estimation with value bootstrap - - Multiple training steps per rollout (ppo_epochs) + - Multiple actor and critic training steps per rollout - Configurable policy training start epoch """ timer = Timer() @@ -1273,6 +1294,7 @@ def ppo_train( current_epoch = ppo_save_state["current_epoch"] max_num_epochs = master_config.ppo.max_num_epochs ppo_epochs = master_config.ppo.ppo_epochs + critic_ppo_epochs = master_config.ppo.critic_ppo_epochs # Number of PPO steps to train only the critic before starting policy # training. Despite the legacy name, this is compared against total_steps # (not current_epoch) to match veRL's critic_warmup semantics. @@ -1668,16 +1690,18 @@ def ppo_train( # PPO: Multiple training steps per rollout memory_tracker.snapshot_start_of_stage("Policy train", dir()) - for step in range(ppo_epochs): + + # Actor and critic share the training GPUs. Keep each model + # resident for all of its PPO epochs so their update phases need + # only one onload/offload cycle apiece. + print("▶ Training value...", flush=True) + with timer.time("value_training_prep"): + value_model.prepare_for_training() + for critic_epoch in range(critic_ppo_epochs): print( - f"▶ Step {step + 1}/{ppo_epochs}...", + f"▶ Value epoch {critic_epoch + 1}/{critic_ppo_epochs}...", flush=True, ) - - # Train value model first (critic before actor, matching veRL). - with timer.time("value_training_prep"): - value_model.prepare_for_training() - with ( timer.time("value_training"), managed_span( @@ -1687,32 +1711,36 @@ def ppo_train( **{"rl.iteration": total_steps + 1}, ), ): - print("▶ Training value...", flush=True) value_results = value_model.train( train_data, value_loss_fn, timer=timer, ) + with timer.time("value_training"): + value_model.finish_training() - value_model.finish_training() - - train_results = None - if total_steps >= policy_training_start_step: - if ( - total_steps == policy_training_start_step - and policy_training_start_step > 0 - ): - print( - f" ✓ Critic warmup complete ({policy_training_start_step} steps). " - f"Starting policy training.", - flush=True, - ) - print("▶ Preparing for training...", flush=True) - with timer.time("training_prep"): - policy.prepare_for_training() - POLICY_GENERATION_STALE = True + train_results = None + if total_steps >= policy_training_start_step: + if ( + total_steps == policy_training_start_step + and policy_training_start_step > 0 + ): + print( + f" ✓ Critic warmup complete ({policy_training_start_step} steps). " + f"Starting policy training.", + flush=True, + ) + print("▶ Preparing for training...", flush=True) + with timer.time("training_prep"): + policy.prepare_for_training() + POLICY_GENERATION_STALE = True - print("▶ Training policy...", flush=True) + print("▶ Training policy...", flush=True) + for policy_epoch in range(ppo_epochs): + print( + f"▶ Policy epoch {policy_epoch + 1}/{ppo_epochs}...", + flush=True, + ) with ( timer.time("policy_training"), managed_span( @@ -1727,17 +1755,15 @@ def ppo_train( loss_fn, timer=timer, ) - if step < ppo_epochs - 1: - policy.offload_to_cpu() - if train_results is not None: - print( - f" • Policy loss: {train_results['loss'].mean().item():.4f}" - ) - if value_results is not None: - print( - f" • Value loss: {value_results['loss'].mean().item():.4f}" - ) + if train_results is not None: + print( + f" • Policy loss: {train_results['loss'].mean().item():.4f}" + ) + if value_results is not None: + print( + f" • Value loss: {value_results['loss'].mean().item():.4f}" + ) # Recompute KV scales after policy training if needed if sync_kv_scales: @@ -2166,8 +2192,7 @@ def async_ppo_train( max_trajectory_age_steps = async_config.max_trajectory_age_steps warmup_generation_lead_steps = async_config.resolved_warmup_generation_lead_steps policy_training_start_step = master_config.ppo.policy_training_start_step - if master_config.ppo.ppo_epochs < 1: - raise ValueError("ppo.ppo_epochs must be at least 1") + critic_ppo_epochs = master_config.ppo.critic_ppo_epochs if max_trajectory_age_steps > 1: print( "⚠️ WARNING: max_trajectory_age_steps > 1 increases off-policy " @@ -2678,47 +2703,49 @@ def _raise_if_collector_stopped(waiting_for: str) -> None: if returns is not None: train_data["returns"] = returns - # ---- 7. ppo_epochs inner loop (critic, then actor) ---- - # Each epoch: value on GPU -> train -> off. Then, once past critic - # warmup, policy on GPU -> train -> off (except the last epoch, - # which leaves the policy on GPU for the refit broadcast below). + # ---- 7. Grouped critic epochs, then grouped actor epochs ---- + # Actor and critic share the training GPUs. Keep each model + # resident for its complete update phase to avoid per-epoch + # onload/offload cycles. The policy remains resident after its + # final epoch for the refit broadcast below. # During warmup (step < policy_training_start_step) the policy is # frozen: it is never loaded/trained here, exactly as in sync # ppo_train, so train_results stays None for the step. is_policy_training_step = step >= policy_training_start_step train_results = None value_results = None - for epoch in range(ppo_epochs): - print(f"▶ PPO epoch {epoch + 1}/{ppo_epochs}...") - with timer.time("value_training_prep"): - value_model.prepare_for_training() + + with timer.time("value_training_prep"): + value_model.prepare_for_training() + for critic_epoch in range(critic_ppo_epochs): + print(f"▶ Value epoch {critic_epoch + 1}/{critic_ppo_epochs}...") with timer.time("value_training"): value_results = value_model.train( train_data, value_loss_fn, timer=timer, ) - value_model.finish_training() + with timer.time("value_training"): + value_model.finish_training() - if is_policy_training_step: - if ( - step == policy_training_start_step - and policy_training_start_step > 0 - and epoch == 0 - ): - print( - f" ✓ Critic warmup complete ({policy_training_start_step} " - "steps). Starting policy training.", - flush=True, - ) - with timer.time("training_prep"): - policy.prepare_for_training() + if is_policy_training_step: + if ( + step == policy_training_start_step + and policy_training_start_step > 0 + ): + print( + f" ✓ Critic warmup complete ({policy_training_start_step} " + "steps). Starting policy training.", + flush=True, + ) + with timer.time("training_prep"): + policy.prepare_for_training() + for policy_epoch in range(ppo_epochs): + print(f"▶ Policy epoch {policy_epoch + 1}/{ppo_epochs}...") with timer.time("policy_training"): train_results = policy.train( train_data, loss_fn, timer=timer ) - if epoch < ppo_epochs - 1: - policy.offload_to_cpu() # ---- 8. Refit once after all PPO epochs ---- # Warmup still advances the replay-buffer version, but skips the diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 23f8ad00c5..7e36fb8bfd 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -27,7 +27,8 @@ → _advantage_stage(meta) → dp_client.get_samples(...) → adv_estimator.compute_advantage(...) → dp_client.put_samples(...) - → _value_train(meta) (PPO only) → value.train_from_meta(...) + → _value_train_epochs(meta) (PPO only) + → value.train_from_meta(...) Value → dp_client.get_samples(...) (via its own client) → trainer.begin/train_microbatches/finish_train_step (split API, driver-side TQPolicy via asyncio.to_thread) @@ -202,6 +203,9 @@ def __init__( self._is_ppo: bool = is_ppo_run(master_config) # GRPO has no epoch knob: it makes one optimizer step per RL step. self._ppo_epochs: int = self._algo_cfg.ppo_epochs if self._is_ppo else 1 + self._critic_ppo_epochs: int = ( + self._algo_cfg.critic_ppo_epochs if self._is_ppo else 1 + ) self._message_level_advantage_penalties_enabled = ( self._algo_cfg.invalid_tool_call_advantage is not None or self._algo_cfg.malformed_thinking_advantage is not None @@ -1670,9 +1674,11 @@ async def _train_pump(self) -> None: chunk -- the value workers have no split train API yet (#2625). b. Policy model: train_microbatches_from_meta, which only accumulates gradients. - c. PPO only: 3a-3b repeat ppo.ppo_epochs times, and the policy's - optimizer step closes here rather than in 5 -- a PPO step is - one chunk, so there is nothing to accumulate across chunks. + c. PPO only: all critic updates run before all policy updates. + Their counts are ppo.critic_ppo_epochs and ppo.ppo_epochs, + respectively. Each policy optimizer step closes here rather + than in 5 -- a PPO step is one chunk, so there is nothing to + accumulate across chunks. 4. Clear the batch. dp_client.clear_samples on the consumed sample_ids. 5. Train the policy model (GRPO) -- finish_train_step all_reduces the accumulated gradients, rescales, and runs optimizer.step. @@ -1840,35 +1846,37 @@ async def _train_pump(self) -> None: # optimizer step with the next chunk. # GRPO runs one iteration: F/B only, its optimizer step is in 5. - # For PPO, each epoch is a full optimizer step for both models. + # For PPO, each actor epoch and critic epoch is a full optimizer + # step. Group each model's epochs under one residency cycle so + # the colocated models do not move between CPU and GPU per epoch. # TODO(#2625): value_result, policy_result only record the last epoch's metrics. # That matches ppo.py for the losses; total_flops is additive and undercounted. - for epoch in range(self._ppo_epochs): - # Value model first, then policy, as in the legacy PPO epoch loop. - if self._is_ppo: - with self._timer.time("value_training"): - value_result = await self._value_train(train_meta) - - if is_policy_training_step: - if ( - self._is_ppo - and self._train_steps == policy_training_start_step - and policy_training_start_step > 0 - and epoch == 0 - ): - print( - f" ✓ Critic warmup complete ({policy_training_start_step} " - "steps). Starting policy training.", - flush=True, - ) - # Always restore training mode because log-prob inference may have - # switched the model to inference mode. - with self._timer.time("training_prep"): - await asyncio.to_thread( - self._trainer.prepare_for_training - ) + if self._is_ppo: + with self._timer.time("value_training"): + value_result = await self._value_train_epochs( + train_meta, + num_epochs=self._critic_ppo_epochs, + ) - if has_valid_training_tokens: + if is_policy_training_step: + if ( + self._is_ppo + and self._train_steps == policy_training_start_step + and policy_training_start_step > 0 + ): + print( + f" ✓ Critic warmup complete ({policy_training_start_step} " + "steps). Starting policy training.", + flush=True, + ) + # Always restore training mode because log-prob inference may have + # switched the model to inference mode. Keep it resident + # across every PPO actor epoch. + with self._timer.time("training_prep"): + await asyncio.to_thread(self._trainer.prepare_for_training) + + if has_valid_training_tokens: + for _ in range(self._ppo_epochs): with self._timer.time("policy_training"): if not step_open: await asyncio.to_thread( @@ -1887,12 +1895,6 @@ async def _train_pump(self) -> None: self._trainer.finish_train_step ) step_open = False - if epoch < self._ppo_epochs - 1: - # The next epoch's critic train must not - # share the training GPUs with the policy. - await asyncio.to_thread( - self._trainer.offload_to_cpu - ) if train_meta.sequence_lengths: self._step_log_dict["sequence_lengths"].extend( @@ -3122,19 +3124,25 @@ async def _value_stage(self, meta: KVBatchMeta) -> KVBatchMeta: await asyncio.to_thread(self._value.finish_inference) return meta.with_fields([self._advantage_cfg.values_field]) - async def _value_train(self, meta: KVBatchMeta) -> dict[str, Any]: - """Run one value model optimizer step against this chunk's GAE returns. + async def _value_train_epochs( + self, meta: KVBatchMeta, *, num_epochs: int + ) -> dict[str, Any]: + """Run consecutive critic epochs under one model onload/offload cycle. Returns: - The aggregated value model train result, shaped like Value.train's. + The final epoch's ``train_from_meta`` output; earlier epochs' + results are discarded. """ await asyncio.to_thread(self._value.prepare_for_training) - result = await asyncio.to_thread( - self._value.train_from_meta, - meta, - self._value_loss_fn, # pyrefly: ignore - ) + result: dict[str, Any] | None = None + for _ in range(num_epochs): + result = await asyncio.to_thread( + self._value.train_from_meta, + meta, + self._value_loss_fn, # pyrefly: ignore + ) await asyncio.to_thread(self._value.finish_training) + assert result is not None return result async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index fa2104b688..ac0f118e7c 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -865,21 +865,19 @@ def _validate_algo_settings(master_config: MasterConfig) -> None: "carry TQWorkerMixin, so it has no data-plane setup to call (#2625)." ) - if algo_cfg.ppo_epochs < 1: - raise ValueError("ppo.ppo_epochs must be at least 1") - - # Without it the critic steps once per chunk and the policy once per step, - # which is two effective learning rates from one config, and no error. + # Each PPO epoch must consume the complete RL batch. Without this guard, every + # chunk would independently run the configured actor and critic optimizer steps. if async_config.min_groups_for_streaming_train != algo_cfg.num_prompts_per_step: raise ValueError( "PPO on the SingleController path requires " "async_rl.min_groups_for_streaming_train " f"({async_config.min_groups_for_streaming_train}) == " f"num_prompts_per_step ({algo_cfg.num_prompts_per_step}) so that each RL " - "step is assembled from a single chunk: the critic steps its " - "optimizer once per chunk and the policy once per step. Streaming " - "PPO needs a split train API on the value workers, which they do " - "not have yet (#2625)." + "step is assembled from a single chunk. Otherwise each chunk would " + "run ppo.critic_ppo_epochs critic optimizer steps and ppo.ppo_epochs " + "policy optimizer steps on only part of the RL batch. Streaming PPO " + "needs a split train API on the value workers, which they do not have " + "yet (#2625)." ) failure_config = async_config.rollout_failure @@ -932,8 +930,8 @@ def _validate_algo_settings(master_config: MasterConfig) -> None: raise ValueError( "num_prompts_per_step * num_generations_per_prompt " f"({rl_step_samples}) must equal value.train_global_batch_size " - f"({value_global_batch_size}) so that one RL step maps to exactly one " - "critic optimizer.step." + f"({value_global_batch_size}) so that each critic epoch consumes one " + "complete RL batch." ) diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 4f9080a242..73a7254969 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -712,23 +712,33 @@ def _clamp_max_num_steps( def _maybe_inject_megatron_train_iters(master_config: MasterConfig) -> None: """Set train_iters from max_num_steps after its dataloader clamp.""" algo_cfg = algo_config(master_config) - is_ppo = is_ppo_run(master_config) - # train_iters is a scheduler-tick budget, and each PPO epoch steps both - # optimizers once, so the configured warmup/decay horizon has to be scaled. - ppo_epochs = algo_cfg.ppo_epochs if is_ppo else 1 - train_iters = algo_cfg.max_num_steps * ppo_epochs + ppo_config = master_config.ppo if is_ppo_run(master_config) else None + # train_iters is a scheduler-tick budget. Policy and value need separate + # budgets when their epoch counts or training start steps differ. + policy_epochs = ppo_config.ppo_epochs if ppo_config is not None else 1 + policy_training_steps = algo_cfg.max_num_steps + if ppo_config is not None: + policy_training_steps = max( + policy_training_steps - ppo_config.policy_training_start_step, + 0, + ) + # Megatron-Bridge requires a positive scheduler horizon at setup. A PPO + # policy scheduler is never advanced when critic warmup spans the whole run. + policy_train_iters = max(policy_training_steps * policy_epochs, 1) # policy policy_config = master_config.policy if policy_config.get("megatron_cfg", {}).get("enabled", False): - policy_config["megatron_cfg"]["train_iters"] = train_iters + policy_config["megatron_cfg"]["train_iters"] = policy_train_iters # value - if not is_ppo: + if ppo_config is None: return value_config = master_config.value if value_config.get("megatron_cfg", {}).get("enabled", False): - value_config["megatron_cfg"]["train_iters"] = train_iters # type: ignore[index] + value_config["megatron_cfg"]["train_iters"] = ( # type: ignore[index] + algo_cfg.max_num_steps * ppo_config.critic_ppo_epochs + ) def _maybe_attach_fleet_health( diff --git a/tests/unit/algorithms/test_ppo.py b/tests/unit/algorithms/test_ppo.py index 2674d85f5e..926a4d73ac 100644 --- a/tests/unit/algorithms/test_ppo.py +++ b/tests/unit/algorithms/test_ppo.py @@ -645,6 +645,40 @@ def test_ppo_schema_rejects_unsupported_estimator_name(): PPOConfig(adv_estimator={"name": "grpo"}) +def test_ppo_config_defaults_both_epoch_counts_to_four(): + config = PPOConfig() + + assert config.ppo_epochs == 4 + assert config.critic_ppo_epochs == 4 + + +def test_ppo_config_accepts_more_critic_epochs(): + config = PPOConfig(ppo_epochs=1, critic_ppo_epochs=3) + + assert config.critic_ppo_epochs == 3 + + +def test_ppo_config_accepts_independent_actor_and_critic_epoch_counts(): + config = PPOConfig(ppo_epochs=3, critic_ppo_epochs=1) + + assert config.ppo_epochs == 3 + assert config.critic_ppo_epochs == 1 + + +def test_ppo_config_rejects_zero_ppo_epochs(): + from pydantic import ValidationError + + with pytest.raises(ValidationError, match="ppo_epochs must be at least 1"): + PPOConfig(ppo_epochs=0) + + +def test_ppo_config_rejects_zero_critic_ppo_epochs(): + from pydantic import ValidationError + + with pytest.raises(ValidationError, match="critic_ppo_epochs must be at least 1"): + PPOConfig(critic_ppo_epochs=0) + + def test_create_advantage_estimator_rejects_unsupported_name(): """The factory still guards names that skipped schema validation. @@ -726,6 +760,7 @@ def _run_mock_ppo_train( max_num_steps: int, ppo_epochs: int, seq_logprob_error_threshold: float | None, + critic_ppo_epochs: int | None = None, policy_training_start_step: int = 0, warmup_generation_lead_steps: int | None = None, overlong_filtering: bool = False, @@ -847,6 +882,9 @@ def __len__(self): "value_inference_finish" ) value_model.finish_training.side_effect = lambda: events.append("value_finish") + value_model.prepare_for_training.side_effect = lambda: events.append( + "value_train_prep" + ) value_model.get_values.return_value = {"values": torch.zeros(2, 3, 1)} value_model.train.side_effect = lambda *_args, **_kwargs: ( events.append("value_train") or value_result @@ -908,6 +946,9 @@ def fake_rollout(*_args, input_batch, **_kwargs): overlong_filtering=overlong_filtering, policy_training_start_step=policy_training_start_step, ppo_epochs=ppo_epochs, + critic_ppo_epochs=( + ppo_epochs if critic_ppo_epochs is None else critic_ppo_epochs + ), reward_scaling={"enabled": False}, reward_shaping=RewardShapingConfig(enabled=False), seq_logprob_error_threshold=seq_logprob_error_threshold, @@ -1045,7 +1086,11 @@ def test_ppo_train_noncolocated_refit_offload_lifecycle(monkeypatch): assert harness.refit.call_count == 2 assert harness.policy.train.call_count == 4 assert harness.value_model.train.call_count == 4 - assert harness.policy.offload_to_cpu.call_count == 4 + assert harness.value_model.prepare_for_training.call_count == 2 + assert harness.policy.prepare_for_training.call_count == 2 + # One offload after each rollout step; there is no longer an extra policy + # offload between PPO epochs. + assert harness.policy.offload_to_cpu.call_count == 2 harness.policy_generation.prepare_for_generation.assert_not_called() assert harness.policy_generation.finish_generation.call_count == 2 @@ -1067,6 +1112,68 @@ def test_ppo_train_noncolocated_refit_offload_lifecycle(monkeypatch): ] +@pytest.mark.parametrize("async_mode", [False, True]) +def test_ppo_train_runs_extra_critic_epochs_without_extra_actor_updates( + monkeypatch, async_mode +): + harness = _run_mock_ppo_train( + monkeypatch, + async_mode=async_mode, + max_num_steps=1, + ppo_epochs=2, + critic_ppo_epochs=3, + seq_logprob_error_threshold=None, + ) + + assert harness.value_model.train.call_count == 3 + assert harness.policy.train.call_count == 2 + assert harness.value_model.prepare_for_training.call_count == 1 + assert harness.policy.prepare_for_training.call_count == 1 + value_phase = harness.events[ + harness.events.index("value_train_prep") : harness.events.index( + "policy_train_prep" + ) + ] + assert value_phase == [ + "value_train_prep", + "value_train", + "value_train", + "value_train", + "value_finish", + ] + policy_prep_index = harness.events.index("policy_train_prep") + policy_train_indices = [ + index for index, event in enumerate(harness.events) if event == "policy_train" + ] + assert harness.events[policy_prep_index : policy_train_indices[-1] + 1] == [ + "policy_train_prep", + "policy_train", + "policy_train", + ] + + +@pytest.mark.parametrize("async_mode", [False, True]) +def test_ppo_train_critic_keeps_extra_epochs_during_policy_warmup( + monkeypatch, async_mode +): + harness = _run_mock_ppo_train( + monkeypatch, + async_mode=async_mode, + max_num_steps=2, + ppo_epochs=2, + critic_ppo_epochs=3, + seq_logprob_error_threshold=None, + policy_training_start_step=1, + ) + + # Step 0 is critic-only warmup; step 1 trains both. The critic always + # runs all 3 epochs, independent of the policy warmup gate. + assert harness.value_model.train.call_count == 6 + assert harness.policy.train.call_count == 2 + assert harness.value_model.prepare_for_training.call_count == 2 + assert harness.policy.prepare_for_training.call_count == 1 + + @pytest.mark.parametrize("async_mode", [False, True]) def test_ppo_train_critic_warmup_reuses_generation_until_policy_update( monkeypatch, async_mode @@ -1947,13 +2054,30 @@ def test_noncolocated_vllm_builds_separate_clusters_and_collective(monkeypatch): @pytest.mark.parametrize( - ("async_enabled", "expected_train_iters"), - [(False, 3), (True, 30)], + ( + "async_enabled", + "critic_ppo_epochs", + "policy_training_start_step", + "expected_policy_train_iters", + "expected_value_train_iters", + ), + [ + (False, 3, 0, 3, 3), + (False, 5, 0, 3, 5), + (True, 3, 0, 30, 30), + (True, 5, 2, 24, 50), + (True, 5, 10, 1, 50), + ], ) def test_megatron_train_iters_matches_ppo_training_limit( - monkeypatch, async_enabled, expected_train_iters + monkeypatch, + async_enabled, + critic_ppo_epochs, + policy_training_start_step, + expected_policy_train_iters, + expected_value_train_iters, ): - """Async PPO cycles data until max_num_steps; sync PPO also honors epochs.""" + """Each model's scheduler budget follows its own number of epochs.""" from nemo_rl.algorithms.ppo import AsyncPPOConfig config = _make_noncolocated_setup_config() @@ -1962,12 +2086,14 @@ def test_megatron_train_iters_matches_ppo_training_limit( config.ppo.max_num_steps = 10 config.ppo.max_num_epochs = -1 if async_enabled else 1 config.ppo.ppo_epochs = 3 + config.ppo.critic_ppo_epochs = critic_ppo_epochs + config.ppo.policy_training_start_step = policy_training_start_step config.ppo.async_ppo = AsyncPPOConfig(enabled=async_enabled) _run_noncolocated_setup(monkeypatch, config) - assert config.policy["megatron_cfg"]["train_iters"] == expected_train_iters - assert config.value["megatron_cfg"]["train_iters"] == expected_train_iters + assert config.policy["megatron_cfg"]["train_iters"] == expected_policy_train_iters + assert config.value["megatron_cfg"]["train_iters"] == expected_value_train_iters def test_ppo_setup_rejects_a_warm_start_that_does_not_resolve(monkeypatch, tmp_path): @@ -2196,23 +2322,11 @@ def test_async_ppo_launcher_entry_guards(mutate, message): _validate_async_ppo_entry_config(config) -@pytest.mark.parametrize( - ("mutate", "message"), - [ - (lambda cfg: setattr(cfg.ppo, "ppo_epochs", 0), "ppo_epochs"), - ( - lambda cfg: ( - setattr(cfg.ppo, "skip_reference_policy_logprobs_calculation", True), - setattr(cfg.loss_fn, "reference_policy_kl_penalty", 0.1), - ), - "Skipping reference logprobs", - ), - ], -) -def test_async_ppo_training_loop_guards(mutate, message): +def test_async_ppo_training_loop_rejects_skipped_reference_logprobs_with_kl_penalty(): config = _make_async_ppo_config() - mutate(config) - with pytest.raises(ValueError, match=message): + config.ppo.skip_reference_policy_logprobs_calculation = True + config.loss_fn.reference_policy_kl_penalty = 0.1 + with pytest.raises(ValueError, match="Skipping reference logprobs"): _call_async_ppo_until_guard(config) diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml index 9887b85d8b..62c9d47080 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml @@ -8,7 +8,9 @@ ppo: max_rollout_turns: 1 max_num_epochs: 100000 max_num_steps: 100000 - ppo_epochs: 4 + ppo_epochs: 4 # actor passes over each rollout batch + # Critic passes over each rollout batch; follows ppo_epochs unless overridden. + critic_ppo_epochs: ${ppo.ppo_epochs} 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 diff --git a/tests/unit/single_controller/test_ppo_setup.py b/tests/unit/single_controller/test_ppo_setup.py index 6e1847c098..4ff9a1fd4d 100644 --- a/tests/unit/single_controller/test_ppo_setup.py +++ b/tests/unit/single_controller/test_ppo_setup.py @@ -263,8 +263,7 @@ def test_rejects_value_block_without_ppo_block(self): validate_single_controller_config(mc) def test_rejects_multi_chunk_streaming(self): - """The critic has no split API, so it would step once per chunk while - the policy steps once per RL step.""" + """PPO cannot spread one full-batch optimizer epoch across chunks.""" mc = _ppo_master_config(min_groups_for_streaming_train=1) with pytest.raises( @@ -283,16 +282,6 @@ def test_rejects_value_global_batch_size_mismatch(self): ): validate_single_controller_config(mc) - def test_rejects_a_ppo_epoch_count_below_one(self): - mc = _ppo_master_config( - ppo=PPOConfig.model_construct( - max_num_steps=100, ppo_epochs=0, **_STEP_CONFIG - ) - ) - - with pytest.raises(ValueError, match="ppo_epochs must be at least 1"): - validate_single_controller_config(mc) - @pytest.mark.parametrize( "sampler_config", [WindowedSamplerConfig(), ReadyFirstSamplerConfig(), WeightFifoSamplerConfig()], @@ -471,20 +460,59 @@ def test_grpo_delegates_to_the_group_relative_factory(self): class TestMegatronTrainIters: - @pytest.mark.parametrize(("ppo_epochs", "expected"), [(1, 7), (3, 21)]) - def test_injects_into_both_policy_and_value(self, ppo_epochs, expected): - """Each epoch steps both optimizers, so each is a scheduler tick.""" + @pytest.mark.parametrize( + ( + "ppo_epochs", + "policy_training_start_step", + "expected_policy", + "expected_value", + ), + [ + (1, 0, 7, 7), + (3, 0, 21, 21), + (3, 2, 15, 21), + (3, 7, 1, 21), + ], + ) + def test_injects_into_both_policy_and_value( + self, + ppo_epochs, + policy_training_start_step, + expected_policy, + expected_value, + ): + """Each scheduler budget matches its model's optimizer update count.""" + mc = _ppo_master_config( + megatron_enabled=True, + ppo=PPOConfig.model_construct( + max_num_steps=7, + ppo_epochs=ppo_epochs, + critic_ppo_epochs=ppo_epochs, + policy_training_start_step=policy_training_start_step, + **_STEP_CONFIG, + ), + ) + + sc_setup_mod._maybe_inject_megatron_train_iters(mc) + + assert mc.policy["megatron_cfg"]["train_iters"] == expected_policy + assert mc.value["megatron_cfg"]["train_iters"] == expected_value + + def test_injects_distinct_policy_and_value_budgets(self): mc = _ppo_master_config( megatron_enabled=True, ppo=PPOConfig.model_construct( - max_num_steps=7, ppo_epochs=ppo_epochs, **_STEP_CONFIG + max_num_steps=7, + ppo_epochs=1, + critic_ppo_epochs=3, + **_STEP_CONFIG, ), ) sc_setup_mod._maybe_inject_megatron_train_iters(mc) - assert mc.policy["megatron_cfg"]["train_iters"] == expected - assert mc.value["megatron_cfg"]["train_iters"] == expected + assert mc.policy["megatron_cfg"]["train_iters"] == 7 + assert mc.value["megatron_cfg"]["train_iters"] == 21 def test_skips_a_critic_on_a_non_megatron_backend(self): mc = _ppo_master_config(megatron_enabled=False, max_num_steps=7) diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index d6ec4c7b5d..f3b3dc4a16 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -1109,6 +1109,7 @@ def _train_pump_controller(*, sampler) -> object: ctrl._trainer = _NoOpTrainer() ctrl._is_ppo = False ctrl._ppo_epochs = 1 + ctrl._critic_ppo_epochs = 1 ctrl._value = None ctrl._value_loss_fn = None ctrl._gen = SimpleNamespace(requires_kv_scale_sync=False) @@ -1484,11 +1485,15 @@ def _ppo_train_pump_controller( policy_training_start_step: int = 0, value: _NoOpValue | None = None, ppo_epochs: int = 1, + critic_ppo_epochs: int | None = None, ) -> tuple[object, _NoOpValue]: ctrl = _train_pump_controller(sampler=sampler) value = _NoOpValue() if value is None else value ctrl._is_ppo = True ctrl._ppo_epochs = ppo_epochs + ctrl._critic_ppo_epochs = ( + ppo_epochs if critic_ppo_epochs is None else critic_ppo_epochs + ) ctrl._value = value ctrl._value_loss_fn = MagicMock(name="value_loss_fn") ctrl._master_config.grpo = None @@ -1631,11 +1636,13 @@ def test_train_pump_freezes_the_policy_during_critic_warmup( """Below policy_training_start_step the critic trains alone: no optimizer step, and no weight transfer to generation either. The frozen policy does not shorten the critic's own epoch loop.""" + critic_ppo_epochs = 3 meta = _single_group_meta() ctrl, value = _ppo_train_pump_controller( sampler=_OneThenEmptySampler(meta), policy_training_start_step=1, ppo_epochs=ppo_epochs, + critic_ppo_epochs=critic_ppo_epochs, ) trainer = MagicMock(spec=_NoOpTrainer) ctrl._trainer = trainer @@ -1644,7 +1651,7 @@ def test_train_pump_freezes_the_policy_during_critic_warmup( asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0)) - assert value.calls.count("train_from_meta") == ppo_epochs + assert value.calls.count("train_from_meta") == critic_ppo_epochs trainer.prepare_for_training.assert_not_called() trainer.begin_train_step.assert_not_called() trainer.finish_train_step.assert_not_called() @@ -1681,11 +1688,11 @@ def test_train_pump_trains_the_policy_once_warmup_is_over(monkeypatch, capsys) - assert capsys.readouterr().out.count("Critic warmup complete") == 1 -def test_train_pump_offloads_the_policy_between_ppo_epochs(monkeypatch) -> None: - """ppo_epochs repeats the whole train stage over the step's own batch. +def test_train_pump_groups_ppo_epochs_by_model(monkeypatch) -> None: + """Each model stays resident for all of its PPO epochs. - The two models share the training GPUs, so every critic train runs with the - policy on CPU -- including the ones after the first epoch.""" + The critic still finishes and leaves the shared training GPUs before the + policy is loaded, but the models no longer move between epochs.""" meta = _single_group_meta() calls: list[str] = [] ctrl, _ = _ppo_train_pump_controller( @@ -1708,24 +1715,55 @@ def test_train_pump_offloads_the_policy_between_ppo_epochs(monkeypatch) -> None: "critic.finish_inference", "critic.prepare_for_training", "critic.train_from_meta", + "critic.train_from_meta", "critic.finish_training", "policy.prepare_for_training", "policy.begin_train_step", "policy.train_microbatches_from_meta", "policy.finish_train_step", + "policy.begin_train_step", + "policy.train_microbatches_from_meta", + "policy.finish_train_step", + ] + # Still one RL step, so one refit and one version bump. + ctrl._sync_weights.assert_awaited_once_with(calibration_data=None) + assert ctrl._trainer_version == 1 + + +def test_train_pump_runs_all_critic_epochs_before_actor_epochs(monkeypatch) -> None: + """Independent critic epochs share one residency and do not update policy.""" + meta = _single_group_meta() + calls: list[str] = [] + ctrl, _ = _ppo_train_pump_controller( + sampler=_OneThenEmptySampler(meta), + value=_NoOpValue(calls=calls, prefix="critic."), + ppo_epochs=1, + critic_ppo_epochs=3, + ) + ctrl._trainer = _EpochRecordingTrainer(calls) + ctrl._advantage_stage = AsyncMock(return_value=(meta, True)) + monkeypatch.setattr(single_controller.ray, "cluster_resources", lambda: {}) + + asyncio.run(asyncio.wait_for(ctrl._train_pump(), timeout=1.0)) + + assert calls == [ + # Neither logprob is required here, so the policy is parked up front. "policy.offload_to_cpu", + "policy.finish_inference", + "critic.prepare_for_inference", + "critic.get_values_from_meta", + "critic.finish_inference", "critic.prepare_for_training", "critic.train_from_meta", + "critic.train_from_meta", + "critic.train_from_meta", "critic.finish_training", "policy.prepare_for_training", "policy.begin_train_step", "policy.train_microbatches_from_meta", - # No offload after the last epoch: the refit needs the policy resident. "policy.finish_train_step", ] - # Still one RL step, so one refit and one version bump. ctrl._sync_weights.assert_awaited_once_with(calibration_data=None) - assert ctrl._trainer_version == 1 def test_advantage_stage_writes_gae_returns_alongside_advantages() -> None: diff --git a/tests/unit/test_config_v2.py b/tests/unit/test_config_v2.py index 335f33db32..f03a7a034f 100644 --- a/tests/unit/test_config_v2.py +++ b/tests/unit/test_config_v2.py @@ -38,7 +38,11 @@ from nemo_rl.algorithms.rm import MasterConfig as RMMasterConfig from nemo_rl.algorithms.sft import MasterConfig as SFTMasterConfig from nemo_rl.evals.eval import MasterConfig as EvalMasterConfig -from nemo_rl.utils.config import load_config, register_omegaconf_resolvers +from nemo_rl.utils.config import ( + load_config, + parse_hydra_overrides, + register_omegaconf_resolvers, +) # All tests in this module should run first pytestmark = pytest.mark.run_first @@ -54,6 +58,15 @@ ) +def test_ppo_recipe_critic_epochs_follow_actor_override(): + config = load_config(real_configs_dir / "ppo_math_1B.yaml") + config = parse_hydra_overrides(config, ["ppo.ppo_epochs=7"]) + + resolved = OmegaConf.to_container(config, resolve=True) + assert resolved["ppo"]["ppo_epochs"] == 7 + assert resolved["ppo"]["critic_ppo_epochs"] == 7 + + def _collect_mismatched_keys(real: dict, reference: dict, path: str = ""): """Return keys present in real but absent in reference, and keys with differing values.