diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index edce5196e93..e4f346d2279 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 and sample filtering — `overlong_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. +- Reward shaping and sample filtering — `overlong_filtering`, `reward_shaping`, and `use_dynamic_sampling` are implemented on neither algorithm block, so setup rejects them rather than silently skipping the shaping. `reward_scaling` is applied, in the advantage stage, before advantages are computed. - 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/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 405c9ec436a..0abc1273336 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -77,8 +77,10 @@ ) from nemo_rl.algorithms.grpo import ( GRPOSaveState, + _clip_grpo_advantages, _write_latest_checkpoint_status, compute_and_apply_seq_logprob_error_masking, + scale_rewards, ) from nemo_rl.algorithms.metric_utils import SetupTimingMetrics from nemo_rl.algorithms.ppo import _compute_critic_metrics @@ -3170,6 +3172,20 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: tensor_field(data, adv_cfg.sample_mask_field) ).float() + # Before anything reads the rewards, matching grpo.py's ordering. A + # no-op unless reward_scaling.enabled, and it is the same helper, so the + # source-range clamp and the warning stay identical. + # ``getattr``: not every algorithm block defines reward_scaling -- the + # OPD test doubles build a bare namespace, and DistillationConfig has + # no reward to scale. A missing block is a disabled one, which is the + # shipped default anyway. + reward_scaling_cfg = getattr(self._algo_cfg, "reward_scaling", None) + if reward_scaling_cfg is not None: + rewards = scale_rewards( + BatchedDataDict({"total_reward": rewards}), + reward_scaling_cfg, + )["total_reward"] + seq_logprob_error_threshold = self._algo_cfg.seq_logprob_error_threshold # Match the legacy path: whenever real policy logprobs are available, # report sequence-level generation/training mismatch. A threshold adds @@ -3276,6 +3292,21 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: self._opd_stat_sumsq += float((valid * valid).sum()) self._opd_stat_count += int(valid.numel()) + # Clip after logging, not before: grpo.py logs at :3471 and clips at + # :3487, grpo_sync.py at :903 then :915. Both report pre-clip + # advantages, so clipping first would make SC's advantages/mean|max|min + # mean something different from every other driver's. + # + # A no-op unless advantage_clip_low/high are set, and GRPO-only: + # _clip_grpo_advantages has no counterpart in ppo.py. + # ``hasattr``: only GRPOConfig carries the two knobs. DistillationConfig + # does not, and the OPD test doubles build a bare namespace -- both mean + # "no clipping configured". ``not self._is_ppo`` is load-bearing on top + # of that, because PPOConfig is ``extra="allow"``: a user who sets + # ppo.advantage_clip_low would otherwise get GRPO clipping on a PPO run. + if not self._is_ppo and hasattr(self._algo_cfg, "advantage_clip_low"): + advantages = _clip_grpo_advantages(advantages, self._algo_cfg) + fields_to_put = {adv_cfg.output_field: advantages} if seq_logprob_error_threshold is not None: fields_to_put[adv_cfg.sample_mask_field] = sample_mask diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 03b2998a0eb..5174aaeb387 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -775,14 +775,15 @@ 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; async GRPO + # rejects three of them the same way. An entry leaves this list when the SC + # path starts implementing it -- rejecting a knob is only right while + # nobody honours it, which is why reward_scaling is no longer here. unsupported = [ name for name, enabled in ( ("overlong_filtering", algo_cfg.overlong_filtering), ("use_dynamic_sampling", algo_cfg.use_dynamic_sampling), - ("reward_scaling", algo_cfg.reward_scaling.enabled), ("reward_shaping", algo_cfg.reward_shaping.enabled), ) if enabled diff --git a/tests/unit/single_controller/test_ppo_setup.py b/tests/unit/single_controller/test_ppo_setup.py index bcc3d791f38..2c03232319e 100644 --- a/tests/unit/single_controller/test_ppo_setup.py +++ b/tests/unit/single_controller/test_ppo_setup.py @@ -331,13 +331,11 @@ def test_accepts_varying_ckpt_structure_with_warmup(self): [ lambda cfg: setattr(cfg, "overlong_filtering", True), lambda cfg: setattr(cfg, "use_dynamic_sampling", True), - lambda cfg: setattr(cfg.reward_scaling, "enabled", True), lambda cfg: setattr(cfg.reward_shaping, "enabled", True), ], ids=[ "overlong_filtering", "use_dynamic_sampling", - "reward_scaling", "reward_shaping", ], ) @@ -740,3 +738,15 @@ def test_worker_group_slots( else: # The critic never lands on the inference cluster. assert inference.kwargs["max_colocated_worker_groups"] == 1 + + def test_reward_scaling_is_implemented_not_rejected(self): + """It is off the unsupported list because _advantage_stage applies it. + + The list exists so an enabled knob cannot silently do nothing. Once the + stage calls the same ``scale_rewards`` helper grpo.py uses, rejecting it + would refuse a run the path now handles. + """ + mc = _ppo_master_config() + mc.ppo.reward_scaling.enabled = True + + validate_single_controller_config(mc) diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index d9c7d359b21..2daa0f42a21 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -44,6 +44,37 @@ from nemo_rl.utils.timer import TimeoutChecker, Timer +def _grpo_stub(**overrides): + """The grpo config fields ``_advantage_stage`` reads, with real defaults. + + Kept in one place: the stage reads a widening set of algorithm-config knobs, + and four tests were previously constructing that namespace by hand. + """ + fields = { + "seq_logprob_error_threshold": None, + "reward_scaling": GRPOConfig.model_fields["reward_scaling"].default_factory(), + "advantage_clip_low": None, + "advantage_clip_high": None, + } + fields.update(overrides) + return SimpleNamespace(**fields) + + +def _ppo_stub(**overrides): + """The ppo config fields ``_advantage_stage`` reads. + + Deliberately without ``advantage_clip_low``/``_high``: real PPOConfig + declares neither, so a stub that carries them would hide whether the + ``not self._is_ppo`` half of the clip guard does anything. + """ + fields = { + "seq_logprob_error_threshold": None, + "reward_scaling": GRPOConfig.model_fields["reward_scaling"].default_factory(), + } + fields.update(overrides) + return SimpleNamespace(**fields) + + class FakeWeightSynchronizer: pass @@ -567,7 +598,7 @@ def test_advantage_stage_applies_seq_logprob_error_mask_before_streaming_train( ctrl._teacher_logprobs_required = False ctrl._is_ppo = False ctrl._master_config = SimpleNamespace( - grpo=SimpleNamespace(seq_logprob_error_threshold=2.0) + grpo=_grpo_stub(seq_logprob_error_threshold=2.0) ) ctrl._algo_cfg = ctrl._master_config.grpo ctrl._step_log_dict = { @@ -636,7 +667,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) + grpo=_grpo_stub(seq_logprob_error_threshold=None) ) ctrl._algo_cfg = ctrl._master_config.grpo ctrl._step_log_dict = { @@ -699,7 +730,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) + grpo=_grpo_stub(seq_logprob_error_threshold=2.0) ) ctrl._algo_cfg = ctrl._master_config.grpo ctrl._step_log_dict = { @@ -755,7 +786,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) + grpo=_grpo_stub(seq_logprob_error_threshold=None) ) ctrl._algo_cfg = ctrl._master_config.grpo ctrl._step_log_dict = { @@ -1739,9 +1770,7 @@ def compute_advantage(self, *, rewards, mask, **kwargs): ctrl._reference_logprobs_required = False ctrl._teacher_logprobs_required = False ctrl._is_ppo = True - ctrl._master_config = SimpleNamespace( - ppo=SimpleNamespace(seq_logprob_error_threshold=None) - ) + ctrl._master_config = SimpleNamespace(ppo=_ppo_stub()) ctrl._algo_cfg = ctrl._master_config.ppo ctrl._step_log_dict = { "rewards": [], @@ -1769,3 +1798,155 @@ def compute_advantage(self, *, rewards, mask, **kwargs): ) assert "returns" in (result_meta.fields or []) assert "advantages" in (result_meta.fields or []) + + +# --------------------------------------------------------------------------- +# grpo.* knobs that every other driver applies and SC did not. +# grpo.py:3321 / grpo.py:4883 / grpo_sync.py:915 clip advantages; grpo.py:3024 +# scales rewards. The SC advantage stage is the same point in the pipeline. +# --------------------------------------------------------------------------- + + +class _RewardRecordingAdvantageEstimator: + """Advantage == reward, broadcast, so the written column is readable.""" + + def __init__(self) -> None: + self.rewards: torch.Tensor | None = None + + def compute_advantage(self, *, rewards, mask, **kwargs): + del kwargs + self.rewards = rewards.clone() + return rewards.unsqueeze(-1).expand_as(mask).clone() + + +def _knob_ctrl(estimator, grpo_stub, *, is_ppo=False): + batch, seq = 2, 4 + columns = { + "prompt_ids_for_adv": torch.zeros(batch, seq, dtype=torch.long), + "total_reward": torch.tensor([-4.0, 6.0]), + "token_mask": torch.ones(batch, seq), + "sample_mask": torch.ones(batch), + } + if is_ppo: + # The PPO branch fetches the critic's values column. + columns["values"] = torch.zeros(batch, seq) + data = TensorDict(columns, batch_size=[batch]) + data_plane = _AdvantageDataPlane(data) + 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._master_config = ( + SimpleNamespace(ppo=grpo_stub) if is_ppo else SimpleNamespace(grpo=grpo_stub) + ) + ctrl._algo_cfg = grpo_stub + ctrl._is_ppo = is_ppo + ctrl._step_log_dict = { + "rewards": [], + "masked_advantages": [], + "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)], + fields=list(data.keys()), + ) + return ctrl, data_plane, meta + + +def test_advantage_clip_bounds_are_applied_before_the_write() -> None: + estimator = _RewardRecordingAdvantageEstimator() + ctrl, data_plane, meta = _knob_ctrl( + estimator, _grpo_stub(advantage_clip_low=-1.0, advantage_clip_high=2.0) + ) + + asyncio.run(ctrl._advantage_stage(meta)) + + written = data_plane.written_fields["advantages"] + assert written.min().item() == pytest.approx(-1.0) + assert written.max().item() == pytest.approx(2.0) + + # ...but the logged advantages stay pre-clip, because grpo.py:3471 and + # grpo_sync.py:903 both log before they clip. Clipping first would leave + # SC's advantages/mean|max|min meaning something no other driver's does. + logged = torch.cat([t.flatten() for t in ctrl._step_log_dict["masked_advantages"]]) + assert logged.min().item() == pytest.approx(-4.0) + assert logged.max().item() == pytest.approx(6.0) + + +def test_a_ppo_run_ignores_advantage_clip_knobs_set_on_its_own_block() -> None: + """``not self._is_ppo`` is the guard, not ``hasattr``. + + PPOConfig is ``extra="allow"``, so a user who sets ``ppo.advantage_clip_low`` + gets the field -- and ``hasattr`` alone would then run GRPO clipping on a + PPO run, which ppo.py never does. + """ + + class _GaeLikeEstimator: + def compute_advantage(self, *, rewards, mask, **kwargs): + del kwargs + adv = rewards.unsqueeze(-1).expand_as(mask).clone() + return adv, adv + 1.0 + + ctrl, data_plane, meta = _knob_ctrl( + _GaeLikeEstimator(), + _ppo_stub(advantage_clip_low=-1.0, advantage_clip_high=2.0), + is_ppo=True, + ) + + asyncio.run(ctrl._advantage_stage(meta)) + + written = data_plane.written_fields["advantages"] + assert written.min().item() == pytest.approx(-4.0) + assert written.max().item() == pytest.approx(6.0) + + +def test_advantages_are_untouched_when_no_clip_bounds_are_set() -> None: + estimator = _RewardRecordingAdvantageEstimator() + ctrl, data_plane, meta = _knob_ctrl(estimator, _grpo_stub()) + + asyncio.run(ctrl._advantage_stage(meta)) + + written = data_plane.written_fields["advantages"] + assert written.min().item() == pytest.approx(-4.0) + assert written.max().item() == pytest.approx(6.0) + + +def test_only_the_configured_clip_bound_applies() -> None: + estimator = _RewardRecordingAdvantageEstimator() + ctrl, data_plane, meta = _knob_ctrl(estimator, _grpo_stub(advantage_clip_high=2.0)) + + asyncio.run(ctrl._advantage_stage(meta)) + + written = data_plane.written_fields["advantages"] + assert written.max().item() == pytest.approx(2.0) + assert written.min().item() == pytest.approx(-4.0), "low bound was not set" + + +def test_reward_scaling_reaches_the_estimator() -> None: + scaling = GRPOConfig.model_fields["reward_scaling"].default_factory() + scaling.enabled = True + scaling.source_min, scaling.source_max = -4.0, 6.0 + scaling.target_min, scaling.target_max = 0.0, 1.0 + + estimator = _RewardRecordingAdvantageEstimator() + ctrl, _, meta = _knob_ctrl(estimator, _grpo_stub(reward_scaling=scaling)) + + asyncio.run(ctrl._advantage_stage(meta)) + + torch.testing.assert_close(estimator.rewards, torch.tensor([0.0, 1.0])) + + +def test_reward_scaling_disabled_leaves_rewards_alone() -> None: + estimator = _RewardRecordingAdvantageEstimator() + ctrl, _, meta = _knob_ctrl(estimator, _grpo_stub()) + + asyncio.run(ctrl._advantage_stage(meta)) + + torch.testing.assert_close(estimator.rewards, torch.tensor([-4.0, 6.0]))