From 748251432c4432e9b13e355b4e6bae334c8fb4c0 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sat, 29 Aug 2026 21:55:16 -0700 Subject: [PATCH] fix(sc): address sample masking review feedback Signed-off-by: Anish Mahishi --- docs/guides/single-controller.md | 2 +- nemo_rl/algorithms/single_controller.py | 39 ++++-- .../single_controller_utils/config.py | 3 +- .../single_controller_utils/utils.py | 8 +- nemo_rl/data_plane/schema.py | 6 + nemo_rl/experience/payload.py | 16 ++- nemo_rl/experience/rollouts.py | 21 ++- tests/unit/data_plane/test_codec_mooncake.py | 22 +++ tests/unit/experience/test_payload.py | 55 +++++++- .../unit/experience/test_reward_penalties.py | 6 +- tests/unit/single_controller/_dp_fakes.py | 2 + .../unit/single_controller/test_ppo_setup.py | 14 +- .../single_controller/test_rollout_pump.py | 2 + .../test_single_controller_actor.py | 126 +++++++++++++++--- .../single_controller/test_train_pump_e2e.py | 4 + tests/unit/single_controller/test_utils.py | 2 + 16 files changed, 270 insertions(+), 58 deletions(-) diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index 6e6fd693221..6b483977d6e 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 — `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. - 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 9d5859551a9..9cd8cde012a 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -434,6 +434,7 @@ def __init__( self._step_log_dict: dict[str, list] = { "rewards": [], "masked_advantages": [], + "num_mask_sample_filtered": [], "sequence_lengths": [], "seq_logprob_error_metrics": [], **{key: [] for key in VIOLATION_TAG_KEYS}, @@ -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 seq_logprob_error_threshold, " + "overlong_filtering, and environment mask_sample flags " + "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 seq_logprob_error_threshold, " + "overlong_filtering, and environment mask_sample flags " + "to avoid an optimizer step with an empty batch." ) with self._timer.time("policy_training"): @@ -3203,6 +3206,18 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: sample_mask = squeeze_trailing_unit_dim( tensor_field(data, adv_cfg.sample_mask_field) ).float() + mask_sample = squeeze_trailing_unit_dim( + tensor_field(data, adv_cfg.mask_sample_field) + ).bool() + truncated = squeeze_trailing_unit_dim( + tensor_field(data, adv_cfg.truncated_field) + ).bool() + + num_mask_sample_filtered = int(mask_sample.sum().item()) + self._step_log_dict["num_mask_sample_filtered"].append(num_mask_sample_filtered) + final_sample_mask = sample_mask * (~mask_sample).to(sample_mask.dtype) + if self._algo_cfg.overlong_filtering: + final_sample_mask = final_sample_mask * (~truncated).to(sample_mask.dtype) seq_logprob_error_threshold = self._algo_cfg.seq_logprob_error_threshold # Match the legacy path: whenever real policy logprobs are available, @@ -3212,7 +3227,7 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: masking_data = BatchedDataDict( { "token_mask": token_mask, - "sample_mask": sample_mask, + "sample_mask": final_sample_mask, "prev_logprobs": tensor_field( data, adv_cfg.policy_logprobs_field, @@ -3224,7 +3239,7 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: } ) num_valid_seqs_before = float( - ((token_mask[:, 1:] * sample_mask.unsqueeze(-1)).sum(dim=-1) > 0) + ((token_mask[:, 1:] * final_sample_mask.unsqueeze(-1)).sum(dim=-1) > 0) .sum() .item() ) @@ -3233,9 +3248,9 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: rewards=rewards, seq_logprob_error_threshold=seq_logprob_error_threshold, ) - sample_mask = masking_data["sample_mask"] + final_sample_mask = masking_data["sample_mask"] num_valid_seqs_after = float( - ((token_mask[:, 1:] * sample_mask.unsqueeze(-1)).sum(dim=-1) > 0) + ((token_mask[:, 1:] * final_sample_mask.unsqueeze(-1)).sum(dim=-1) > 0) .sum() .item() ) @@ -3246,7 +3261,7 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: seq_error_metrics["_num_valid_seqs_after"] = num_valid_seqs_after self._step_log_dict["seq_logprob_error_metrics"].append(seq_error_metrics) - mask = token_mask * sample_mask.unsqueeze(-1) + mask = token_mask * final_sample_mask.unsqueeze(-1) repeated_batch: dict[str, torch.Tensor] = { "total_reward": rewards, @@ -3332,8 +3347,8 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: self._opd_stat_count += int(valid.numel()) 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 + if not torch.equal(final_sample_mask, sample_mask): + fields_to_put[adv_cfg.sample_mask_field] = final_sample_mask new_fields = [adv_cfg.output_field] if returns is not None: fields_to_put[adv_cfg.returns_field] = returns @@ -3360,6 +3375,8 @@ def _advantage_input_fields(self) -> list[str]: adv_cfg.token_mask_field, adv_cfg.sample_mask_field, *adv_cfg.repeated_batch_fields, + adv_cfg.mask_sample_field, + adv_cfg.truncated_field, ] if self._message_level_advantage_penalties_enabled: fields.extend( diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 3ecfc0a257e..e02083b02b9 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -791,7 +791,6 @@ def _validate_algo_settings(master_config: MasterConfig) -> None: 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), @@ -1156,6 +1155,8 @@ class AdvantageConfig: sample_mask_field: str = "sample_mask" invalid_tool_call_mask_field: str = INVALID_TOOL_CALL_MASK malformed_thinking_mask_field: str = MALFORMED_THINKING_MASK + mask_sample_field: str = "mask_sample" + truncated_field: str = "truncated" repeated_batch_fields: list[str] = field(default_factory=list) policy_logprobs_field: str = "prev_logprobs" generation_logprobs_field: str = "generation_logprobs" diff --git a/nemo_rl/algorithms/single_controller_utils/utils.py b/nemo_rl/algorithms/single_controller_utils/utils.py index 1b4f6dab4ee..8fd5804dccc 100644 --- a/nemo_rl/algorithms/single_controller_utils/utils.py +++ b/nemo_rl/algorithms/single_controller_utils/utils.py @@ -96,6 +96,7 @@ def reduce_advantage_pump_metrics( sequence_lengths: list[int], *, seq_logprob_error_metrics: list[dict[str, float]] | None = None, + num_mask_sample_filtered: list[int] | None = None, num_invalid_tool_calls: list[int] | None = None, num_malformed_thinking: list[int] | None = None, num_assistant_messages: list[int] | None = None, @@ -108,13 +109,16 @@ def reduce_advantage_pump_metrics( sequence_lengths: All input_lengths trained on this step. seq_logprob_error_metrics: Sequence-error metrics and their aggregation counts, one record per streaming chunk. + num_mask_sample_filtered: Environment-flagged sample counts, one per + streaming chunk. num_invalid_tool_calls: Per-sample invalid tool-call counts. num_malformed_thinking: Per-sample malformed-thinking counts. num_assistant_messages: Per-sample assistant message counts (rate denominator). Returns: Step-level reward, advantage, token-count, optional sequence - log-probability error metrics, and per-sample violation counts. + log-probability error metrics, the num_mask_sample_filtered count, and + per-sample violation counts. """ out: dict[str, float] = {} @@ -132,6 +136,8 @@ def reduce_advantage_pump_metrics( out["advantages/min"] = 0.0 if sequence_lengths: out["total_num_tokens"] = float(sum(sequence_lengths)) + if num_mask_sample_filtered is not None: + out["num_mask_sample_filtered"] = float(sum(num_mask_sample_filtered)) if seq_logprob_error_metrics: out.update(_reduce_seq_logprob_error_metrics(seq_logprob_error_metrics)) n_asst = sum(num_assistant_messages or []) diff --git a/nemo_rl/data_plane/schema.py b/nemo_rl/data_plane/schema.py index 4a7dbd8d5e0..6ed1fb64317 100644 --- a/nemo_rl/data_plane/schema.py +++ b/nemo_rl/data_plane/schema.py @@ -28,6 +28,8 @@ INPUT_IDS = "input_ids" INPUT_LENGTHS = "input_lengths" SAMPLE_MASK = "sample_mask" +MASK_SAMPLE = "mask_sample" +TRUNCATED = "truncated" META_IDX = "meta_idx" # Token-aligned message-violation fields consumed by SingleController advantages. @@ -55,6 +57,8 @@ # TransferQueue's lazy field-name registration race. SC_ROLLOUT_SCHEMA_FIELDS = ( *DP_TRAIN_FIELDS, + MASK_SAMPLE, + TRUNCATED, "prompt_ids_for_adv", "total_reward", "values", @@ -117,8 +121,10 @@ PROMOTE_1D_FIELDS: frozenset[str] = frozenset( { INPUT_LENGTHS, + MASK_SAMPLE, "total_reward", SAMPLE_MASK, + TRUNCATED, } ) diff --git a/nemo_rl/experience/payload.py b/nemo_rl/experience/payload.py index b3b34432316..b0441937942 100644 --- a/nemo_rl/experience/payload.py +++ b/nemo_rl/experience/payload.py @@ -27,7 +27,9 @@ from nemo_rl.data_plane.schema import ( INVALID_TOOL_CALL_MASK, MALFORMED_THINKING_MASK, + MASK_SAMPLE, ROUTED_EXPERTS_FIELD, + TRUNCATED, ) from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.experience.interfaces import PromptGroupRecord @@ -102,8 +104,9 @@ def record_to_train_batch( flags for configured advantage penalties. Returns: - BatchedDataDict with input_ids, input_lengths, generation_logprobs, token_mask, - sample_mask, prompt_ids_for_adv, total_reward, violation counts, and optional + BatchedDataDict with input_ids, input_lengths, generation_logprobs, + token_mask, an all-ones sample_mask, the raw mask_sample and truncated + flags, prompt_ids_for_adv, total_reward, violation counts, and optional routed experts and message-violation masks. """ # Lazy imports: grpo and llm_message_utils transitively pull @@ -113,7 +116,10 @@ def record_to_train_batch( extract_initial_prompt_messages, ) from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message - from nemo_rl.experience.rollouts import backfill_missing_routed_experts + from nemo_rl.experience.rollouts import ( + _mask_sample_flags, + backfill_missing_routed_experts, + ) completions = record.completions n = len(completions) @@ -146,6 +152,8 @@ def record_to_train_batch( total_reward = torch.tensor( [float(c.reward) for c in completions], dtype=torch.float32 ) + mask_sample = _mask_sample_flags(c.env_extras for c in completions) + truncated = torch.tensor([c.truncated for c in completions], dtype=torch.bool) sample_mask = torch.ones(n, dtype=torch.float32) train_data: dict[str, Any] = { @@ -155,6 +163,8 @@ def record_to_train_batch( "token_mask": flat["token_loss_mask"], "sample_mask": sample_mask, "prompt_ids_for_adv": prompt_flat["token_ids"], + MASK_SAMPLE: mask_sample, + TRUNCATED: truncated, "total_reward": total_reward, _VIOLATION_COUNTS_KEY: violation_counts, } diff --git a/nemo_rl/experience/rollouts.py b/nemo_rl/experience/rollouts.py index 0ca333aa242..898bf8e55d8 100644 --- a/nemo_rl/experience/rollouts.py +++ b/nemo_rl/experience/rollouts.py @@ -21,7 +21,7 @@ import statistics import warnings from collections import defaultdict -from collections.abc import AsyncGenerator, Mapping, Sequence +from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence from dataclasses import dataclass from typing import Any, Optional @@ -49,6 +49,7 @@ attach_image_model_inputs_to_message, extract_input_images_from_responses_messages, ) +from nemo_rl.data_plane.schema import MASK_SAMPLE from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.environments.interfaces import ( EnvironmentInterface, @@ -279,16 +280,12 @@ def _add_r3_fallback_metrics( ) -def _extract_mask_sample_flags(results: list[dict[str, Any]]) -> torch.Tensor: +def _mask_sample_flags(extras: Iterable[dict[str, Any] | None]) -> torch.Tensor: """Return True for samples the environment asks GRPO to mask from loss.""" return torch.tensor( [ - bool( - (result["full_result"].get("instance_config") or {}).get( - "mask_sample", False - ) - ) - for result in results + bool(((extra or {}).get("instance_config") or {}).get(MASK_SAMPLE, False)) + for extra in extras ], dtype=torch.bool, ) @@ -2848,10 +2845,12 @@ def _postprocess_single_nemo_gym_group( ), } ) - # Env/agent mask flag: flagged samples are dropped from the loss but still - # count for advantages. env.should_mask_flagged_samples=false skips this. + # Carry the raw env/agent flag downstream; the advantage stage composes it + # into sample_mask. env.should_mask_flagged_samples=false skips this. if mask_env_flagged_samples: - final_batch["mask_sample"] = _extract_mask_sample_flags(results) + final_batch[MASK_SAMPLE] = _mask_sample_flags( + result["full_result"] for result in results + ) rollout_metrics.update(_effort_shaping_metrics(shaping)) diff --git a/tests/unit/data_plane/test_codec_mooncake.py b/tests/unit/data_plane/test_codec_mooncake.py index f2392701111..bced5792717 100644 --- a/tests/unit/data_plane/test_codec_mooncake.py +++ b/tests/unit/data_plane/test_codec_mooncake.py @@ -75,6 +75,28 @@ def test_promote_1d_roundtrip_via_from_wire() -> None: assert torch.equal(back["input_lengths"], original) +@pytest.mark.parametrize("field_name", ["mask_sample", "truncated"]) +def test_raw_sample_filter_fields_roundtrip_as_dense_1d(field_name: str) -> None: + """Raw loss-filter fields use the Mooncake scalar wire workaround.""" + from tensordict import TensorDict + + from nemo_rl.data_plane.adapters.transfer_queue import ( + _from_wire, + _promote_1d_leaves, + ) + + n = 4 + original = torch.tensor([False, True, False, True]) + td = TensorDict({field_name: original}, batch_size=[n]) + + wire = _promote_1d_leaves(td) + assert wire[field_name].shape == (n, 1) + + back = _from_wire(wire) + assert back[field_name].shape == (n,) + assert torch.equal(back[field_name], original) + + def test_from_wire_densifies_uniform_nested_rows() -> None: """TQ v0.1.9's uniform nested reads are restored to dense tensors.""" from tensordict import TensorDict diff --git a/tests/unit/experience/test_payload.py b/tests/unit/experience/test_payload.py index c1c32016bfa..a30e8f86c9e 100644 --- a/tests/unit/experience/test_payload.py +++ b/tests/unit/experience/test_payload.py @@ -42,6 +42,8 @@ def _completion( *, env_token_ids: tuple[int, ...] = (30,), with_routes: bool = True, + mask_sample: bool | None = None, + truncated: bool = False, ) -> Completion: message_log = [ { @@ -67,10 +69,15 @@ def _completion( if not with_routes: for message in message_log: message.pop("routed_experts") + env_extras = ( + None + if mask_sample is None + else {"instance_config": {"mask_sample": mask_sample}} + ) return Completion( message_log=message_log, - env_extras=None, - truncated=False, + env_extras=env_extras, + truncated=truncated, reward=reward, ) @@ -235,6 +242,50 @@ def test_record_to_train_batch_omits_routed_experts_when_absent() -> None: assert "routed_experts" not in fields +def test_record_to_train_batch_carries_raw_masks_without_applying_them() -> None: + record = _record( + [ + _completion( + route_start=10, + reward=1.0, + mask_sample=True, + ), + _completion( + route_start=30, + reward=2.0, + mask_sample=False, + 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}, + include_message_violation_fields=False, + ) + + assert torch.equal(train_batch["sample_mask"], torch.ones(3)) + assert torch.equal( + train_batch["mask_sample"], + torch.tensor([True, False, False]), + ) + assert torch.equal( + train_batch["truncated"], + torch.tensor([False, True, False]), + ) + + _, fields, _ = pack_payload( + train_batch, + weight_version=3, + group_id="group", + prompt_idx=17, + ) + assert torch.equal(fields["mask_sample"], train_batch["mask_sample"]) + assert torch.equal(fields["truncated"], train_batch["truncated"]) + + def _failed_completion() -> Completion: """A trajectory whose first generation raised: prompt only, no routes.""" return Completion( diff --git a/tests/unit/experience/test_reward_penalties.py b/tests/unit/experience/test_reward_penalties.py index ea5f4921519..ad03ba44728 100644 --- a/tests/unit/experience/test_reward_penalties.py +++ b/tests/unit/experience/test_reward_penalties.py @@ -19,7 +19,7 @@ from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.experience.rollouts import ( - _extract_mask_sample_flags, + _mask_sample_flags, _postprocess_single_nemo_gym_group, apply_reward_penalties, resolve_reward_penalty_config, @@ -91,7 +91,7 @@ def encode(self, text, add_special_tokens=False): return self.token_map[text] -class TestExtractMaskSampleFlags: +class TestMaskSampleFlags: def test_reads_mask_sample_from_instance_config(self): results = [ {"full_result": {"instance_config": {"mask_sample": True}}}, @@ -101,7 +101,7 @@ def test_reads_mask_sample_from_instance_config(self): {"full_result": {"instance_config": None}}, ] - mask_sample = _extract_mask_sample_flags(results) + mask_sample = _mask_sample_flags(r["full_result"] for r in results) assert mask_sample.dtype == torch.bool assert torch.equal( diff --git a/tests/unit/single_controller/_dp_fakes.py b/tests/unit/single_controller/_dp_fakes.py index 3d63f5915ef..1ec25860564 100644 --- a/tests/unit/single_controller/_dp_fakes.py +++ b/tests/unit/single_controller/_dp_fakes.py @@ -33,6 +33,8 @@ "generation_logprobs", "token_mask", "sample_mask", + "mask_sample", + "truncated", "prompt_ids_for_adv", "total_reward", ] diff --git a/tests/unit/single_controller/test_ppo_setup.py b/tests/unit/single_controller/test_ppo_setup.py index 4ff9a1fd4d9..2541709b3fa 100644 --- a/tests/unit/single_controller/test_ppo_setup.py +++ b/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,12 @@ 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("algorithm", ["grpo", "ppo"]) + def test_accepts_overlong_filtering(self, algorithm: str): + mc = _make_master_config() if algorithm == "grpo" else _ppo_master_config() + getattr(mc, algorithm).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.""" diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index b0a7285b0c4..b260be12f4b 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -1272,6 +1272,8 @@ def test_rollout_pump_writes_expected_tq_data( bulk["sample_mask"].float(), torch.ones(expected_samples, dtype=torch.float32), ) + assert not bulk["mask_sample"].bool().any() + assert not bulk["truncated"].bool().any() # Same deterministic prompt as test_async_rollout_manager: the model # solves the calculator task every time -> reward == 1.0 and decoded diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index 6788af5bff3..ca7c3732077 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -547,13 +547,13 @@ def compute_advantage(self, *, rewards, mask, **kwargs) -> torch.Tensor: return rewards.unsqueeze(-1).expand_as(mask).clone() -def test_advantage_stage_applies_seq_logprob_error_mask_before_streaming_train( +def test_advantage_stage_composes_all_filters_before_computing_advantages( capsys: pytest.CaptureFixture[str], ) -> None: batch_size, sequence_length = 4, 5 generation_logprobs = torch.zeros(batch_size, sequence_length) - # exp(abs(1 - 0)) > the configured threshold of 2, so only row 2 - # should be removed from the loss while the other rows remain trainable. + # Rows 1, 2, and 3 are removed by the environment, sequence-error, + # and overlong masks respectively. Row 0 remains trainable. generation_logprobs[2, 1:] = 1.0 data = TensorDict( { @@ -563,10 +563,12 @@ def test_advantage_stage_applies_seq_logprob_error_mask_before_streaming_train( "total_reward": torch.tensor([0.0, 0.0, 1.0, 0.0]), "token_mask": torch.ones(batch_size, sequence_length), "sample_mask": torch.ones(batch_size), + "mask_sample": torch.tensor([False, True, False, False]), + "truncated": torch.tensor([False, False, False, True]), "prev_logprobs": torch.zeros(batch_size, sequence_length), "generation_logprobs": generation_logprobs, - # The filtered row is also flagged. Its penalty must not overwrite - # the sequence-error mask and leak back into streaming training. + # The sequence-error- and overlong-filtered rows are also flagged. + # Their penalties must not leak back into streaming training. "invalid_tool_call_mask": torch.tensor( [[False] * sequence_length] * 2 + [[True] * sequence_length] * 2 ), @@ -589,6 +591,7 @@ def test_advantage_stage_applies_seq_logprob_error_mask_before_streaming_train( ctrl._master_config = SimpleNamespace( grpo=SimpleNamespace( seq_logprob_error_threshold=2.0, + overlong_filtering=True, invalid_tool_call_advantage=-5.0, malformed_thinking_advantage=None, ) @@ -599,6 +602,7 @@ def test_advantage_stage_applies_seq_logprob_error_mask_before_streaming_train( "rewards": [], "masked_advantages": [], "sequence_lengths": [], + "num_mask_sample_filtered": [], "seq_logprob_error_metrics": [], } meta = KVBatchMeta( @@ -617,21 +621,22 @@ def test_advantage_stage_applies_seq_logprob_error_mask_before_streaming_train( assert "invalid_tool_call_mask" in data_plane.selected_fields assert "generation_logprobs" in data_plane.selected_fields assert data_plane.written_fields is not None - # The estimator's value remains, but the penalty did not overwrite it with - # -5; sample_mask below is what excludes this row from streaming training. + # The estimator's values remain, but the penalty did not overwrite them with + # -5; sample_mask below is what excludes these rows from streaming training. torch.testing.assert_close( data_plane.written_fields["advantages"][2], torch.ones(5) ) torch.testing.assert_close( - data_plane.written_fields["advantages"][3], torch.full((5,), -5.0) + data_plane.written_fields["advantages"][3], torch.zeros(5) ) assert torch.equal( data_plane.written_fields["sample_mask"], - torch.tensor([1.0, 1.0, 0.0, 1.0]), + torch.tensor([1.0, 0.0, 0.0, 0.0]), ) assert estimator.mask is not None - assert estimator.mask[2].count_nonzero() == 0 - assert estimator.mask[[0, 1, 3]].all() + assert estimator.mask[0].all() + assert estimator.mask[1:].count_nonzero() == 0 + assert ctrl._step_log_dict["num_mask_sample_filtered"] == [1] metrics = ctrl._step_log_dict["seq_logprob_error_metrics"] assert len(metrics) == 1 assert metrics[0]["num_masked_seqs_by_logprob_error"] == 1 @@ -640,6 +645,78 @@ def test_advantage_stage_applies_seq_logprob_error_mask_before_streaming_train( assert "advantages" in (result_meta.fields or []) +@pytest.mark.parametrize( + "overlong_filtering, mask_sample, truncated, expected_sample_mask", + [ + (False, [True, False], [True, True], [0.0, 1.0]), + (True, [False, False], [False, True], [1.0, 0.0]), + ], + ids=["env_mask_only", "overlong_only"], +) +def test_advantage_stage_writes_each_sample_filter_without_seq_threshold( + overlong_filtering: bool, + mask_sample: list[bool], + truncated: list[bool], + expected_sample_mask: list[float], +) -> None: + batch_size, sequence_length = 2, 5 + data = TensorDict( + { + "prompt_ids_for_adv": torch.zeros( + batch_size, sequence_length, dtype=torch.long + ), + "total_reward": torch.tensor([1.0, 0.0]), + "token_mask": torch.ones(batch_size, sequence_length), + "sample_mask": torch.ones(batch_size), + "mask_sample": torch.tensor(mask_sample), + "truncated": torch.tensor(truncated), + }, + 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._message_level_advantage_penalties_enabled = False + ctrl._algo_cfg = SimpleNamespace( + seq_logprob_error_threshold=None, + overlong_filtering=overlong_filtering, + ) + 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()), + ) + + _, has_valid_training_tokens = asyncio.run(ctrl._advantage_stage(meta)) + + expected = torch.tensor(expected_sample_mask) + assert has_valid_training_tokens + assert data_plane.written_fields is not None + assert torch.equal(data_plane.written_fields["sample_mask"], expected) + assert estimator.mask is not None + assert torch.equal( + estimator.mask, + data["token_mask"] * expected.unsqueeze(-1), + ) + + def test_advantage_stage_reports_seq_logprob_metrics_without_masking() -> None: batch_size, sequence_length = 2, 5 generation_logprobs = torch.zeros(batch_size, sequence_length) @@ -654,6 +731,8 @@ def test_advantage_stage_reports_seq_logprob_metrics_without_masking() -> None: "sample_mask": torch.ones(batch_size), "prev_logprobs": torch.zeros(batch_size, sequence_length), "generation_logprobs": generation_logprobs, + "mask_sample": torch.zeros(batch_size, dtype=torch.bool), + "truncated": torch.tensor([False, True]), }, batch_size=[batch_size], ) @@ -670,13 +749,14 @@ 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=SimpleNamespace(seq_logprob_error_threshold=None, overlong_filtering=False) ) 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": [], } @@ -700,6 +780,7 @@ def test_advantage_stage_reports_seq_logprob_metrics_without_masking() -> None: metrics = ctrl._step_log_dict["seq_logprob_error_metrics"] assert len(metrics) == 1 assert metrics[0]["num_masked_seqs_by_logprob_error"] == 0 + assert ctrl._step_log_dict["num_mask_sample_filtered"] == [0] assert metrics[0]["max_seq_mult_prob_error"] == pytest.approx(math.e) assert metrics[0]["max_seq_mult_prob_error_after_mask"] == pytest.approx(math.e) @@ -718,6 +799,8 @@ def test_advantage_stage_skips_estimator_when_seq_mask_removes_whole_chunk( "sample_mask": torch.ones(batch_size), "prev_logprobs": torch.zeros(batch_size, sequence_length), "generation_logprobs": torch.ones(batch_size, sequence_length), + "mask_sample": torch.zeros(batch_size, dtype=torch.bool), + "truncated": torch.zeros(batch_size, dtype=torch.bool), }, batch_size=[batch_size], ) @@ -734,13 +817,14 @@ 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=SimpleNamespace(seq_logprob_error_threshold=2.0, overlong_filtering=False) ) 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": [], } @@ -775,6 +859,8 @@ def test_advantage_stage_skips_preexisting_empty_mask_without_seq_threshold() -> "total_reward": torch.tensor([1.0, 0.0]), "token_mask": torch.ones(batch_size, sequence_length), "sample_mask": torch.zeros(batch_size), + "mask_sample": torch.zeros(batch_size, dtype=torch.bool), + "truncated": torch.zeros(batch_size, dtype=torch.bool), }, batch_size=[batch_size], ) @@ -791,13 +877,14 @@ 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=SimpleNamespace(seq_logprob_error_threshold=None, overlong_filtering=False) ) 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": [], } @@ -849,6 +936,8 @@ def get_samples(self, sample_ids, partition_id, select_fields): "total_reward": torch.zeros(2), "token_mask": torch.tensor([[1.0, 1.0, 1.0], [1.0, 0.0, 0.0]]), "sample_mask": torch.ones(2), + "mask_sample": torch.zeros(2, dtype=torch.bool), + "truncated": torch.zeros(2, dtype=torch.bool), "generation_logprobs": torch.full((2, 3), 0.5), "prev_logprobs": torch.full((2, 3), 0.5), "teacher_reference_logprobs": torch.full((2, 3), 0.75), @@ -868,7 +957,7 @@ 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) + grpo=SimpleNamespace(seq_logprob_error_threshold=None, overlong_filtering=False) ) ctrl._algo_cfg = ctrl._master_config.grpo ctrl._message_level_advantage_penalties_enabled = False @@ -877,6 +966,7 @@ def put_samples(self, sample_ids, partition_id, fields): "masked_advantages": [], "sequence_lengths": [], "seq_logprob_error_metrics": [], + "num_mask_sample_filtered": [], } ctrl._opd_stat_sum = 0.0 ctrl._opd_stat_sumsq = 0.0 @@ -1176,6 +1266,7 @@ def _train_pump_controller(*, sampler) -> object: "rewards": [], "masked_advantages": [], "sequence_lengths": [], + "num_mask_sample_filtered": [], "seq_logprob_error_metrics": [], } ctrl._opd_stat_sum = 0.0 @@ -1883,6 +1974,8 @@ def test_advantage_stage_writes_gae_returns_alongside_advantages() -> None: "token_mask": torch.ones(batch_size, sequence_length), "sample_mask": torch.ones(batch_size), "values": torch.zeros(batch_size, sequence_length), + "mask_sample": torch.zeros(batch_size, dtype=torch.bool), + "truncated": torch.zeros(batch_size, dtype=torch.bool), }, batch_size=[batch_size], ) @@ -1908,7 +2001,7 @@ def compute_advantage(self, *, rewards, mask, **kwargs): ctrl._teacher_logprobs_required = False ctrl._is_ppo = True ctrl._master_config = SimpleNamespace( - ppo=SimpleNamespace(seq_logprob_error_threshold=None) + ppo=SimpleNamespace(seq_logprob_error_threshold=None, overlong_filtering=False) ) ctrl._algo_cfg = ctrl._master_config.ppo ctrl._message_level_advantage_penalties_enabled = False @@ -1916,6 +2009,7 @@ def compute_advantage(self, *, rewards, mask, **kwargs): "rewards": [], "masked_advantages": [], "sequence_lengths": [], + "num_mask_sample_filtered": [], "seq_logprob_error_metrics": [], } meta = KVBatchMeta( diff --git a/tests/unit/single_controller/test_train_pump_e2e.py b/tests/unit/single_controller/test_train_pump_e2e.py index 68cc4fd8ea9..fcc107d5323 100644 --- a/tests/unit/single_controller/test_train_pump_e2e.py +++ b/tests/unit/single_controller/test_train_pump_e2e.py @@ -58,6 +58,8 @@ "advantages", "token_mask", "sample_mask", + "mask_sample", + "truncated", "total_reward", "prompt_ids_for_adv", ] @@ -94,6 +96,8 @@ def _populate_group( "input_lengths": torch.tensor([seq_len] * group_size).long(), "token_mask": torch.ones(group_size, seq_len, dtype=torch.long), "sample_mask": torch.ones(group_size, dtype=torch.long), + "mask_sample": torch.zeros(group_size, dtype=torch.bool), + "truncated": torch.zeros(group_size, dtype=torch.bool), "generation_logprobs": torch.zeros( group_size, seq_len, dtype=torch.float32 ), diff --git a/tests/unit/single_controller/test_utils.py b/tests/unit/single_controller/test_utils.py index 584b1039d11..702a3c17454 100644 --- a/tests/unit/single_controller/test_utils.py +++ b/tests/unit/single_controller/test_utils.py @@ -139,12 +139,14 @@ def test_reward_and_advantages_and_tokens(self) -> None: rewards=[torch.tensor([1.0, 3.0])], masked_advantages=[torch.tensor([-1.0, 0.0, 2.0])], sequence_lengths=[4, 6], + num_mask_sample_filtered=[1, 2], ) assert out["reward"] == pytest.approx(2.0) assert out["advantages/mean"] == pytest.approx(1.0 / 3.0) assert out["advantages/max"] == pytest.approx(2.0) assert out["advantages/min"] == pytest.approx(-1.0) assert out["total_num_tokens"] == pytest.approx(10.0) + assert out["num_mask_sample_filtered"] == pytest.approx(3.0) def test_empty_advantages_tensor_yields_zeros(self) -> None: out = reduce_advantage_pump_metrics(