From 788736bcbf32a07fe9479650a9152acf38e1adf0 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 20 Aug 2026 12:27:11 -0700 Subject: [PATCH 1/4] feat(grpo): support message penalties in single controller Signed-off-by: Anish Mahishi --- ...po_math_1B_megatron_single_controller.yaml | 4 ++ ...po_math_1B_megatron_single_controller.yaml | 4 ++ nemo_rl/algorithms/ppo.py | 6 ++ nemo_rl/algorithms/single_controller.py | 32 +++++++++ .../single_controller_utils/config.py | 18 +++++ .../single_controller_utils/utils.py | 68 +++++++++++++++++++ nemo_rl/data_plane/column_io.py | 9 ++- nemo_rl/data_plane/schema.py | 4 ++ nemo_rl/experience/payload.py | 45 +++++++++++- 9 files changed, 187 insertions(+), 3 deletions(-) diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index 2b50750038e..c07838bcd75 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -7,6 +7,10 @@ grpo: # SC reads async_rl instead, and the entrypoint raises if this block is set. async_grpo: null val_period: 0 + # Optional NeMo-Gym message-level penalties. Values overwrite computed + # advantages on tokens belonging to the corresponding assistant message. + invalid_tool_call_advantage: null + malformed_thinking_advantage: null async_rl: sampler: diff --git a/examples/configs/ppo_math_1B_megatron_single_controller.yaml b/examples/configs/ppo_math_1B_megatron_single_controller.yaml index f94bcc7a39e..40ff951e6f1 100644 --- a/examples/configs/ppo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/ppo_math_1B_megatron_single_controller.yaml @@ -15,6 +15,10 @@ ppo: enabled: false reward_scaling: enabled: false + # Optional NeMo-Gym message-level penalties. Values overwrite computed + # advantages on tokens belonging to the corresponding assistant message. + invalid_tool_call_advantage: null + malformed_thinking_advantage: null async_rl: sampler: diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index 14de4c6797f..ad1afeea4f1 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -205,6 +205,12 @@ class PPOConfig(BaseModel, extra="allow"): # Nullable sequence-level multiplicative probability-error threshold. # None logs metrics without masking; values above the threshold are excluded. seq_logprob_error_threshold: float | None = None + # Advantage value assigned to invalid-tool-call tokens; None disables it. + invalid_tool_call_advantage: float | None = None + # Advantage value assigned to malformed-thinking tokens; None disables it. + malformed_thinking_advantage: float | None = None + + # Asynchronous PPO uses a replay buffer with non-colocated generation. # Legacy async config block; SC reads its async knobs from `async_rl` instead. async_ppo: AsyncPPOConfig | None = Field(default_factory=AsyncPPOConfig) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 405c9ec436a..1b2a6471ead 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -201,6 +201,10 @@ def __init__( self._is_ppo: bool = is_ppo_run(master_config) # GRPO has no epoch knob: it makes one optimizer step per RL step. self._ppo_epochs: int = self._algo_cfg.ppo_epochs if self._is_ppo else 1 + self._message_level_advantage_penalties_enabled = ( + self._algo_cfg.invalid_tool_call_advantage is not None + or self._algo_cfg.malformed_thinking_advantage is not None + ) self._policy_logprobs_required = not ( master_config.loss_fn.force_on_policy_ratio @@ -3265,6 +3269,27 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: if self._is_ppo: returns = torch.zeros_like(mask) + if self._message_level_advantage_penalties_enabled: + # Sequence-error filtering and the pre-existing sample mask remain + # authoritative: a message penalty must not make a filtered token + # trainable again. + valid_tokens = mask.bool() + advantages = apply_message_level_advantage_penalties( + advantages, + invalid_tool_call_mask=( + tensor_field(data, adv_cfg.invalid_tool_call_mask_field).bool() + & valid_tokens + ), + malformed_thinking_mask=( + tensor_field(data, adv_cfg.malformed_thinking_mask_field).bool() + & valid_tokens + ), + invalid_tool_call_advantage=self._algo_cfg.invalid_tool_call_advantage, + malformed_thinking_advantage=( + self._algo_cfg.malformed_thinking_advantage + ), + ) + response_advantages = torch.masked_select(advantages, mask.bool()) self._step_log_dict["rewards"].append(rewards.detach().cpu()) self._step_log_dict["masked_advantages"].append( @@ -3306,6 +3331,13 @@ def _advantage_input_fields(self) -> list[str]: adv_cfg.sample_mask_field, *adv_cfg.repeated_batch_fields, ] + if self._message_level_advantage_penalties_enabled: + fields.extend( + [ + adv_cfg.invalid_tool_call_mask_field, + adv_cfg.malformed_thinking_mask_field, + ] + ) if self._policy_logprobs_required: fields.append(adv_cfg.policy_logprobs_field) if self._policy_logprobs_required: diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 03b2998a0eb..fcda24b8a48 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -42,7 +42,12 @@ from nemo_rl.algorithms.ppo import PPOConfig from nemo_rl.data import DataConfig from nemo_rl.data_plane.interfaces import DataPlaneConfig +from nemo_rl.data_plane.schema import ( + INVALID_TOOL_CALL_MASK, + MALFORMED_THINKING_MASK, +) from nemo_rl.distributed.virtual_cluster import ClusterConfig +from nemo_rl.environments.nemo_gym import should_use_nemo_gym from nemo_rl.models.policy import PolicyConfig from nemo_rl.models.value import ValueConfig from nemo_rl.utils.checkpoint import CheckpointingConfig @@ -1024,6 +1029,17 @@ def validate_single_controller_config(master_config: MasterConfig) -> None: # configs can omit it. Only apply rollout-path validation when it is present. env_config = getattr(master_config, "env", None) + penalties_enabled = ( + algo_cfg.invalid_tool_call_advantage is not None + or algo_cfg.malformed_thinking_advantage is not None + ) + if penalties_enabled and not should_use_nemo_gym(master_config): + raise ValueError( + "invalid_tool_call_advantage and malformed_thinking_advantage on the " + "active algorithm block require the NeMo-Gym rollout path " + "(env.should_use_nemo_gym=true) on SingleController." + ) + opd_enabled = opd_module.is_opd_enabled(master_config) if opd_enabled and is_ppo_run(master_config): raise ValueError( @@ -1110,6 +1126,8 @@ class AdvantageConfig: reward_field: str = "total_reward" token_mask_field: str = "token_mask" sample_mask_field: str = "sample_mask" + invalid_tool_call_mask_field: str = INVALID_TOOL_CALL_MASK + malformed_thinking_mask_field: str = MALFORMED_THINKING_MASK 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 62e596184a6..1b4f6dab4ee 100644 --- a/nemo_rl/algorithms/single_controller_utils/utils.py +++ b/nemo_rl/algorithms/single_controller_utils/utils.py @@ -198,6 +198,74 @@ def reduce_range( return reduced +def apply_message_level_advantage_penalties( + advantages: torch.Tensor, + *, + invalid_tool_call_mask: torch.Tensor, + malformed_thinking_mask: torch.Tensor, + invalid_tool_call_advantage: float | None, + malformed_thinking_advantage: float | None, +) -> torch.Tensor: + """Overwrite flagged token advantages while leaving valid tokens unchanged. + + Invalid-tool-call penalties take precedence when both masks select the same + token, matching the legacy GRPO message-level implementation. + + Args: + advantages: Per-token advantages of shape (batch, seq). + invalid_tool_call_mask: Bool mask of the same shape as ``advantages``; + True at tokens produced by an invalid tool call. + malformed_thinking_mask: Bool mask of the same shape as ``advantages``; + True at tokens produced by malformed thinking. + invalid_tool_call_advantage: Value to overwrite flagged tokens with, or + ``None`` to leave the invalid-tool-call branch disabled. + malformed_thinking_advantage: Value to overwrite flagged tokens with, or + ``None`` to leave the malformed-thinking branch disabled. + + Returns: + New tensor with penalties applied. The original ``advantages`` object is + returned unchanged when both advantages are ``None``. + + Raises: + ValueError: If either mask shape does not match ``advantages``. + """ + if invalid_tool_call_mask.shape != advantages.shape: + raise ValueError( + "invalid_tool_call_mask shape " + f"{tuple(invalid_tool_call_mask.shape)} does not match advantages " + f"{tuple(advantages.shape)}" + ) + if malformed_thinking_mask.shape != advantages.shape: + raise ValueError( + "malformed_thinking_mask shape " + f"{tuple(malformed_thinking_mask.shape)} does not match advantages " + f"{tuple(advantages.shape)}" + ) + + result = advantages + if malformed_thinking_advantage is not None: + result = torch.where( + malformed_thinking_mask.bool(), + torch.as_tensor( + malformed_thinking_advantage, + dtype=advantages.dtype, + device=advantages.device, + ), + result, + ) + if invalid_tool_call_advantage is not None: + result = torch.where( + invalid_tool_call_mask.bool(), + torch.as_tensor( + invalid_tool_call_advantage, + dtype=advantages.dtype, + device=advantages.device, + ), + result, + ) + return result + + def tensor_field(data: TensorDict, field_name: str) -> torch.Tensor: """Read a tensor column from a TensorDict, depadding if nested. diff --git a/nemo_rl/data_plane/column_io.py b/nemo_rl/data_plane/column_io.py index 9bbdb5616c0..f9c49790fdf 100644 --- a/nemo_rl/data_plane/column_io.py +++ b/nemo_rl/data_plane/column_io.py @@ -37,7 +37,12 @@ from nemo_rl.data.llm_message_utils import attach_message_log_view from nemo_rl.data_plane.codec import materialize, pack_jagged_fields from nemo_rl.data_plane.interfaces import DataPlaneClient, KVBatchMeta -from nemo_rl.data_plane.schema import GLOBAL_FORWARD_PAD_SEQLEN, Layout +from nemo_rl.data_plane.schema import ( + GLOBAL_FORWARD_PAD_SEQLEN, + INVALID_TOOL_CALL_MASK, + MALFORMED_THINKING_MASK, + Layout, +) from nemo_rl.distributed.batched_data_dict import BatchedDataDict TOKEN_ALIGNED_FIELDS = frozenset( @@ -53,6 +58,8 @@ "token_mask", "sample_mask", "routed_experts", + INVALID_TOOL_CALL_MASK, + MALFORMED_THINKING_MASK, } ) diff --git a/nemo_rl/data_plane/schema.py b/nemo_rl/data_plane/schema.py index 9c1f9d1e2ef..a8059796acf 100644 --- a/nemo_rl/data_plane/schema.py +++ b/nemo_rl/data_plane/schema.py @@ -30,6 +30,10 @@ SAMPLE_MASK = "sample_mask" META_IDX = "meta_idx" +# Token-aligned message-violation fields consumed by SingleController advantages. +INVALID_TOOL_CALL_MASK = "invalid_tool_call_mask" +MALFORMED_THINKING_MASK = "malformed_thinking_mask" + # Tensor fields in the train partition. Rollout writes the input # subset on first put; later stages add prev_logprobs / # reference_policy_logprobs (workers) and advantages (driver). diff --git a/nemo_rl/experience/payload.py b/nemo_rl/experience/payload.py index e10a9026fa0..67edf22d412 100644 --- a/nemo_rl/experience/payload.py +++ b/nemo_rl/experience/payload.py @@ -24,7 +24,11 @@ from nemo_rl.data.interfaces import LLMMessageLogType, VLMMessageLogType from nemo_rl.data_plane.codec import pack_jagged_fields from nemo_rl.data_plane.column_io import TOKEN_ALIGNED_FIELDS -from nemo_rl.data_plane.schema import ROUTED_EXPERTS_FIELD +from nemo_rl.data_plane.schema import ( + INVALID_TOOL_CALL_MASK, + MALFORMED_THINKING_MASK, + ROUTED_EXPERTS_FIELD, +) from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.experience.interfaces import PromptGroupRecord @@ -54,21 +58,53 @@ def _violation_counts( return counts +def _add_message_violation_masks( + message_logs: list[LLMMessageLogType | VLMMessageLogType], +) -> None: + """Attach token-aligned masks for generated assistant violations. + + This must run before the generic message normalizer fills missing + ``generation_logprobs`` on prompt and environment messages, because field + presence distinguishes generated assistant turns. + """ + for message_log in message_logs: + for message in message_log: + token_ids = message["token_ids"] + is_generated_assistant = ( + message["role"] == "assistant" and "generation_logprobs" in message + ) + is_invalid = is_generated_assistant and bool( + message.get("is_invalid_tool_call", False) + ) + is_malformed = is_generated_assistant and bool( + message.get("has_malformed_thinking", False) + ) + message[INVALID_TOOL_CALL_MASK] = torch.full_like( + token_ids, is_invalid, dtype=torch.bool + ) + message[MALFORMED_THINKING_MASK] = torch.full_like( + token_ids, is_malformed, dtype=torch.bool + ) + + def record_to_train_batch( record: PromptGroupRecord, *, pad_value_dict: Mapping[str, int], + 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. + include_message_violation_fields: Whether to tensorize message violation + 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 - routed_experts. + routed experts and message-violation masks. """ # Lazy imports: grpo and llm_message_utils transitively pull # experience.rollouts, so importing at module top risks a cycle. @@ -85,6 +121,8 @@ def record_to_train_batch( message_logs = [c.message_log for c in completions] prompt_token_count = sum(len(m["token_ids"]) for m in record.prompt) + if include_message_violation_fields: + _add_message_violation_masks(message_logs) prompt_lengths = torch.full((n,), prompt_token_count, dtype=torch.long) # Must precede the prompt extraction: it reuses the same message dicts, so @@ -121,6 +159,9 @@ def record_to_train_batch( } if ROUTED_EXPERTS_FIELD in flat: train_data[ROUTED_EXPERTS_FIELD] = flat[ROUTED_EXPERTS_FIELD] + if include_message_violation_fields: + train_data[INVALID_TOOL_CALL_MASK] = flat[INVALID_TOOL_CALL_MASK] + train_data[MALFORMED_THINKING_MASK] = flat[MALFORMED_THINKING_MASK] return BatchedDataDict[Any](train_data) From 971332f99230c1c76ce53a255031946eb22cace2 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 20 Aug 2026 13:21:21 -0700 Subject: [PATCH 2/4] fix(grpo): refine single controller message penalties Signed-off-by: Anish Mahishi --- .../algorithms/async_utils/replay_buffer.py | 8 +- nemo_rl/algorithms/ppo.py | 1 - nemo_rl/algorithms/single_controller.py | 1 + .../single_controller_utils/setup.py | 4 + nemo_rl/data_plane/schema.py | 2 + nemo_rl/experience/payload.py | 7 +- tests/unit/experience/test_payload.py | 77 ++++++++++++++++++- .../_checkpoint_scenarios.py | 1 + .../test_checkpoint_dispatch_races.py | 2 + .../unit/single_controller/test_ppo_setup.py | 10 +++ .../single_controller/test_rollout_pump.py | 1 + .../test_single_controller_actor.py | 29 ++++++- .../test_tq_replay_buffer.py | 9 ++- .../single_controller/test_train_pump_e2e.py | 1 + tests/unit/single_controller/test_utils.py | 70 +++++++++++++++++ 15 files changed, 214 insertions(+), 9 deletions(-) diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 2cf2d09df08..3f4e622d435 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -981,11 +981,13 @@ def __init__( partition_id: str, *, pad_value_dict: Mapping[str, int], + 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._include_message_violation_fields = include_message_violation_fields self._require_routed_experts = require_routed_experts self.meta_list: list[Optional[KVBatchMeta]] = [] self.start_weight_list: list[int] = [] @@ -1090,7 +1092,11 @@ async def commit( "TQReplayBuffer must be bound to the controller data-plane " "checkpoint barrier before committing samples" ) - train_batch = record_to_train_batch(record, pad_value_dict=self._pad_value_dict) + train_batch = record_to_train_batch( + record, + pad_value_dict=self._pad_value_dict, + include_message_violation_fields=self._include_message_violation_fields, + ) sample_ids, fields, tags = pack_payload( train_batch, weight_version=start_weight_version, diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index ad1afeea4f1..a7902e208b9 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -210,7 +210,6 @@ class PPOConfig(BaseModel, extra="allow"): # Advantage value assigned to malformed-thinking tokens; None disables it. malformed_thinking_advantage: float | None = None - # Asynchronous PPO uses a replay buffer with non-colocated generation. # Legacy async config block; SC reads its async knobs from `async_rl` instead. async_ppo: AsyncPPOConfig | None = Field(default_factory=AsyncPPOConfig) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 1b2a6471ead..23f8ad00c57 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -93,6 +93,7 @@ from nemo_rl.algorithms.single_controller_utils.setup import SingleControllerActorArgs from nemo_rl.algorithms.single_controller_utils.utils import ( aggregate_step_metrics, + apply_message_level_advantage_penalties, fields_for_put, reduce_advantage_pump_metrics, squeeze_trailing_unit_dim, diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 4f77ba4c802..9f595486375 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -1393,6 +1393,10 @@ def _build_generation_then_trainer( dp_client, partition_id=partition_id, pad_value_dict={"token_ids": pad_id, "input_ids": pad_id}, + include_message_violation_fields=( + algo_cfg.invalid_tool_call_advantage is not None + or algo_cfg.malformed_thinking_advantage is not None + ), require_routed_experts=router_replay_enabled(policy_config), ) rollout_manager = RolloutManager( diff --git a/nemo_rl/data_plane/schema.py b/nemo_rl/data_plane/schema.py index a8059796acf..4a7dbd8d5e0 100644 --- a/nemo_rl/data_plane/schema.py +++ b/nemo_rl/data_plane/schema.py @@ -60,6 +60,8 @@ "values", "returns", "teacher_reference_logprobs", + INVALID_TOOL_CALL_MASK, + MALFORMED_THINKING_MASK, ) # Subset fetched by logprob / ref-logprob workers. diff --git a/nemo_rl/experience/payload.py b/nemo_rl/experience/payload.py index 67edf22d412..b3b34432316 100644 --- a/nemo_rl/experience/payload.py +++ b/nemo_rl/experience/payload.py @@ -15,7 +15,7 @@ """Producer-side payload helpers for the async-RL TQ path.""" from collections.abc import Mapping -from typing import Any +from typing import Any, cast import numpy as np import torch @@ -69,7 +69,7 @@ def _add_message_violation_masks( """ for message_log in message_logs: for message in message_log: - token_ids = message["token_ids"] + token_ids = cast(torch.Tensor, message["token_ids"]) is_generated_assistant = ( message["role"] == "assistant" and "generation_logprobs" in message ) @@ -120,6 +120,7 @@ def record_to_train_batch( assert n > 0, "PromptGroupRecord has no completions" message_logs = [c.message_log for c in completions] + violation_counts = [_violation_counts(message_log) for message_log in message_logs] prompt_token_count = sum(len(m["token_ids"]) for m in record.prompt) if include_message_violation_fields: _add_message_violation_masks(message_logs) @@ -155,7 +156,7 @@ def record_to_train_batch( "sample_mask": sample_mask, "prompt_ids_for_adv": prompt_flat["token_ids"], "total_reward": total_reward, - _VIOLATION_COUNTS_KEY: [_violation_counts(ml) for ml in message_logs], + _VIOLATION_COUNTS_KEY: violation_counts, } if ROUTED_EXPERTS_FIELD in flat: train_data[ROUTED_EXPERTS_FIELD] = flat[ROUTED_EXPERTS_FIELD] diff --git a/tests/unit/experience/test_payload.py b/tests/unit/experience/test_payload.py index 5029d3ebb63..c1c32016bfa 100644 --- a/tests/unit/experience/test_payload.py +++ b/tests/unit/experience/test_payload.py @@ -16,6 +16,10 @@ import torch +from nemo_rl.data_plane.schema import ( + INVALID_TOOL_CALL_MASK, + MALFORMED_THINKING_MASK, +) from nemo_rl.experience.interfaces import Completion, PromptGroupRecord from nemo_rl.experience.payload import pack_payload, record_to_train_batch @@ -103,6 +107,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}, + include_message_violation_fields=False, ) expected_routes = [ @@ -144,14 +149,82 @@ def test_record_to_train_batch_preserves_routed_experts_in_tq_payload() -> None: ] +def test_record_to_train_batch_preserves_message_violation_masks() -> None: + invalid = _completion(route_start=10, reward=1.0) + invalid.message_log[1]["is_invalid_tool_call"] = True + + malformed = _completion( + route_start=30, + reward=2.0, + env_token_ids=(30, 31), + ) + malformed.message_log[1]["has_malformed_thinking"] = True + + train_batch = record_to_train_batch( + _record([invalid, malformed]), + pad_value_dict={"token_ids": 0, "input_ids": 0}, + include_message_violation_fields=True, + ) + + assert train_batch[INVALID_TOOL_CALL_MASK].dtype == torch.bool + assert train_batch[MALFORMED_THINKING_MASK].dtype == torch.bool + assert train_batch[INVALID_TOOL_CALL_MASK][0, :5].tolist() == [ + False, + False, + True, + True, + False, + ] + assert not train_batch[MALFORMED_THINKING_MASK][0, :5].any() + assert not train_batch[INVALID_TOOL_CALL_MASK][1, :6].any() + assert train_batch[MALFORMED_THINKING_MASK][1, :6].tolist() == [ + False, + False, + True, + True, + False, + False, + ] + + _, fields, tags = pack_payload( + train_batch, + weight_version=3, + group_id="group", + prompt_idx=17, + ) + invalid_rows = list(fields[INVALID_TOOL_CALL_MASK].unbind()) + malformed_rows = list(fields[MALFORMED_THINKING_MASK].unbind()) + assert invalid_rows[0].tolist() == [False, False, True, True, False] + assert malformed_rows[1].tolist() == [False, False, True, True, False, False] + assert tags[0]["num_invalid_tool_calls"] == 1 + assert tags[1]["num_malformed_thinking"] == 1 + + +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}, + include_message_violation_fields=True, + ) + + assert not train_batch[INVALID_TOOL_CALL_MASK].any() + assert not train_batch[MALFORMED_THINKING_MASK].any() + + def test_record_to_train_batch_omits_routed_experts_when_absent() -> None: - record = _record([_completion(route_start=10, reward=1.0, with_routes=False)]) + completion = _completion(route_start=10, reward=1.0, with_routes=False) + completion.message_log[1]["is_invalid_tool_call"] = True + completion.message_log[1]["has_malformed_thinking"] = True + record = _record([completion]) train_batch = record_to_train_batch( record, pad_value_dict={"token_ids": 0, "input_ids": 0}, + include_message_violation_fields=False, ) assert "routed_experts" not in train_batch + assert INVALID_TOOL_CALL_MASK not in train_batch + assert MALFORMED_THINKING_MASK not in train_batch _, fields, _ = pack_payload( train_batch, @@ -185,6 +258,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}, + include_message_violation_fields=False, ) assert train_batch["input_lengths"].tolist() == [5, 2] @@ -220,6 +294,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}, + include_message_violation_fields=False, ) _, fields, tags = pack_payload( train_batch, diff --git a/tests/unit/single_controller/_checkpoint_scenarios.py b/tests/unit/single_controller/_checkpoint_scenarios.py index 5e657615746..8151ce189be 100644 --- a/tests/unit/single_controller/_checkpoint_scenarios.py +++ b/tests/unit/single_controller/_checkpoint_scenarios.py @@ -230,6 +230,7 @@ def _new_buffer(dp: NoOpDataPlaneClient) -> TQReplayBuffer: dp, partition_id=PARTITION, pad_value_dict={"input_ids": 0}, + include_message_violation_fields=False, require_routed_experts=False, ) buf.set_data_plane_checkpoint_barrier(DataPlaneCheckpointBarrier()) diff --git a/tests/unit/single_controller/test_checkpoint_dispatch_races.py b/tests/unit/single_controller/test_checkpoint_dispatch_races.py index 062b2595d62..eb80541c928 100644 --- a/tests/unit/single_controller/test_checkpoint_dispatch_races.py +++ b/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}, + include_message_violation_fields=False, require_routed_experts=False, ) group_id = buffer.reserve( @@ -782,6 +783,7 @@ async def exercise() -> None: dp_client, partition_id="rollout_data", pad_value_dict={"input_ids": 0}, + include_message_violation_fields=False, require_routed_experts=False, ) group_id = buffer.reserve( diff --git a/tests/unit/single_controller/test_ppo_setup.py b/tests/unit/single_controller/test_ppo_setup.py index bcc3d791f38..6e1847c0980 100644 --- a/tests/unit/single_controller/test_ppo_setup.py +++ b/tests/unit/single_controller/test_ppo_setup.py @@ -238,6 +238,16 @@ class TestPPOValidation: def test_accepts_a_well_formed_ppo_config(self): validate_single_controller_config(_ppo_master_config()) + def test_rejects_message_penalties_without_nemo_gym(self): + mc = _ppo_master_config() + assert mc.ppo is not None + mc.ppo.invalid_tool_call_advantage = -5.0 + + with pytest.raises( + ValueError, match="active algorithm block require the NeMo-Gym" + ): + validate_single_controller_config(mc) + @pytest.mark.parametrize("missing", ["value", "value_loss_fn"]) def test_rejects_ppo_without_its_critic_blocks(self, missing): mc = _ppo_master_config(**{missing: None}) diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index a57b0b9f13a..b0a7285b0c4 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/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)}, + include_message_violation_fields=False, ) rollout_manager = RolloutManager( tokenizer=tokenizer, diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index d9c7d359b21..d6ec4c7b5d5 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -551,6 +551,12 @@ def test_advantage_stage_applies_seq_logprob_error_mask_before_streaming_train( "sample_mask": torch.ones(batch_size), "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. + "invalid_tool_call_mask": torch.tensor( + [[False] * sequence_length] * 2 + [[True] * sequence_length] * 2 + ), + "malformed_thinking_mask": torch.zeros(batch_size, sequence_length), }, batch_size=[batch_size], ) @@ -567,9 +573,14 @@ 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=SimpleNamespace( + seq_logprob_error_threshold=2.0, + invalid_tool_call_advantage=-5.0, + malformed_thinking_advantage=None, + ) ) ctrl._algo_cfg = ctrl._master_config.grpo + ctrl._message_level_advantage_penalties_enabled = True ctrl._step_log_dict = { "rewards": [], "masked_advantages": [], @@ -589,8 +600,17 @@ def test_advantage_stage_applies_seq_logprob_error_mask_before_streaming_train( assert has_valid_training_tokens assert data_plane.selected_fields is not None assert "prev_logprobs" in data_plane.selected_fields + 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. + 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) + ) assert torch.equal( data_plane.written_fields["sample_mask"], torch.tensor([1.0, 1.0, 0.0, 1.0]), @@ -639,6 +659,7 @@ def test_advantage_stage_reports_seq_logprob_metrics_without_masking() -> None: grpo=SimpleNamespace(seq_logprob_error_threshold=None) ) ctrl._algo_cfg = ctrl._master_config.grpo + ctrl._message_level_advantage_penalties_enabled = False ctrl._step_log_dict = { "rewards": [], "masked_advantages": [], @@ -702,6 +723,7 @@ def test_advantage_stage_skips_estimator_when_seq_mask_removes_whole_chunk( grpo=SimpleNamespace(seq_logprob_error_threshold=2.0) ) ctrl._algo_cfg = ctrl._master_config.grpo + ctrl._message_level_advantage_penalties_enabled = False ctrl._step_log_dict = { "rewards": [], "masked_advantages": [], @@ -758,6 +780,7 @@ def test_advantage_stage_skips_preexisting_empty_mask_without_seq_threshold() -> grpo=SimpleNamespace(seq_logprob_error_threshold=None) ) ctrl._algo_cfg = ctrl._master_config.grpo + ctrl._message_level_advantage_penalties_enabled = False ctrl._step_log_dict = { "rewards": [], "masked_advantages": [], @@ -834,6 +857,7 @@ def put_samples(self, sample_ids, partition_id, fields): grpo=SimpleNamespace(seq_logprob_error_threshold=None) ) ctrl._algo_cfg = ctrl._master_config.grpo + ctrl._message_level_advantage_penalties_enabled = False ctrl._step_log_dict = { "rewards": [], "masked_advantages": [], @@ -1058,6 +1082,7 @@ def _train_pump_controller(*, sampler) -> object: checkpointing={"enabled": False, "save_period": 10}, ) ctrl._algo_cfg = ctrl._master_config.grpo + ctrl._message_level_advantage_penalties_enabled = False ctrl._async_cfg = SimpleNamespace( min_groups_for_streaming_train=1, rollout_failure=SimpleNamespace(min_step_batch_fraction=0.9), @@ -1474,6 +1499,7 @@ def _ppo_train_pump_controller( seq_logprob_error_threshold=None, ) ctrl._algo_cfg = ctrl._master_config.ppo + ctrl._message_level_advantage_penalties_enabled = False ctrl._sync_weights = AsyncMock(return_value=0) ctrl._logger = MagicMock() return ctrl, value @@ -1743,6 +1769,7 @@ def compute_advantage(self, *, rewards, mask, **kwargs): ppo=SimpleNamespace(seq_logprob_error_threshold=None) ) ctrl._algo_cfg = ctrl._master_config.ppo + ctrl._message_level_advantage_penalties_enabled = False ctrl._step_log_dict = { "rewards": [], "masked_advantages": [], diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index edfa15a0b48..52cb5dc57bd 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -41,9 +41,12 @@ def _stub_record_to_train_batch( - record: PromptGroupRecord, *, pad_value_dict: Any + record: PromptGroupRecord, + *, + pad_value_dict: Any, + include_message_violation_fields: bool, ) -> BatchedDataDict[Any]: - del record, pad_value_dict + del record, pad_value_dict, include_message_violation_fields return BatchedDataDict[Any]( { "input_ids": torch.ones((_N_GENS, 3), dtype=torch.long), @@ -183,6 +186,7 @@ def _make_buffer( dp, partition_id="rollout_data", pad_value_dict={"token_ids": 0}, + include_message_violation_fields=False, require_routed_experts=require_routed_experts, ) buffer.set_data_plane_checkpoint_barrier( @@ -568,6 +572,7 @@ def test_remove_with_dp_clear_fails_without_bound_checkpoint_barrier(self): dp, partition_id="rollout_data", pad_value_dict={"token_ids": 0}, + include_message_violation_fields=False, ) with pytest.raises(RuntimeError, match="must be bound"): diff --git a/tests/unit/single_controller/test_train_pump_e2e.py b/tests/unit/single_controller/test_train_pump_e2e.py index f489d722e59..68cc4fd8ea9 100644 --- a/tests/unit/single_controller/test_train_pump_e2e.py +++ b/tests/unit/single_controller/test_train_pump_e2e.py @@ -279,6 +279,7 @@ def test_train_pump_drives_mcore_training_step( dp_client, partition_id=_PARTITION_ID, pad_value_dict={"input_ids": int(tokenizer.pad_token_id or 0)}, + include_message_violation_fields=False, ) for step in range(train_steps): for g in range(num_prompts): diff --git a/tests/unit/single_controller/test_utils.py b/tests/unit/single_controller/test_utils.py index 24f15b89b80..584b1039d11 100644 --- a/tests/unit/single_controller/test_utils.py +++ b/tests/unit/single_controller/test_utils.py @@ -24,6 +24,7 @@ from nemo_rl.algorithms.single_controller_utils.utils import ( aggregate_step_metrics, + apply_message_level_advantage_penalties, fields_for_put, reduce_advantage_pump_metrics, squeeze_trailing_unit_dim, @@ -233,6 +234,75 @@ def test_no_assistant_messages_omits_violation_metrics(self) -> None: ) +class TestApplyMessageLevelAdvantagePenalties: + def test_overwrites_only_flagged_tokens(self) -> None: + advantages = torch.tensor([[1.0, 2.0, 3.0, 4.0]]) + result = apply_message_level_advantage_penalties( + advantages, + invalid_tool_call_mask=torch.tensor([[False, True, False, False]]), + malformed_thinking_mask=torch.tensor([[False, False, True, False]]), + invalid_tool_call_advantage=-5.0, + malformed_thinking_advantage=-7.0, + ) + torch.testing.assert_close(result, torch.tensor([[1.0, -5.0, -7.0, 4.0]])) + torch.testing.assert_close(advantages, torch.tensor([[1.0, 2.0, 3.0, 4.0]])) + + def test_invalid_tool_call_takes_precedence_on_overlap(self) -> None: + result = apply_message_level_advantage_penalties( + torch.zeros(1, 2), + invalid_tool_call_mask=torch.tensor([[False, True]]), + malformed_thinking_mask=torch.tensor([[False, True]]), + invalid_tool_call_advantage=-5.0, + malformed_thinking_advantage=-7.0, + ) + torch.testing.assert_close(result, torch.tensor([[0.0, -5.0]])) + + def test_only_invalid_tool_call_penalty_leaves_malformed_untouched( + self, + ) -> None: + result = apply_message_level_advantage_penalties( + torch.tensor([[1.0, 2.0, 3.0, 4.0]]), + invalid_tool_call_mask=torch.tensor([[False, True, False, False]]), + malformed_thinking_mask=torch.tensor([[False, False, True, False]]), + invalid_tool_call_advantage=-5.0, + malformed_thinking_advantage=None, + ) + torch.testing.assert_close(result, torch.tensor([[1.0, -5.0, 3.0, 4.0]])) + + def test_only_malformed_thinking_penalty_leaves_invalid_untouched( + self, + ) -> None: + result = apply_message_level_advantage_penalties( + torch.tensor([[1.0, 2.0, 3.0, 4.0]]), + invalid_tool_call_mask=torch.tensor([[False, True, False, False]]), + malformed_thinking_mask=torch.tensor([[False, False, True, False]]), + invalid_tool_call_advantage=None, + malformed_thinking_advantage=-7.0, + ) + torch.testing.assert_close(result, torch.tensor([[1.0, 2.0, -7.0, 4.0]])) + + def test_disabled_penalties_leave_advantages_unchanged(self) -> None: + advantages = torch.tensor([[1.0, 2.0]]) + result = apply_message_level_advantage_penalties( + advantages, + invalid_tool_call_mask=torch.tensor([[True, False]]), + malformed_thinking_mask=torch.tensor([[False, True]]), + invalid_tool_call_advantage=None, + malformed_thinking_advantage=None, + ) + assert result is advantages + + def test_rejects_misaligned_mask(self) -> None: + with pytest.raises(ValueError, match="invalid_tool_call_mask shape"): + apply_message_level_advantage_penalties( + torch.zeros(1, 2), + invalid_tool_call_mask=torch.zeros(1, 3, dtype=torch.bool), + malformed_thinking_mask=torch.zeros(1, 2, dtype=torch.bool), + invalid_tool_call_advantage=-5.0, + malformed_thinking_advantage=None, + ) + + class TestFieldsForPut: def test_no_sequence_lengths_packs_contiguous(self) -> None: meta = _meta(2, sequence_lengths=None) From dfe2d2dcb5e0833a06d8a4fc2aed64d7b840141d Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Mon, 31 Aug 2026 09:01:47 -0700 Subject: [PATCH 3/4] fix(config): sync PPO reference penalty defaults Signed-off-by: Anish Mahishi --- tests/unit/reference_configs/ppo_math_1B_megatron.yaml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml index 706fde4c3df..9887b85d8b4 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml @@ -56,6 +56,10 @@ ppo: source_max: 1.0 target_min: -1.0 # DAPO: scale rewards to [-1, 1] target_max: 1.0 + # Advantage value assigned to invalid tool call tokens (e.g. -5.0 to penalize). null disables. + invalid_tool_call_advantage: null + # Advantage value assigned to tokens with malformed / tags (e.g. -5.0). null disables. + malformed_thinking_advantage: null seq_logprob_error_threshold: null loss_fn: From 9cd5e2e99a6662be1da469f7c29c5763ee9c726f Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Mon, 31 Aug 2026 11:13:14 -0700 Subject: [PATCH 4/4] fix(sc): update checkpoint converter stub Signed-off-by: Anish Mahishi --- tests/unit/single_controller/_checkpoint_scenarios.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/unit/single_controller/_checkpoint_scenarios.py b/tests/unit/single_controller/_checkpoint_scenarios.py index 8151ce189be..2fbe41696a6 100644 --- a/tests/unit/single_controller/_checkpoint_scenarios.py +++ b/tests/unit/single_controller/_checkpoint_scenarios.py @@ -197,8 +197,13 @@ def _record() -> PromptGroupRecord: ) -def _stub_converter(record: PromptGroupRecord, *, pad_value_dict: Any): - del record, pad_value_dict +def _stub_converter( + record: PromptGroupRecord, + *, + pad_value_dict: Any, + include_message_violation_fields: bool, +): + del record, pad_value_dict, include_message_violation_fields return BatchedDataDict[Any]( { "input_ids": torch.ones((ROLLOUTS_PER_GROUP, 3), dtype=torch.long),