diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index 97e536cfb02..7a521a13989 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -262,6 +262,8 @@ def validate_sync( return {}, {} timer = Timer() + # >= 1 is validated in setup(). + val_num_generations_per_prompt = master_config.grpo.val_num_generations_per_prompt total_rewards: list[float] = [] total_lengths: list[float] = [] all_message_logs: list[list[dict[str, str]]] = [] @@ -276,8 +278,10 @@ def validate_sync( for batch_idx, val_batch in enumerate(val_dataloader): if batch_idx >= max_batches: break - n_prompts = int(val_batch.size) - policy.prepare_val_partition(n_prompts, partition_id=partition_id) + if val_num_generations_per_prompt > 1: + val_batch = val_batch.repeat_interleave(val_num_generations_per_prompt) + n_rollouts = int(val_batch.size) + policy.prepare_val_partition(n_rollouts, partition_id=partition_id) meta, driver_carry, rollout_metrics, _ = ray.get( rollout_actor.rollout_to_tq.remote( val_batch, @@ -294,7 +298,7 @@ def validate_sync( total_lengths.append(rollout_metrics["mean_gen_tokens_per_sample"]) all_message_logs.extend( [{"role": r, "content": c} for r, c in zip(roles[i], contents[i])] - for i in range(n_prompts) + for i in range(n_rollouts) ) if capture_extras: additional_metrics = rollout_metrics @@ -305,12 +309,35 @@ def validate_sync( if total_rewards else 0.0 ) + # Grouped validation (val_num_generations_per_prompt > 1) additionally + # reports pass@k over each prompt's k rollouts, mirroring + # nemo_rl.algorithms.grpo.validate. + pass_k = None + if total_rewards and val_num_generations_per_prompt > 1: + assert len(total_rewards) % val_num_generations_per_prompt == 0, ( + "Validation rewards must be divisible by " + "grpo.val_num_generations_per_prompt" + ) + pass_k = ( + ( + torch.tensor(total_rewards, dtype=torch.float32).view( + -1, val_num_generations_per_prompt + ) + > 0 + ) + .any(dim=1) + .float() + .mean() + .item() + ) avg_length = sum(total_lengths) / len(total_lengths) if total_lengths else 0.0 val_metrics = { "accuracy": accuracy, "avg_length": avg_length, **additional_metrics, } + if pass_k is not None: + val_metrics["pass_k"] = pass_k try: print_message_log_samples( all_message_logs, diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 346353fcd45..8e314b05bd6 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -56,7 +56,11 @@ setup, validate, ) -from nemo_rl.algorithms.grpo_sync import _train_fields_for_step, grpo_train_sync +from nemo_rl.algorithms.grpo_sync import ( + _train_fields_for_step, + grpo_train_sync, + validate_sync, +) from nemo_rl.algorithms.loss import ClippedPGLossConfig, ClippedPGLossFn from nemo_rl.algorithms.reward_functions import ( RewardShapingConfig, @@ -4037,6 +4041,69 @@ def run_rollout(_policy, repeated_batch, *_args, **_kwargs): assert val_metrics["accuracy"] == pytest.approx(0.125) assert val_metrics["pass_k"] == pytest.approx(0.5) + def test_sync_grouped_validation_reports_pass_k(self, mock_grpo_components): + mock_batch = BatchedDataDict[DatumSpec]( + { + "message_log": [ + [{"role": "user", "content": "a", "token_ids": torch.tensor([1])}], + [{"role": "user", "content": "b", "token_ids": torch.tensor([2])}], + ], + "task_name": ["math", "math"], + "extra_env_info": [{}, {}], + "loss_multiplier": torch.tensor([1.0, 1.0]), + "idx": torch.tensor([0, 1]), + "length": torch.tensor([1, 1]), + "total_reward": torch.tensor([0.0, 0.0]), + } + ) + mock_dataloader = MagicMock(spec=StatefulDataLoader) + mock_dataloader.__iter__ = MagicMock(return_value=iter([mock_batch])) + mock_config = mock_grpo_components["master_config"] + mock_config.grpo.max_val_samples = 2 + mock_config.grpo.val_batch_size = 2 + mock_config.grpo.val_num_generations_per_prompt = 4 + + def rollout_to_tq_remote(repeated_batch, **_kwargs): + # Each prompt is repeated k=4 times, contiguously. + assert repeated_batch["idx"].tolist() == [0, 0, 0, 0, 1, 1, 1, 1] + driver_carry = { + # Prompt 0 passes once out of 4; prompt 1 never passes. + "total_reward": torch.tensor( + [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + ), + "turn_roles": [["user"]] * 8, + "turn_contents": [["x"]] * 8, + } + return (MagicMock(), driver_carry, {"mean_gen_tokens_per_sample": 1.0}, {}) + + rollout_actor = MagicMock() + rollout_actor.rollout_to_tq.remote.side_effect = rollout_to_tq_remote + policy = MagicMock() + + with ( + patch("nemo_rl.algorithms.grpo_sync.ray.get", side_effect=lambda x: x), + patch( + "nemo_rl.algorithms.grpo_sync._should_use_nemo_gym", + return_value=False, + ), + patch("nemo_rl.algorithms.grpo_sync.print_message_log_samples"), + ): + val_metrics, _ = validate_sync( + rollout_actor=rollout_actor, + policy=policy, + val_dataloader=mock_dataloader, + val_task_to_env={"math": MagicMock(spec=EnvironmentInterface)}, + step=0, + master_config=mock_config, + ) + + # The val partition is sized for the expanded batch (2 prompts x k=4). + policy.prepare_val_partition.assert_called_once_with(8, partition_id="val") + # accuracy stays the plain mean over all 8 rollouts; pass@4 counts + # prompts with at least one passing rollout (1 of 2). + assert val_metrics["accuracy"] == pytest.approx(0.125) + assert val_metrics["pass_k"] == pytest.approx(0.5) + def test_validation_uses_val_sampling_params_on_gym_path( self, mock_grpo_components ):