diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index 6b483977d6e..d50bac5d280 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -241,6 +241,6 @@ The SC path is still under active development. Feature gaps are tracked in [issu - Generation backend: vLLM and Megatron generation are supported; SGLang and TRT-LLM have not been tested on SC. - Validation is not yet supported (setup raises on `val_period > 0`, `val_at_start`, or `val_at_end`); checkpointing is. - (PPO) Rollout drop budgets — `async_rl.rollout_failure.max_skipped_prompts` and `max_consecutive_dropped_prompts` must both be `0`. A drop shortens the step, and the critic shards it against the configured `value.train_global_batch_size` rather than its actual size, so setup rejects a non-zero budget. The resiliency layer stays available on GRPO. -- Reward shaping — `reward_shaping`, `reward_scaling`, and `use_dynamic_sampling` are implemented on neither algorithm block, so setup rejects them rather than silently skipping the shaping. Environment-flagged sample masking and `overlong_filtering` are supported. +- Reward shaping and sample filtering — `reward_shaping`, `reward_scaling`, and `use_dynamic_sampling` are implemented on neither algorithm block, so setup rejects them rather than silently skipping the shaping. Environment-flagged sample masking and `overlong_filtering` are supported; truncated completions are excluded from the loss through `sample_mask`, and a step in which every completion is filtered is rejected rather than skipped. - The `windowed` sampler has no `over_sampling_ratio` cap — over-produced groups aged past the window are evicted, wasting rollout compute. - The drain gate in refit is not yet supported. diff --git a/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml b/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml index bbddb1beccf..02e7492f6b8 100644 --- a/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml +++ b/examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async-single-controller.yaml @@ -14,9 +14,6 @@ ppo: # SC does not support validation yet. val_period: 0 val_at_start: false - # SC does not implement it, so leaving it enabled would describe filtering - # this run does not do. - overlong_filtering: false # override from the inherited recipe. num_prompts_per_step: 256 max_num_epochs: 1000 diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index d5d5ec58233..bf725e4868d 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -77,7 +77,9 @@ create_sampler, ) from nemo_rl.algorithms.grpo import ( + GRPOConfig, GRPOSaveState, + _clip_grpo_advantages, _write_latest_checkpoint_status, aggregate_rollout_metrics, compute_and_apply_seq_logprob_error_masking, @@ -3377,15 +3379,26 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: response_advantages = torch.masked_select(advantages, mask.bool()) self._step_log_dict["rewards"].append(rewards.detach().cpu()) - self._step_log_dict["masked_advantages"].append( - response_advantages.detach().cpu() - ) if self._teacher_logprobs_required: valid = response_advantages.detach().double() self._opd_stat_sum += float(valid.sum()) self._opd_stat_sumsq += float((valid * valid).sum()) self._opd_stat_count += int(valid.numel()) + # OPD accumulates its statistics from the estimator output above. The + # ordinary advantage metrics and policy training use the clipped values, + # matching the legacy paths. + if not self._is_ppo: + assert isinstance(self._algo_cfg, GRPOConfig) + advantages = _clip_grpo_advantages( + advantages, + self._algo_cfg, + ) + response_advantages = torch.masked_select(advantages, mask.bool()) + self._step_log_dict["masked_advantages"].append( + response_advantages.detach().cpu() + ) + fields_to_put = {adv_cfg.output_field: advantages} if not torch.equal(final_sample_mask, sample_mask): fields_to_put[adv_cfg.sample_mask_field] = final_sample_mask diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 511d242e936..7188285517c 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -776,8 +776,8 @@ def _validate_algo_settings(master_config: MasterConfig) -> None: """Reject algorithm blocks the SingleController path cannot honour. Both directions on the critic: one the PPO path needs and does not have, and - one a GRPO run carries and would never build. Plus the reward shaping and - filtering knobs SC reads on neither path. + one a GRPO run carries and would never build. Plus the reward-shaping and + sampling knobs SC reads on neither path. """ algo_cfg = algo_config(master_config) @@ -791,8 +791,9 @@ def _validate_algo_settings(master_config: MasterConfig) -> None: "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. + # An enabled one here describes shaping this run does not do. An entry leaves + # this list once the SC path implements it; overlong_filtering is applied in + # the advantage stage from the raw completion flags in the TransferQueue. unsupported = [ name for name, enabled in ( diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index 90e36099589..814f1c047c6 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -56,6 +56,9 @@ SingleControllerActorArgs, setup_single_controller, ) +from nemo_rl.algorithms.single_controller_utils.config import ( + validate_single_controller_config, +) from nemo_rl.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION from nemo_rl.data_plane.schema import SC_ROLLOUT_SCHEMA_FIELDS from nemo_rl.experience.rollouts import EffortLevelsConfig @@ -408,6 +411,7 @@ def test_single_controller_mopd_recipe_resolves_to_runtime_contract(): assert isinstance(resolved, dict) config = MasterConfig.model_validate(resolved) + validate_single_controller_config(config) assert config.grpo.async_grpo is None assert config.grpo.adv_estimator.name == "opd" assert config.grpo.skip_reference_policy_logprobs_calculation is True @@ -428,6 +432,24 @@ def test_single_controller_mopd_recipe_resolves_to_runtime_contract(): ) +def test_single_controller_ppo_recipe_inherits_overlong_filtering(): + """The SC nightly exercises the overlong filtering inherited from its parent.""" + register_omegaconf_resolvers() + repo_root = Path(__file__).resolve().parents[3] + recipe = repo_root / ( + "examples/configs/recipes/llm/" + "ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-" + "noncolocated-async-single-controller.yaml" + ) + resolved = OmegaConf.to_container(load_config(recipe), resolve=True) + + assert isinstance(resolved, dict) + config = MasterConfig.model_validate(resolved) + validate_single_controller_config(config) + assert config.ppo is not None + assert config.ppo.overlong_filtering is True + + @pytest.mark.parametrize( ("reference_policy_kl_penalty", "expected_init_reference_model"), [(0.0, False), (0.01, True)], diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index 5f3f41d471e..50e361f5a1d 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -591,7 +591,7 @@ def test_advantage_stage_composes_all_filters_before_computing_advantages( ctrl._teacher_logprobs_required = False ctrl._is_ppo = False ctrl._master_config = SimpleNamespace( - grpo=SimpleNamespace( + grpo=GRPOConfig( seq_logprob_error_threshold=2.0, overlong_filtering=True, invalid_tool_call_advantage=-5.0, @@ -688,7 +688,7 @@ def test_advantage_stage_writes_each_sample_filter_without_seq_threshold( ctrl._teacher_logprobs_required = False ctrl._is_ppo = False ctrl._message_level_advantage_penalties_enabled = False - ctrl._algo_cfg = SimpleNamespace( + ctrl._algo_cfg = GRPOConfig( seq_logprob_error_threshold=None, overlong_filtering=overlong_filtering, ) @@ -751,7 +751,7 @@ def test_advantage_stage_reports_seq_logprob_metrics_without_masking() -> None: ctrl._teacher_logprobs_required = False ctrl._is_ppo = False ctrl._master_config = SimpleNamespace( - grpo=SimpleNamespace(seq_logprob_error_threshold=None, overlong_filtering=False) + grpo=GRPOConfig(seq_logprob_error_threshold=None) ) ctrl._algo_cfg = ctrl._master_config.grpo ctrl._message_level_advantage_penalties_enabled = False @@ -787,6 +787,68 @@ def test_advantage_stage_reports_seq_logprob_metrics_without_masking() -> None: assert metrics[0]["max_seq_mult_prob_error_after_mask"] == pytest.approx(math.e) +def test_advantage_stage_clips_training_values_and_metrics() -> None: + batch_size, sequence_length = 2, 4 + data = TensorDict( + { + "prompt_ids_for_adv": torch.zeros( + batch_size, sequence_length, dtype=torch.long + ), + "total_reward": torch.tensor([-4.0, 6.0]), + "token_mask": torch.ones(batch_size, sequence_length), + "sample_mask": torch.ones(batch_size), + "mask_sample": torch.zeros(batch_size, dtype=torch.bool), + "truncated": torch.zeros(batch_size, dtype=torch.bool), + }, + batch_size=[batch_size], + ) + data_plane = _AdvantageDataPlane(data) + estimator = _MaskRecordingAdvantageEstimator() + + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + ctrl = object.__new__(controller_cls) + ctrl._dp_client = data_plane + ctrl._advantage_cfg = AdvantageConfig() + ctrl._advantage_estimator = estimator + ctrl._policy_logprobs_required = False + ctrl._reference_logprobs_required = False + ctrl._teacher_logprobs_required = False + ctrl._is_ppo = False + ctrl._master_config = SimpleNamespace( + grpo=GRPOConfig( + seq_logprob_error_threshold=None, + advantage_clip_low=-1.0, + advantage_clip_high=2.0, + ) + ) + ctrl._algo_cfg = ctrl._master_config.grpo + ctrl._message_level_advantage_penalties_enabled = False + ctrl._step_log_dict = { + "rewards": [], + "masked_advantages": [], + "num_mask_sample_filtered": [], + "sequence_lengths": [], + "seq_logprob_error_metrics": [], + } + meta = KVBatchMeta( + partition_id="rollout_data", + task_name="train", + sample_ids=[f"sample-{i}" for i in range(batch_size)], + fields=list(data.keys()), + ) + + asyncio.run(ctrl._advantage_stage(meta)) + + assert data_plane.written_fields is not None + torch.testing.assert_close( + data_plane.written_fields["advantages"], + torch.tensor([[-1.0] * sequence_length, [2.0] * sequence_length]), + ) + logged = torch.cat(ctrl._step_log_dict["masked_advantages"]) + assert logged.min().item() == pytest.approx(-1.0) + assert logged.max().item() == pytest.approx(2.0) + + def test_advantage_stage_skips_estimator_when_seq_mask_removes_whole_chunk( capsys: pytest.CaptureFixture[str], ) -> None: @@ -819,7 +881,7 @@ def test_advantage_stage_skips_estimator_when_seq_mask_removes_whole_chunk( ctrl._teacher_logprobs_required = False ctrl._is_ppo = False ctrl._master_config = SimpleNamespace( - grpo=SimpleNamespace(seq_logprob_error_threshold=2.0, overlong_filtering=False) + grpo=GRPOConfig(seq_logprob_error_threshold=2.0) ) ctrl._algo_cfg = ctrl._master_config.grpo ctrl._message_level_advantage_penalties_enabled = False @@ -879,7 +941,7 @@ def test_advantage_stage_skips_preexisting_empty_mask_without_seq_threshold() -> ctrl._teacher_logprobs_required = False ctrl._is_ppo = False ctrl._master_config = SimpleNamespace( - grpo=SimpleNamespace(seq_logprob_error_threshold=None, overlong_filtering=False) + grpo=GRPOConfig(seq_logprob_error_threshold=None) ) ctrl._algo_cfg = ctrl._master_config.grpo ctrl._message_level_advantage_penalties_enabled = False @@ -959,7 +1021,10 @@ def put_samples(self, sample_ids, partition_id, fields): ctrl._is_ppo = False ctrl._dp_client = FakeDataPlane() ctrl._master_config = SimpleNamespace( - grpo=SimpleNamespace(seq_logprob_error_threshold=None, overlong_filtering=False) + grpo=GRPOConfig( + seq_logprob_error_threshold=None, + advantage_clip_high=0.1, + ) ) ctrl._algo_cfg = ctrl._master_config.grpo ctrl._message_level_advantage_penalties_enabled = False @@ -1001,6 +1066,14 @@ def put_samples(self, sample_ids, partition_id, fields): assert ctrl._opd_stat_sum == pytest.approx(1.0) assert ctrl._opd_stat_sumsq == pytest.approx(0.25) assert ctrl._opd_stat_count == 4 + assert ctrl._dp_client.put_fields is not None + written_advantages = ctrl._dp_client.put_fields["advantages"] + torch.testing.assert_close( + written_advantages, + torch.full_like(written_advantages, 0.1), + ) + logged = torch.cat(ctrl._step_log_dict["masked_advantages"]) + torch.testing.assert_close(logged, torch.full((4,), 0.1)) def test_pooled_opd_metrics_weight_unequal_chunks_by_valid_token_count() -> None: