From cb4733ce304c8fd4a0efb8f1d42d65b57190c248 Mon Sep 17 00:00:00 2001 From: Ehsan Hosseini Asl Date: Mon, 24 Aug 2026 13:07:47 -0700 Subject: [PATCH] feat(length-reward): support global reasoning profile bands - clamp profile-band multipliers at the configured floor after the upper bound - derive reasoning length from generated assistant token IDs without prompt history - support global profile bands with row and agent overrides - cover async configuration propagation and profile-band behavior Signed-off-by: Ehsan Hosseini Asl --- nemo_rl/utils/length_adjustments.md | 32 +++- nemo_rl/utils/length_adjustments.py | 120 +++++++++++- tests/unit/algorithms/test_async_utils.py | 53 ++++++ tests/unit/utils/test_length_adjustments.py | 196 ++++++++++++++++++++ 4 files changed, 391 insertions(+), 10 deletions(-) create mode 100644 tests/unit/utils/test_length_adjustments.py diff --git a/nemo_rl/utils/length_adjustments.md b/nemo_rl/utils/length_adjustments.md index 35dd936a8dc..987fefb23a0 100644 --- a/nemo_rl/utils/length_adjustments.md +++ b/nemo_rl/utils/length_adjustments.md @@ -49,6 +49,8 @@ grpo: profile_band_total: false profile_band_reasoning: false profile_band_answer: false + reasoning_end_token_id: null + profile_band: null agent_overrides: math_with_judge_simple_agent: enabled: true @@ -117,6 +119,8 @@ grpo: | `profile_band_total` | Enables per-prompt `{a,b,f}` multiplier on total length for correct rollouts. | | `profile_band_reasoning` | Enables per-prompt `{a,b,f}` multiplier on reasoning length for correct rollouts. | | `profile_band_answer` | Enables per-prompt `{a,b,f}` multiplier on answer length for correct rollouts. | +| `reasoning_end_token_id` | Optional model-specific reasoning-end token used to split newly generated assistant token IDs exactly. | +| `profile_band` | Optional global `{total,reasoning,answer}` band used when a dataset row has no `profile_band` metadata. | | `group_length_penalty_profile_gate` | Gates group-relative length coefficients using a per-prompt `profile_band` threshold. | | `group_length_penalty_profile_gate_channel` | Selects which profile-band channel to gate on: `reasoning`, `answer`, or `total`. | | `group_length_penalty_profile_gate_field` | Selects which field from the chosen profile-band channel to use as the gate threshold, usually `a`. | @@ -415,7 +419,8 @@ Config keys: - `profile_band_reasoning` - `profile_band_answer` -This uses per-row `profile_band` metadata with channel-specific `{a, b, f}` values: +This uses channel-specific `{a, b, f}` values. Per-row `profile_band` metadata +takes precedence when present: ```json { @@ -427,15 +432,34 @@ This uses per-row `profile_band` metadata with channel-specific `{a, b, f}` valu } ``` +When rows do not carry that metadata, configure a global fallback under +`grpo.length_bonus.default`. Agent overrides may replace the fallback: + +```yaml +grpo: + length_bonus: + default: + enabled: true + profile_band_reasoning: true + reasoning_end_token_id: 13 + profile_band: + reasoning: {a: 1024, b: 4096, f: 0.95} +``` + +When `reasoning_end_token_id` is set, reasoning and answer lengths are computed +from newly generated assistant token IDs only. The original prompt prefix is +excluded, and the first matching end token in each generated assistant turn is +the boundary. This supports models such as Nemotron Omni whose prompt opens the +reasoning block and whose generation emits only the closing token. + For an enabled channel, the multiplier is: ```text length <= a: multiplier = 1 -length > a: multiplier = max(0, 1 - (length - a) / (b - a) * (1 - f)) +a < length < b: multiplier = 1 - (length - a) / (b - a) * (1 - f) +length >= b: multiplier = f ``` -So the multiplier is `f` at `b`, then the same slope continues past `b` until clamped at `0`. - Profile-band multipliers are applied only to rollouts whose original environment reward is positive. diff --git a/nemo_rl/utils/length_adjustments.py b/nemo_rl/utils/length_adjustments.py index ff68335e07f..1644bc485a3 100644 --- a/nemo_rl/utils/length_adjustments.py +++ b/nemo_rl/utils/length_adjustments.py @@ -33,8 +33,11 @@ import logging import statistics +from numbers import Real from typing import Any +import torch + logger = logging.getLogger(__name__) # MAD/median floor for zMAD (fixed; matches ``flag_reasoning_length_outliers`` default). @@ -65,6 +68,8 @@ "profile_band_total", "profile_band_reasoning", "profile_band_answer", + "profile_band", + "reasoning_end_token_id", "group_length_penalty_profile_gate", "group_length_penalty_profile_gate_channel", "group_length_penalty_profile_gate_field", @@ -87,6 +92,10 @@ "group_length_penalty_profile_gate_field", }) +_MAPPING_PARAM_KEYS = frozenset({"profile_band"}) + +_OPTIONAL_INT_PARAM_KEYS = frozenset({"reasoning_end_token_id"}) + _GDPO_LENGTH_FEATURE_PARAM_KEYS = frozenset({ "reasoning_bonus", "answer_bonus", @@ -176,6 +185,62 @@ def _extract_reasoning_and_answer_text(result: dict[str, Any]) -> tuple[str, str return reasoning_text, answer_text +def _generated_assistant_token_lengths( + result: dict[str, Any], reasoning_end_token_id: int +) -> tuple[int, int]: + """Return reasoning and answer token lengths for newly generated turns. + + NeMo Gym returns the full conversation in ``message_log`` and the original + prompt as its prefix in ``input_message_log``. Only messages after that + prefix are generated by the current rollout. Each generated assistant turn + is split at its first reasoning-end token; an unfinished turn is counted + entirely as reasoning. + """ + input_message_count = len(result.get("input_message_log", [])) + generated_messages = result.get("message_log", [])[input_message_count:] + + reasoning_length = 0 + answer_length = 0 + for message in generated_messages: + if message.get("role") != "assistant": + continue + token_ids = message.get("token_ids") + if token_ids is None: + raise ValueError( + "Generated assistant messages must include token_ids when " + "reasoning_end_token_id is configured" + ) + if isinstance(token_ids, torch.Tensor): + if token_ids.ndim != 1: + raise ValueError( + "Generated assistant token_ids must be one-dimensional when " + "reasoning_end_token_id is configured, got " + f"shape={tuple(token_ids.shape)}" + ) + token_id_list = token_ids.tolist() + elif isinstance(token_ids, list) and all( + isinstance(token_id, int) and not isinstance(token_id, bool) + for token_id in token_ids + ): + token_id_list = token_ids + else: + raise TypeError( + "Generated assistant token_ids must be a one-dimensional integer " + "sequence when " + "reasoning_end_token_id is configured" + ) + + try: + split_idx = token_id_list.index(reasoning_end_token_id) + except ValueError: + reasoning_length += len(token_id_list) + else: + reasoning_length += split_idx + answer_length += len(token_id_list) - split_idx - 1 + + return reasoning_length, answer_length + + def _extract_gdpo_length_feature_params(feature_cfg: Any) -> dict[str, Any]: if not isinstance(feature_cfg, dict): return {} @@ -297,6 +362,19 @@ def apply_group_length_adjustments( defaults[k] = default_cfg.get(k, True) elif k == "group_length_penalty_profile_gate_positive_only": defaults[k] = default_cfg.get(k, True) + elif k in _MAPPING_PARAM_KEYS: + defaults[k] = default_cfg.get(k) + elif k in _OPTIONAL_INT_PARAM_KEYS: + value = default_cfg.get(k) + if value is not None and ( + isinstance(value, bool) + or not isinstance(value, int) + or value < 0 + ): + raise ValueError( + f"{k} must be a non-negative integer or null, got {value!r}" + ) + defaults[k] = value elif k in _BOOL_PARAM_KEYS: defaults[k] = default_cfg.get(k, False) elif k == "profiled_length_min_samples": @@ -350,11 +428,19 @@ def apply_group_length_adjustments( group_lt = params.pop("length_type", "tokens") use_tokens = group_lt == "tokens" + reasoning_end_token_id = params.pop("reasoning_end_token_id", None) + configured_profile_band = params.pop("profile_band", None) for k in range(group_size): idx = g + k r_text, a_text = texts[idx] - if use_tokens and tokenizer is not None: + if use_tokens and reasoning_end_token_id is not None: + reasoning_lengths[idx], answer_lengths[idx] = ( + _generated_assistant_token_lengths( + results[idx], reasoning_end_token_id + ) + ) + elif use_tokens and tokenizer is not None: reasoning_lengths[idx] = len(tokenizer.encode(r_text, add_special_tokens=False)) if r_text else 0 answer_lengths[idx] = len(tokenizer.encode(a_text, add_special_tokens=False)) if a_text else 0 else: @@ -366,8 +452,9 @@ def apply_group_length_adjustments( group_total = [r + a for r, a in zip(group_reasoning, group_answer)] total_lengths[g : g + group_size] = group_total[:group_size] group_rewards = original_rewards[g : g + num_gens] + profile_band = results[g].get("profile_band") or configured_profile_band gate_info = _group_length_profile_gate_info( - band=results[g].get("profile_band"), + band=profile_band, params=params, rewards=group_rewards[:group_size], reasoning_lengths=group_reasoning[:group_size], @@ -701,7 +788,7 @@ def _apply_profile_band_multipliers( use_ans = bool(params.get("profile_band_answer", False)) if not (use_total or use_rsn or use_ans): continue - band = results[g].get("profile_band") + band = results[g].get("profile_band") or params.get("profile_band") if not band: continue ch_total = band.get("total") if use_total else None @@ -771,18 +858,25 @@ def _band_multiplier(rl: int, ch: dict[str, Any] | None) -> float: Otherwise: rl <= a -> 1.0 rl == b -> f - rl > a -> same linear slope continues past b, floored at 0.0 + rl >= b -> f """ if not ch: return 1.0 a = ch.get("a") b = ch.get("b") f = ch.get("f") - if a is None or b is None or f is None or b <= a: + if any( + isinstance(value, bool) or not isinstance(value, Real) + for value in (a, b, f) + ): + return 1.0 + if b <= a or not 0.0 <= float(f) <= 1.0: return 1.0 if rl <= a: return 1.0 - return max(0.0, 1.0 - (rl - a) / (b - a) * (1.0 - float(f))) + if rl >= b: + return float(f) + return 1.0 - (rl - a) / (b - a) * (1.0 - float(f)) def _group_length_profile_gate_info( @@ -894,6 +988,20 @@ def _resolve_agent_params( merged[key] = overrides[key] elif key in _BOOL_PARAM_KEYS: merged[key] = bool(overrides[key]) + elif key in _MAPPING_PARAM_KEYS: + value = overrides[key] + merged[key] = None if value is None else dict(value) + elif key in _OPTIONAL_INT_PARAM_KEYS: + value = overrides[key] + if value is not None and ( + isinstance(value, bool) + or not isinstance(value, int) + or value < 0 + ): + raise ValueError( + f"{key} must be a non-negative integer or null, got {value!r}" + ) + merged[key] = value else: merged[key] = float(overrides[key]) return merged diff --git a/tests/unit/algorithms/test_async_utils.py b/tests/unit/algorithms/test_async_utils.py index cacfea384ef..ed2a83a3630 100644 --- a/tests/unit/algorithms/test_async_utils.py +++ b/tests/unit/algorithms/test_async_utils.py @@ -1234,6 +1234,59 @@ async def test_drain_payload_metrics_returns_collector_interval(self, monkeypatc assert metrics["payload_bytes/nemo_gym_return/logical_media"] == 150 assert metrics["payload_ratio/nemo_gym_return/physical_to_logical"] == 0.2 + def test_nemo_gym_forwards_length_adjustment_config(self, monkeypatch): + collector = self.create_local_collector() + collector.master_config.grpo.length_bonus = { + "default": { + "enabled": True, + "profile_band_reasoning": True, + } + } + collector.master_config.policy["generation"] = { + "stop_token_ids": [1], + "stop_strings": ["stop"], + } + repeated_batch = BatchedDataDict( + { + "extra_env_info": [ + {"_ng_task_index": 7}, + {"_ng_task_index": 7}, + ], + "loss_multiplier": torch.ones(2), + } + ) + captured_kwargs = {} + + async def fake_rollouts(**kwargs): + captured_kwargs.update(kwargs) + yield SimpleNamespace( + task_index=7, + final_batch=BatchedDataDict({"loss_multiplier": torch.ones(2)}), + rollout_metrics={}, + ) + + import nemo_rl.experience.rollouts as rollouts_mod + + monkeypatch.setattr(rollouts_mod, "run_async_nemo_gym_rollout", fake_rollouts) + + async def collect_groups(): + return [ + group + async for group in collector._iter_rollout_groups( + repeated_batch=repeated_batch, + num_generations=2, + use_nemo_gym=True, + task_index_to_group_index={7: 0}, + ) + ] + + groups = asyncio.run(collect_groups()) + + assert len(groups) == 1 + assert captured_kwargs["length_adjustment_config"] == ( + collector.master_config.grpo.model_dump() + ) + def test_collection_loop_marks_errored_on_crash(self): """A crash sets errored (not data_exhausted) so driver guards fail fast.""" collector = self.create_local_collector() diff --git a/tests/unit/utils/test_length_adjustments.py b/tests/unit/utils/test_length_adjustments.py new file mode 100644 index 00000000000..fb6ab2b71af --- /dev/null +++ b/tests/unit/utils/test_length_adjustments.py @@ -0,0 +1,196 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from nemo_rl.utils.length_adjustments import ( + _band_multiplier, + _generated_assistant_token_lengths, + apply_group_length_adjustments, +) + + +def _result( + *, + reward: float = 1.0, + agent_name: str = "video_agent", + prompt_history: list[dict] | None = None, + generated_token_ids: list[int] | None = None, + profile_band: dict | None = None, +) -> dict: + prompt_history = prompt_history or [ + {"role": "user", "token_ids": torch.tensor([90])} + ] + generated_token_ids = generated_token_ids or [1, 2, 13, 3, 4] + return { + "agent_ref": {"name": agent_name}, + "input_message_log": prompt_history, + "message_log": [ + *prompt_history, + { + "role": "assistant", + "token_ids": torch.tensor(generated_token_ids), + }, + ], + "profile_band": profile_band, + "full_result": {"reward": reward, "response": {"output": []}}, + } + + +def _config(default: dict, agent_overrides: dict | None = None) -> dict: + length_bonus = {"default": {"enabled": True, **default}} + if agent_overrides is not None: + length_bonus["agent_overrides"] = agent_overrides + return { + "grpo": { + "num_generations_per_prompt": 1, + "length_bonus": length_bonus, + } + } + + +@pytest.mark.parametrize( + ("length", "expected"), + [ + (1024, 1.0), + (2560, 0.975), + (4096, 0.95), + (8192, 0.95), + ], +) +def test_profile_band_multiplier_stays_at_floor_after_b(length, expected): + band = {"a": 1024, "b": 4096, "f": 0.95} + + assert _band_multiplier(length, band) == pytest.approx(expected) + + +def test_generated_reasoning_excludes_prompt_history_and_splits_omni_end_token(): + prompt_history = [ + {"role": "user", "token_ids": torch.tensor([90])}, + {"role": "assistant", "token_ids": torch.tensor([7, 8, 9, 13])}, + {"role": "user", "token_ids": torch.tensor([91])}, + ] + result = _result( + prompt_history=prompt_history, + generated_token_ids=[1, 2, 13, 3, 4, 5], + ) + + assert _generated_assistant_token_lengths(result, 13) == (2, 3) + + +def test_generated_reasoning_rejects_non_vector_token_ids(): + result = _result() + result["message_log"][-1]["token_ids"] = torch.tensor([[1, 2, 13]]) + + with pytest.raises(ValueError, match="one-dimensional"): + _generated_assistant_token_lengths(result, 13) + + +@pytest.mark.parametrize("reasoning_end_token_id", [True, -1, "13"]) +def test_reasoning_end_token_id_must_be_non_negative_integer( + reasoning_end_token_id, +): + config = _config( + { + "profile_band_reasoning": True, + "reasoning_end_token_id": reasoning_end_token_id, + "profile_band": { + "reasoning": {"a": 1, "b": 2, "f": 0.5}, + }, + } + ) + + with pytest.raises(ValueError, match="reasoning_end_token_id"): + apply_group_length_adjustments([], config) + + +def test_reasoning_profile_band_uses_global_config_without_row_metadata(): + result = _result(generated_token_ids=[1, 2, 13, 3, 4, 5]) + config = _config( + { + "profile_band_reasoning": True, + "reasoning_end_token_id": 13, + "profile_band": { + "reasoning": {"a": 1, "b": 2, "f": 0.5}, + }, + } + ) + + apply_group_length_adjustments([result], config) + + assert result["full_result"]["reward"] == pytest.approx(0.5) + feature = result["full_result"]["gdpo_reward_features"]["profile_band_reasoning"] + assert feature["multiplier"] == pytest.approx(0.5) + + +def test_row_profile_band_takes_precedence_over_global_config(): + result = _result( + generated_token_ids=[1, 2, 13, 3, 4, 5], + profile_band={"reasoning": {"a": 10, "b": 20, "f": 0.1}}, + ) + config = _config( + { + "profile_band_reasoning": True, + "reasoning_end_token_id": 13, + "profile_band": { + "reasoning": {"a": 1, "b": 2, "f": 0.5}, + }, + } + ) + + apply_group_length_adjustments([result], config) + + assert result["full_result"]["reward"] == pytest.approx(1.0) + + +def test_agent_override_can_supply_global_profile_band(): + result = _result(generated_token_ids=[1, 2, 13]) + config = _config( + { + "profile_band_reasoning": True, + "reasoning_end_token_id": 13, + "profile_band": { + "reasoning": {"a": 1, "b": 2, "f": 0.5}, + }, + }, + agent_overrides={ + "video_agent": { + "profile_band": { + "reasoning": {"a": 1, "b": 2, "f": 0.8}, + } + } + }, + ) + + apply_group_length_adjustments([result], config) + + assert result["full_result"]["reward"] == pytest.approx(0.8) + + +def test_profile_band_does_not_change_incorrect_reward(): + result = _result(reward=0.0, generated_token_ids=[1, 2, 13]) + config = _config( + { + "profile_band_reasoning": True, + "reasoning_end_token_id": 13, + "profile_band": { + "reasoning": {"a": 1, "b": 2, "f": 0.5}, + }, + } + ) + + apply_group_length_adjustments([result], config) + + assert result["full_result"]["reward"] == pytest.approx(0.0)