fix(sc): apply advantage clipping and overlong filtering

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 — `reward_shaping`, `reward_scaling`, and `use_dynamic_sampling` are implemented on neither algorithm block, so setup rejects them rather than silently skipping the shaping. `overlong_filtering` is supported and excludes truncated completions from the loss through `sample_mask`; a step in which every completion is truncated is still 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.

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

nemo_rl/algorithms/async_utils/replay_buffer.py

@@ -981,12 +981,14 @@ def __init__(
         partition_id: str,
         *,
         pad_value_dict: Mapping[str, int],
+        overlong_filtering: bool,
         include_message_violation_fields: bool,
         require_routed_experts: bool = False,
     ):
         self._dp_client = dp_client
         self._partition_id = partition_id
         self._pad_value_dict = dict(pad_value_dict)
+        self._overlong_filtering = overlong_filtering
         self._include_message_violation_fields = include_message_violation_fields
         self._require_routed_experts = require_routed_experts
         self.meta_list: list[Optional[KVBatchMeta]] = []
@@ -1095,6 +1097,7 @@ async def commit(
         train_batch = record_to_train_batch(
             record,
             pad_value_dict=self._pad_value_dict,
+            overlong_filtering=self._overlong_filtering,
             include_message_violation_fields=self._include_message_violation_fields,
         )
         sample_ids, fields, tags = pack_payload(

nemo_rl/algorithms/single_controller.py

@@ -78,6 +78,7 @@
 )
 from nemo_rl.algorithms.grpo import (
     GRPOSaveState,
+    _clip_grpo_advantages,
     _write_latest_checkpoint_status,
     compute_and_apply_seq_logprob_error_masking,
 )
@@ -1856,8 +1857,9 @@ async def _train_pump(self) -> None:
                     if self._is_ppo and not has_valid_training_tokens:
                         raise RuntimeError(
                             "SingleController has no valid response tokens after "
-                            "filtering. Check ppo.seq_logprob_error_threshold to "
-                            "avoid an optimizer step with an empty batch."
+                            "filtering. Check ppo.overlong_filtering and "
+                            "ppo.seq_logprob_error_threshold to avoid an optimizer "
+                            "step with an empty batch."
                         )
 
                     # ---- 3. Train the model -- train_microbatches_from_meta ----
@@ -1988,8 +1990,9 @@ async def _train_pump(self) -> None:
                     if not step_open:
                         raise RuntimeError(
                             "SingleController has no valid response tokens after "
-                            "filtering. Check grpo.seq_logprob_error_threshold to "
-                            "avoid an optimizer step with an empty batch."
+                            "filtering. Check grpo.overlong_filtering and "
+                            "grpo.seq_logprob_error_threshold to avoid an optimizer "
+                            "step with an empty batch."
                         )
 
                     with self._timer.time("policy_training"):
@@ -3322,15 +3325,25 @@ 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 self._master_config.grpo is not None
+            advantages = _clip_grpo_advantages(
+                advantages,
+                self._master_config.grpo,
+            )
+        self._step_log_dict["masked_advantages"].append(
+            torch.masked_select(advantages, mask.bool()).detach().cpu()
+        )
+
         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

nemo_rl/algorithms/single_controller_utils/config.py

@@ -771,8 +771,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 knobs
+    SC reads on neither path.
     """
     algo_cfg = algo_config(master_config)
 
@@ -786,12 +786,12 @@ 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 while
+    # converting completions into the TransferQueue training payload.
     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),

nemo_rl/algorithms/single_controller_utils/setup.py

@@ -1412,6 +1412,7 @@ def _build_generation_then_trainer(
         dp_client,
         partition_id=partition_id,
         pad_value_dict={"token_ids": pad_id, "input_ids": pad_id},
+        overlong_filtering=algo_cfg.overlong_filtering,
         include_message_violation_fields=(
             algo_cfg.invalid_tool_call_advantage is not None
             or algo_cfg.malformed_thinking_advantage is not None

nemo_rl/experience/payload.py

@@ -91,13 +91,15 @@ def record_to_train_batch(
     record: PromptGroupRecord,
     *,
     pad_value_dict: Mapping[str, int],
+    overlong_filtering: bool,
     include_message_violation_fields: bool,
 ) -> BatchedDataDict[Any]:
     """Convert one prompt group's record into a packed BatchedDataDict of N rows.
 
     Args:
         record: Rollout's PromptGroupRecord with N completions to flatten into rows.
         pad_value_dict: Field-name → pad value used by batched_message_log_to_flat_message.
+        overlong_filtering: Whether truncated completions should be excluded from loss.
         include_message_violation_fields: Whether to tensorize message violation
             flags for configured advantage penalties.
 
@@ -147,6 +149,9 @@ def record_to_train_batch(
         [float(c.reward) for c in completions], dtype=torch.float32
     )
     sample_mask = torch.ones(n, dtype=torch.float32)
+    if overlong_filtering:
+        truncated = torch.tensor([c.truncated for c in completions], dtype=torch.bool)
+        sample_mask[truncated] = 0
 
     train_data: dict[str, Any] = {
         "input_ids": flat["token_ids"],

tests/unit/experience/test_payload.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import pytest
 import torch
 
 from nemo_rl.data_plane.schema import (
@@ -42,6 +43,7 @@ def _completion(
     *,
     env_token_ids: tuple[int, ...] = (30,),
     with_routes: bool = True,
+    truncated: bool = False,
 ) -> Completion:
     message_log = [
         {
@@ -70,7 +72,7 @@ def _completion(
     return Completion(
         message_log=message_log,
         env_extras=None,
-        truncated=False,
+        truncated=truncated,
         reward=reward,
     )
 
@@ -107,6 +109,7 @@ def test_record_to_train_batch_preserves_routed_experts_in_tq_payload() -> None:
     train_batch = record_to_train_batch(
         record,
         pad_value_dict={"token_ids": 0, "input_ids": 0},
+        overlong_filtering=False,
         include_message_violation_fields=False,
     )
 
@@ -163,6 +166,7 @@ def test_record_to_train_batch_preserves_message_violation_masks() -> None:
     train_batch = record_to_train_batch(
         _record([invalid, malformed]),
         pad_value_dict={"token_ids": 0, "input_ids": 0},
+        overlong_filtering=False,
         include_message_violation_fields=True,
     )
 
@@ -204,6 +208,7 @@ def test_record_to_train_batch_preserves_clean_masks_when_enabled() -> None:
     train_batch = record_to_train_batch(
         _record([_completion(route_start=10, reward=1.0)]),
         pad_value_dict={"token_ids": 0, "input_ids": 0},
+        overlong_filtering=False,
         include_message_violation_fields=True,
     )
 
@@ -220,6 +225,7 @@ def test_record_to_train_batch_omits_routed_experts_when_absent() -> None:
     train_batch = record_to_train_batch(
         record,
         pad_value_dict={"token_ids": 0, "input_ids": 0},
+        overlong_filtering=False,
         include_message_violation_fields=False,
     )
     assert "routed_experts" not in train_batch
@@ -235,6 +241,36 @@ def test_record_to_train_batch_omits_routed_experts_when_absent() -> None:
     assert "routed_experts" not in fields
 
 
+@pytest.mark.parametrize(
+    ("overlong_filtering", "expected_sample_mask"),
+    [
+        (False, [1.0, 1.0, 1.0]),
+        (True, [1.0, 0.0, 1.0]),
+    ],
+)
+def test_record_to_train_batch_masks_truncated_completions_when_enabled(
+    overlong_filtering: bool,
+    expected_sample_mask: list[float],
+) -> None:
+    record = _record(
+        [
+            _completion(route_start=10, reward=1.0),
+            _completion(route_start=30, reward=2.0, truncated=True),
+            _completion(route_start=50, reward=3.0),
+        ]
+    )
+
+    train_batch = record_to_train_batch(
+        record,
+        pad_value_dict={"token_ids": 0, "input_ids": 0},
+        overlong_filtering=overlong_filtering,
+        include_message_violation_fields=False,
+    )
+
+    assert train_batch["sample_mask"].dtype == torch.float32
+    assert train_batch["sample_mask"].tolist() == expected_sample_mask
+
+
 def _failed_completion() -> Completion:
     """A trajectory whose first generation raised: prompt only, no routes."""
     return Completion(
@@ -258,6 +294,7 @@ def test_record_to_train_batch_backfills_routes_for_failed_completion() -> None:
     train_batch = record_to_train_batch(
         record,
         pad_value_dict={"token_ids": 0, "input_ids": 0},
+        overlong_filtering=False,
         include_message_violation_fields=False,
     )
 
@@ -294,6 +331,7 @@ def test_pack_payload_stamps_violation_counts_on_tags() -> None:
     train_batch = record_to_train_batch(
         _record(completions),
         pad_value_dict={"token_ids": 0, "input_ids": 0},
+        overlong_filtering=False,
         include_message_violation_fields=False,
     )
     _, fields, tags = pack_payload(

tests/unit/single_controller/_checkpoint_scenarios.py

@@ -201,9 +201,10 @@ def _stub_converter(
     record: PromptGroupRecord,
     *,
     pad_value_dict: Any,
+    overlong_filtering: bool,
     include_message_violation_fields: bool,
 ):
-    del record, pad_value_dict, include_message_violation_fields
+    del record, pad_value_dict, overlong_filtering, include_message_violation_fields
     return BatchedDataDict[Any](
         {
             "input_ids": torch.ones((ROLLOUTS_PER_GROUP, 3), dtype=torch.long),
@@ -235,6 +236,7 @@ def _new_buffer(dp: NoOpDataPlaneClient) -> TQReplayBuffer:
         dp,
         partition_id=PARTITION,
         pad_value_dict={"input_ids": 0},
+        overlong_filtering=False,
         include_message_violation_fields=False,
         require_routed_experts=False,
     )

tests/unit/single_controller/test_checkpoint_dispatch_races.py

@@ -679,6 +679,7 @@ async def exercise() -> None:
             dp_client,
             partition_id="rollout_data",
             pad_value_dict={"input_ids": 0},
+            overlong_filtering=False,
             include_message_violation_fields=False,
             require_routed_experts=False,
         )
@@ -783,6 +784,7 @@ async def exercise() -> None:
             dp_client,
             partition_id="rollout_data",
             pad_value_dict={"input_ids": 0},
+            overlong_filtering=False,
             include_message_violation_fields=False,
             require_routed_experts=False,
         )

tests/unit/single_controller/test_ppo_setup.py

@@ -328,13 +328,11 @@ def test_accepts_varying_ckpt_structure_with_warmup(self):
     @pytest.mark.parametrize(
         "enable",
         [
-            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",
@@ -349,14 +347,16 @@ def test_rejects_shaping_the_sc_path_does_not_implement(self, enable):
         ):
             validate_single_controller_config(mc)
 
-    def test_rejects_shaping_on_a_grpo_run_too(self):
-        mc = _make_master_config()
-        mc.grpo.overlong_filtering = True
+    @pytest.mark.parametrize(
+        "make_config",
+        [_make_master_config, _ppo_master_config],
+        ids=["grpo", "ppo"],
+    )
+    def test_accepts_overlong_filtering(self, make_config):
+        mc = make_config()
+        algo_config(mc).overlong_filtering = True
 
-        with pytest.raises(
-            NotImplementedError, match="overlong_filtering not supported"
-        ):
-            validate_single_controller_config(mc)
+        validate_single_controller_config(mc)
 
     def test_rejects_a_dtensor_critic(self):
         """Only the Megatron value worker carries TQWorkerMixin."""

tests/unit/single_controller/test_rollout_pump.py

@@ -1204,6 +1204,7 @@ def test_rollout_pump_writes_expected_tq_data(
         dp_adapter,
         partition_id=_PARTITION_ID,
         pad_value_dict={"token_ids": int(tokenizer.pad_token_id or 0)},
+        overlong_filtering=False,
         include_message_violation_fields=False,
     )
     rollout_manager = RolloutManager(