diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 8c45f089b9b..49b4154c16d 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -8,6 +8,7 @@ grpo: normalize_rewards: true use_leave_one_out_baseline: true val_period: 10 + val_start_at: -1 # First step eligible for periodic validation; -1 disables the delay val_at_start: false val_at_end: false overlong_filtering: false diff --git a/examples/nemo_gym/grpo_nanov3.yaml b/examples/nemo_gym/grpo_nanov3.yaml index a72f6844c63..55ecacc1b3d 100644 --- a/examples/nemo_gym/grpo_nanov3.yaml +++ b/examples/nemo_gym/grpo_nanov3.yaml @@ -9,6 +9,7 @@ grpo: normalize_rewards: true use_leave_one_out_baseline: true val_period: 5 + val_start_at: -1 val_at_start: False val_at_end: False overlong_filtering: true diff --git a/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml b/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml index a845056e84e..c7fe9c903ce 100644 --- a/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml +++ b/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml @@ -7,6 +7,7 @@ grpo: normalize_rewards: true use_leave_one_out_baseline: true val_period: 10 + val_start_at: -1 val_at_start: true val_at_end: false overlong_filtering: false diff --git a/examples/nemo_gym/nemotron-3-super/stage1_rlvr.yaml b/examples/nemo_gym/nemotron-3-super/stage1_rlvr.yaml index a95d1053f81..a98b0abbb42 100644 --- a/examples/nemo_gym/nemotron-3-super/stage1_rlvr.yaml +++ b/examples/nemo_gym/nemotron-3-super/stage1_rlvr.yaml @@ -22,6 +22,7 @@ grpo: advantage_clip_low: -50 advantage_clip_high: 50 val_period: -1 + val_start_at: -1 val_at_start: false val_at_end: false overlong_filtering: false diff --git a/examples/nemo_gym/nemotron-3-super/stage2_swe1.yaml b/examples/nemo_gym/nemotron-3-super/stage2_swe1.yaml index 5338e66b362..9c7f3cdf68c 100644 --- a/examples/nemo_gym/nemotron-3-super/stage2_swe1.yaml +++ b/examples/nemo_gym/nemotron-3-super/stage2_swe1.yaml @@ -22,6 +22,7 @@ grpo: advantage_clip_low: -100 advantage_clip_high: 100 val_period: 10 + val_start_at: -1 val_at_start: false val_at_end: false overlong_filtering: true diff --git a/examples/nemo_gym/nemotron-3-super/stage2_swe2.yaml b/examples/nemo_gym/nemotron-3-super/stage2_swe2.yaml index 47cefdb5266..d14b0d6685c 100644 --- a/examples/nemo_gym/nemotron-3-super/stage2_swe2.yaml +++ b/examples/nemo_gym/nemotron-3-super/stage2_swe2.yaml @@ -22,6 +22,7 @@ grpo: advantage_clip_low: -100 advantage_clip_high: 100 val_period: 100 + val_start_at: -1 val_at_start: false val_at_end: false overlong_filtering: true diff --git a/examples/nemo_gym/nemotron-3-super/stage3_rlhf.yaml b/examples/nemo_gym/nemotron-3-super/stage3_rlhf.yaml index 7d251f8958b..501cccd82a1 100644 --- a/examples/nemo_gym/nemotron-3-super/stage3_rlhf.yaml +++ b/examples/nemo_gym/nemotron-3-super/stage3_rlhf.yaml @@ -22,6 +22,7 @@ grpo: advantage_clip_low: -50 advantage_clip_high: 50 val_period: 10000 + val_start_at: -1 val_at_start: false val_at_end: false overlong_filtering: false diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 973f6175284..d0e59f1891f 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -253,6 +253,8 @@ class GRPOConfig(TypedDict): advantage_clip_high: NotRequired[float | None] use_leave_one_out_baseline: bool val_period: int + # First training step eligible for periodic validation; -1 disables the delay. + val_start_at: int val_batch_size: int | None # None for NeMo-Gym compatibility val_at_start: bool # Whether to run validation on the last training step. Setting this to True ensures the @@ -2586,6 +2588,7 @@ def grpo_train( val_at_start = master_config.grpo["val_at_start"] val_at_end = master_config.grpo["val_at_end"] val_period = master_config.grpo["val_period"] + val_start_at = master_config.grpo["val_start_at"] colocated_inference = master_config.policy["generation"]["colocated"]["enabled"] refit_buffer_size_gb = master_config.policy.get("refit_buffer_size_gb") @@ -3144,9 +3147,11 @@ def grpo_train( ) # Run validation if it's a validation step or last step with val_at_end - if (val_period > 0 and (total_steps + 1) % val_period == 0) or ( - val_at_end and is_last_step - ): + if ( + val_period > 0 + and (total_steps + 1) >= val_start_at + and (total_steps + 1) % val_period == 0 + ) or (val_at_end and is_last_step): memory_tracker.snapshot_start_of_stage("Validation", dir()) if NEED_REFIT and POLICY_GENERATION_STALE: refit_metrics = refit_policy_generation( @@ -3850,6 +3855,7 @@ def async_grpo_train( "total_valid_tokens", 0 ) # Default to 0 for backward compatibility with older checkpoints val_period = master_config.grpo["val_period"] + val_start_at = master_config.grpo["val_start_at"] val_at_start = master_config.grpo["val_at_start"] val_at_end = master_config.grpo["val_at_end"] colocated_inference = master_config.policy["generation"]["colocated"]["enabled"] @@ -4547,9 +4553,11 @@ def async_grpo_train( is_last_step = step + 1 == master_config.grpo["max_num_steps"] # Run validation if it's a validation step or last step with val_at_end - if (val_period > 0 and (step + 1) % val_period == 0) or ( - val_at_end and is_last_step - ): + if ( + val_period > 0 + and (step + 1) >= val_start_at + and (step + 1) % val_period == 0 + ) or (val_at_end and is_last_step): with timer.time("idle/validation"): # Pause trajectory collection during validation to reduce memory pressure trajectory_collector.pause.remote() diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index 03089443b60..e076bccbc6d 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -443,6 +443,7 @@ def grpo_train_sync( val_at_start = master_config.grpo["val_at_start"] val_at_end = master_config.grpo["val_at_end"] val_period = master_config.grpo["val_period"] + val_start_at = master_config.grpo["val_start_at"] colocated_inference = master_config.policy["generation"]["colocated"]["enabled"] # ── Data-plane setup (mandatory in the sync trainer) ─────────────── @@ -990,9 +991,11 @@ def grpo_train_sync( and (current_step + 1 == len(wrapped_dataloader)) ) - if (val_period > 0 and (total_steps + 1) % val_period == 0) or ( - val_at_end and is_last_step - ): + if ( + val_period > 0 + and (total_steps + 1) >= val_start_at + and (total_steps + 1) % val_period == 0 + ) or (val_at_end and is_last_step): memory_tracker.snapshot_start_of_stage("Validation", dir()) if NEED_REFIT and POLICY_GENERATION_STALE: refit_policy_generation( diff --git a/research/template_project/configs/grpo_math_1B.yaml b/research/template_project/configs/grpo_math_1B.yaml index b172576f95d..c6ab5191a7c 100644 --- a/research/template_project/configs/grpo_math_1B.yaml +++ b/research/template_project/configs/grpo_math_1B.yaml @@ -8,6 +8,7 @@ grpo: normalize_rewards: true use_leave_one_out_baseline: true val_period: 10 + val_start_at: -1 # First step eligible for periodic validation; -1 disables the delay val_at_start: false val_at_end: false overlong_filtering: false diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 7b8b2c703fd..0f983715bda 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from contextlib import contextmanager +from contextlib import ExitStack, contextmanager from typing import Any from unittest.mock import MagicMock, patch @@ -47,7 +47,7 @@ refit_policy_generation, validate, ) -from nemo_rl.algorithms.grpo_sync import _train_fields_for_step +from nemo_rl.algorithms.grpo_sync import _train_fields_for_step, grpo_train_sync from nemo_rl.algorithms.loss import ClippedPGLossConfig, ClippedPGLossFn from nemo_rl.algorithms.reward_functions import ( RewardShapingConfig, @@ -263,6 +263,7 @@ def val_iter(self): "num_generations_per_prompt": 1, "max_rollout_turns": 1, "val_period": 100, + "val_start_at": -1, "val_batch_size": 1, "val_at_start": False, "val_at_end": False, @@ -1000,6 +1001,93 @@ def mock_ray_get(ref): return stack +def mock_sync_grpo_infrastructure(policy): + """Context manager that mocks the TQ/data-plane infrastructure of grpo_train_sync. + + Mirrors ``mock_async_grpo_infrastructure``: the Ray rollout actor and the + TQ round-trips are stubbed so the driver loop runs for real, with small + real tensors standing in for the per-sample slices the driver computes + against. ``validate_sync`` is intentionally left unpatched so tests can + install their own capturing mock. + """ + stack = ExitStack() + + # Slice returned by the stubbed rollout actor; baseline/std are computed + # for real on the driver from these fields. + driver_carry = BatchedDataDict( + { + "total_reward": torch.tensor([1.0]), + "prompt_ids_for_adv": torch.tensor([[1, 2, 3]]), + "input_lengths": torch.tensor([4]), + "loss_multiplier": torch.tensor([1.0]), + "truncated": torch.tensor([False]), + "length": torch.tensor([3]), + } + ) + meta = MagicMock() + meta.fields = ["input_ids"] + rollout_metrics = { + "mean_gen_tokens_per_sample": 10.0, + "max_gen_tokens": 20, + "min_gen_tokens": 5, + } + rollout_actor = MagicMock() + rollout_actor.rollout_to_tq.remote.return_value = ( + meta, + driver_carry, + rollout_metrics, + {}, + ) + rollout_actor_cls = MagicMock() + rollout_actor_cls.options.return_value.remote.return_value = rollout_actor + stack.enter_context( + patch("nemo_rl.algorithms.grpo_sync.SyncRolloutActor", rollout_actor_cls) + ) + stack.enter_context( + patch("nemo_rl.algorithms.grpo_sync.make_actor_runtime_env", return_value={}) + ) + # The only ray.get on the driver path receives the stub actor's plain tuple. + stack.enter_context(patch("ray.get", side_effect=lambda ref: ref)) + + stack.enter_context( + patch("nemo_rl.algorithms.grpo_sync.refit_policy_generation", return_value=None) + ) + stack.enter_context( + patch( + "nemo_rl.algorithms.grpo_sync._compute_seq_logprob_error_metrics", + return_value=(torch.ones(1), _mock_seq_logprob_error_result()), + ) + ) + adv_estimator = MagicMock() + adv_estimator.compute_advantage.return_value = torch.zeros(1, 4) + stack.enter_context( + patch( + "nemo_rl.algorithms.grpo_sync._create_advantage_estimator", + return_value=adv_estimator, + ) + ) + stack.enter_context( + patch("nemo_rl.algorithms.grpo_sync.print_performance_metrics", return_value={}) + ) + + # TQ-mediated policy methods: per-token slices read back from the data + # plane, and train results in the same shape as ``policy.train``. + dp_bank = { + "generation_logprobs": torch.zeros(1, 4), + "token_mask": torch.ones(1, 4), + "prev_logprobs": torch.zeros(1, 4), + "reference_policy_logprobs": torch.zeros(1, 4), + "input_ids": torch.ones(1, 4, dtype=torch.long), + } + policy.read_from_dataplane.side_effect = lambda meta, select_fields, **kw: ( + BatchedDataDict({k: dp_bank[k].clone() for k in select_fields}) + ) + policy.train_from_meta.return_value = policy.train.return_value + policy.tq_partition_id = 0 + + return stack + + @pytest.mark.parametrize( ("generation_config", "expected"), [ @@ -2445,6 +2533,87 @@ def test_grpo_train_skips_prev_logprobs_when_force_on_policy_ratio( ) +@pytest.mark.parametrize("train_func", [grpo_train, async_grpo_train, grpo_train_sync]) +@pytest.mark.parametrize( + ("val_at_end", "expected_validation_steps"), + [(False, [4]), (True, [4, 5])], +) +def test_periodic_validation_starts_at_configured_step( + mock_grpo_components, train_func, val_at_end, expected_validation_steps +): + """All three trainers preserve cadence while honoring the validation lower bound.""" + master_config = mock_grpo_components["master_config"] + master_config.grpo.update( + { + "max_num_steps": 5, + "val_period": 2, + "val_start_at": 3, + "val_at_end": val_at_end, + } + ) + mock_batch = next(iter(mock_grpo_components["train_dataloader"])) + mock_rollout_metrics = { + "mean_gen_tokens_per_sample": 10.0, + "max_gen_tokens": 20, + "min_gen_tokens": 5, + } + + with ExitStack() as stack: + validate_target = "nemo_rl.algorithms.grpo.validate" + if train_func is grpo_train_sync: + master_config.data_plane = {"enabled": True} + stack.enter_context( + mock_sync_grpo_infrastructure(mock_grpo_components["policy"]) + ) + validate_target = "nemo_rl.algorithms.grpo_sync.validate_sync" + elif train_func is async_grpo_train: + master_config.policy["generation"]["colocated"]["enabled"] = False + stack.enter_context( + mock_async_grpo_infrastructure(mock_batch, mock_rollout_metrics) + ) + else: + stack.enter_context( + patch( + "nemo_rl.algorithms.grpo.run_multi_turn_rollout", + return_value=(mock_batch, mock_rollout_metrics), + ) + ) + stack.enter_context( + patch( + "nemo_rl.algorithms.grpo.run_async_multi_turn_rollout", + return_value=(mock_batch, mock_rollout_metrics), + ) + ) + stack.enter_context( + patch( + "nemo_rl.algorithms.grpo.compute_and_apply_seq_logprob_error_masking", + return_value=_mock_seq_logprob_error_result(), + ) + ) + + mock_validate = stack.enter_context( + patch(validate_target, return_value=({}, {})) + ) + train_func( + mock_grpo_components["policy"], + _mock_policy_generation(), + mock_grpo_components["train_dataloader"], + mock_grpo_components["val_dataloader"], + mock_grpo_components["tokenizer"], + mock_grpo_components["loss_fn"], + mock_grpo_components["task_to_env"], + mock_grpo_components["val_task_to_env"], + mock_grpo_components["logger"], + mock_grpo_components["checkpointer"], + _default_grpo_save_state(), + master_config, + ) + + assert [call.kwargs["step"] for call in mock_validate.call_args_list] == ( + expected_validation_steps + ) + + @pytest.mark.parametrize("train_func", [grpo_train, async_grpo_train]) def test_grpo_exit_on_max_steps(mock_grpo_components, train_func): """Test that GRPO training loop exits when max_num_steps is reached""" diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index 2124588dac6..de020edfac0 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -8,6 +8,7 @@ grpo: normalize_rewards: true use_leave_one_out_baseline: true val_period: 10 + val_start_at: -1 val_at_start: false val_at_end: false overlong_filtering: false