From 9a93d65d79269cebb478ed7e7b2836e7a1bc9e76 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 20 Aug 2026 16:26:02 -0400 Subject: [PATCH 1/2] fix: detect rollouts truncated by the per-turn generation cap hit_max_tokens feeds final_batch["truncated"], which grpo.overlong_filtering uses to drop a sample from the loss, but it only tested whether the conversation had filled the context (vllm_cfg.max_model_len). A rollout can also be stopped by policy.generation.max_new_tokens, which NeMo-Gym agents apply per turn, and a run with a small max_new_tokens never fills the context at all -- so every one of its truncated samples kept the near-zero reward that overlong filtering exists to discard. Check both budgets, and compare with >= rather than ==: the message log is post-processed before it reaches here (reasoning content is re-wrapped in thinking tags, for one), so an exact match can be thrown off by a single token. truncation_rate and natural_termination_rate derive from the same flag, so runs whose generation cap bound below the context will now report higher truncation than before. The earlier numbers were under-counting. Bump the Gym submodule for the matching change: verifiers_agent now honors responses_create_params.max_output_tokens, so max_new_tokens actually reaches the agent. The two belong together -- without the Gym half nothing is capped below the context, and without this half the newly honored cap goes undetected. Signed-off-by: Albert Cui --- 3rdparty/Gym-workspace/Gym | 2 +- nemo_rl/experience/rollouts.py | 75 +++++++++++++++------ tests/unit/experience/test_rollouts.py | 93 ++++++++++++++++++++++++++ 3 files changed, 148 insertions(+), 22 deletions(-) diff --git a/3rdparty/Gym-workspace/Gym b/3rdparty/Gym-workspace/Gym index 74a67f6d104..6230e442921 160000 --- a/3rdparty/Gym-workspace/Gym +++ b/3rdparty/Gym-workspace/Gym @@ -1 +1 @@ -Subproject commit 74a67f6d104ff116862a041473f71726ee1697e5 +Subproject commit 6230e442921c4ed7ef33060c63cb8871e4019dea diff --git a/nemo_rl/experience/rollouts.py b/nemo_rl/experience/rollouts.py index 1a6145160b8..5140fbc4344 100644 --- a/nemo_rl/experience/rollouts.py +++ b/nemo_rl/experience/rollouts.py @@ -2313,6 +2313,52 @@ async def _consume_rollout() -> NemoGymRolloutResult: return asyncio.run(_consume_rollout()) +def _nemo_gym_sample_metrics( + result: dict, + *, + max_total_tokens: int, + max_new_tokens: Optional[int], +) -> dict[str, Any]: + """Per-sample rollout metrics, including whether generation was cut short. + + ``hit_max_tokens`` becomes ``final_batch["truncated"]``, which + ``grpo.overlong_filtering`` uses to drop a sample from the loss. A rollout can be + stopped by either of two budgets, so both are checked: + + * the conversation filled the context (``max_total_tokens``), or + * a single turn reached the per-turn generation cap (``max_new_tokens``). + + Checking only the first misses every rollout that stopped below the context. That is + not a corner case: ``max_new_tokens`` is applied per turn, so a run with a small + ``max_new_tokens`` never fills the context at all and its truncated samples keep the + near-zero reward that overlong filtering exists to discard. + + Both comparisons are ``>=`` rather than ``==``. The message log is post-processed + before it reaches here -- reasoning content is re-wrapped in thinking tags, for one -- + so an exact match can be thrown off by a single token. + """ + message_log = result["message_log"] + assistant_lengths = [ + len(m["token_ids"]) for m in message_log if m["role"] == "assistant" + ] + total_tokens = sum(len(m["token_ids"]) for m in message_log) + max_gen_tokens_per_turn = max(assistant_lengths, default=0) + + hit_max_tokens = total_tokens >= max_total_tokens + if max_new_tokens is not None: + hit_max_tokens = hit_max_tokens or max_gen_tokens_per_turn >= max_new_tokens + + return { + "total_reward": result["full_result"]["reward"], + "assistant_tokens": sum(assistant_lengths), + "total_tokens": total_tokens, + "turn_count": sum(1 for m in message_log if m["role"] == "user"), + "hit_max_tokens": hit_max_tokens, + # max_gen_tokens_per_turn: Diagnostic for long single generations + "max_gen_tokens_per_turn": max_gen_tokens_per_turn, + } + + def _postprocess_single_nemo_gym_group( nemo_gym_rows: list[dict], results: list[dict], @@ -2359,28 +2405,15 @@ def _postprocess_single_nemo_gym_group( max_total_tokens_per_sample = policy_generation.cfg[ "max_total_sequence_length" ] + # Absent for backends whose generation config omits it; the context check then + # applies on its own, matching the previous behaviour. + max_new_tokens = policy_generation.cfg.get("max_new_tokens") all_sample_metrics = [ - { - "total_reward": r["full_result"]["reward"], - "assistant_tokens": sum( - len(m["token_ids"]) - for m in r["message_log"] - if m["role"] == "assistant" - ), - "total_tokens": sum(len(m["token_ids"]) for m in r["message_log"]), - "turn_count": sum(1 for m in r["message_log"] if m["role"] == "user"), - "hit_max_tokens": sum(len(m["token_ids"]) for m in r["message_log"]) - == max_total_tokens_per_sample, - # max_gen_tokens_per_turn: Diagnostic for long single generations - "max_gen_tokens_per_turn": max( - ( - len(m["token_ids"]) - for m in r["message_log"] - if m["role"] == "assistant" - ), - default=0, - ), - } + _nemo_gym_sample_metrics( + r, + max_total_tokens=max_total_tokens_per_sample, + max_new_tokens=max_new_tokens, + ) for r in results ] diff --git a/tests/unit/experience/test_rollouts.py b/tests/unit/experience/test_rollouts.py index 038384a5958..3ac1b344efe 100644 --- a/tests/unit/experience/test_rollouts.py +++ b/tests/unit/experience/test_rollouts.py @@ -1299,6 +1299,99 @@ def test_postprocess_nemo_gym_group_returns_task_index(log_full_result_tables): ) is log_full_result_tables +def _gym_result(prompt_tokens, assistant_turns, reward=1.0): + message_log = [ + { + "role": "user", + "content": "prompt", + "token_ids": torch.zeros(prompt_tokens, dtype=torch.long), + } + ] + for turn_tokens in assistant_turns: + message_log.append( + { + "role": "assistant", + "content": "answer", + "token_ids": torch.zeros(turn_tokens, dtype=torch.long), + "generation_logprobs": torch.zeros(turn_tokens), + } + ) + return { + "input_message_log": message_log[:1], + "message_log": message_log, + "full_result": {"reward": reward}, + } + + +@pytest.mark.parametrize( + "prompt_tokens,assistant_turns,max_total_tokens,max_new_tokens,expected", + [ + # The conversation filled the context. Detected before this check also + # considered the per-turn cap, and still detected now. + (10, [90], 100, None, True), + # A turn reached the generation cap well inside the context. This is the case + # that used to go unflagged, so overlong_filtering silently kept the sample. + (10, [64], 1024, 64, True), + # Neither budget reached: no single turn hits the cap, context has room left. + (10, [20, 20], 1024, 64, False), + # `>=`, not `==`: the message log is post-processed, so an exact match can be + # thrown off by a token or two. + (10, [95], 100, None, True), + # No generation cap configured: context check alone, matching prior behaviour. + (10, [64], 1024, None, False), + ], +) +def test_nemo_gym_sample_metrics_detects_both_truncation_budgets( + prompt_tokens, assistant_turns, max_total_tokens, max_new_tokens, expected +): + metrics = rollouts_mod._nemo_gym_sample_metrics( + _gym_result(prompt_tokens, assistant_turns), + max_total_tokens=max_total_tokens, + max_new_tokens=max_new_tokens, + ) + + assert metrics["hit_max_tokens"] is expected + + +def test_nemo_gym_sample_metrics_reports_lengths(): + metrics = rollouts_mod._nemo_gym_sample_metrics( + _gym_result(10, [30, 20], reward=2.5), + max_total_tokens=1024, + max_new_tokens=None, + ) + + assert metrics["total_reward"] == 2.5 + assert metrics["assistant_tokens"] == 50 + assert metrics["total_tokens"] == 60 + assert metrics["turn_count"] == 1 + assert metrics["max_gen_tokens_per_turn"] == 30 + + +def test_postprocess_nemo_gym_group_flags_generation_cap_truncation(): + """final_batch["truncated"] is what grpo.overlong_filtering masks on.""" + rows = [{"agent_ref": {"name": "agent"}} for _ in range(2)] + # First rollout stops at the generation cap; second finishes early. Neither comes + # close to max_model_len, so the context check alone would flag neither. + results = [_gym_result(10, [16]), _gym_result(10, [4])] + + rollout_result = rollouts_mod._postprocess_single_nemo_gym_group( + nemo_gym_rows=rows, + results=results, + timer=rollouts_mod.Timer(), + timer_prefix="timing/rollout", + policy_generation=type( + "_PolicyGeneration", + (), + {"cfg": {"vllm_cfg": {"max_model_len": 1024}, "max_new_tokens": 16}}, + )(), + input_batch=BatchedDataDict({"loss_multiplier": torch.ones(2)}), + tokenizer=type("_Tokenizer", (), {"pad_token_id": 0})(), + log_full_result_tables=False, + ) + + assert rollout_result.final_batch["truncated"].tolist() == [True, False] + + def test_run_nemo_gym_rollout_sync_drains_entire_batch(monkeypatch): input_batch = BatchedDataDict({"loss_multiplier": torch.ones(3)}) expected = rollouts_mod.NemoGymRolloutResult( From a61a49a2673ffeae8515154c896e2e37608f6cff Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 21 Aug 2026 12:42:22 -0400 Subject: [PATCH 2/2] fix: use the agent's reported truncation flag when it has one The previous commit inferred truncation from token counts because the Gym path appeared to have no better signal. It does: verifiers_agent now reports is_truncated, derived from the model server's finish_reason == "length" -- the same ground truth the native generation path already uses (vllm_worker_async.py, "is_truncated = generation_details.finish_reason == 'length'"), and the semantic docs/guides/grpo.md describes for overlong filtering. Prefer it. It needs no knowledge of which budget was binding, so it covers context exhaustion, an environment's own per-turn cap and max_new_tokens alike, and it does not mistake a turn that emitted EOS on its final allowed token for one that ran out of budget. This also converges the NeMo-Gym path with the native path rather than maintaining a second, weaker notion of truncation. Keep the length comparison as a fallback: agents other than verifiers_agent report no flag, and omitting the field selects the previous behaviour. Bumps the Gym submodule for the reporting change. Signed-off-by: Albert Cui --- 3rdparty/Gym-workspace/Gym | 2 +- nemo_rl/experience/rollouts.py | 45 ++++++++++++++++---------- tests/unit/experience/test_rollouts.py | 29 +++++++++++++++-- 3 files changed, 56 insertions(+), 20 deletions(-) diff --git a/3rdparty/Gym-workspace/Gym b/3rdparty/Gym-workspace/Gym index 6230e442921..c4fc3d9d6ca 160000 --- a/3rdparty/Gym-workspace/Gym +++ b/3rdparty/Gym-workspace/Gym @@ -1 +1 @@ -Subproject commit 6230e442921c4ed7ef33060c63cb8871e4019dea +Subproject commit c4fc3d9d6cac8c73beea1f09042e96ab39e5103d diff --git a/nemo_rl/experience/rollouts.py b/nemo_rl/experience/rollouts.py index 5140fbc4344..78118ca2f18 100644 --- a/nemo_rl/experience/rollouts.py +++ b/nemo_rl/experience/rollouts.py @@ -2322,20 +2322,27 @@ def _nemo_gym_sample_metrics( """Per-sample rollout metrics, including whether generation was cut short. ``hit_max_tokens`` becomes ``final_batch["truncated"]``, which - ``grpo.overlong_filtering`` uses to drop a sample from the loss. A rollout can be - stopped by either of two budgets, so both are checked: - - * the conversation filled the context (``max_total_tokens``), or - * a single turn reached the per-turn generation cap (``max_new_tokens``). - - Checking only the first misses every rollout that stopped below the context. That is - not a corner case: ``max_new_tokens`` is applied per turn, so a run with a small - ``max_new_tokens`` never fills the context at all and its truncated samples keep the - near-zero reward that overlong filtering exists to discard. - - Both comparisons are ``>=`` rather than ``==``. The message log is post-processed - before it reaches here -- reasoning content is re-wrapped in thinking tags, for one -- - so an exact match can be thrown off by a single token. + ``grpo.overlong_filtering`` uses to drop a sample from the loss. + + Prefer the agent's own ``is_truncated``. It is the same ground truth the native + generation path uses -- the model server's ``finish_reason == "length"`` -- carried + through the agent, so it distinguishes "the budget stopped this" from "the model + emitted EOS" without knowing which budget applied. NeMo-Gym's ``verifiers_agent`` + reports it; agents that do not simply omit the field. + + Fall back to comparing lengths when it is absent. A rollout can be stopped by either + of two budgets, so both are checked: the conversation filling the context + (``max_total_tokens``), or a single turn reaching the per-turn generation cap + (``max_new_tokens``). Checking only the first misses every rollout that stopped below + the context, which is not a corner case -- ``max_new_tokens`` applies per turn, so a + run with a small ``max_new_tokens`` never fills the context at all and its truncated + samples keep the near-zero reward that overlong filtering exists to discard. + + The fallback compares with ``>=`` rather than ``==``. The message log is + post-processed before it reaches here -- reasoning content is re-wrapped in thinking + tags, for one -- so an exact match can be thrown off by a single token. It also + cannot tell a turn that ran out of budget from one that happened to emit EOS on the + final allowed token, which is why the reported flag wins when there is one. """ message_log = result["message_log"] assistant_lengths = [ @@ -2344,9 +2351,13 @@ def _nemo_gym_sample_metrics( total_tokens = sum(len(m["token_ids"]) for m in message_log) max_gen_tokens_per_turn = max(assistant_lengths, default=0) - hit_max_tokens = total_tokens >= max_total_tokens - if max_new_tokens is not None: - hit_max_tokens = hit_max_tokens or max_gen_tokens_per_turn >= max_new_tokens + reported_truncated = (result.get("full_result") or {}).get("is_truncated") + if reported_truncated is not None: + hit_max_tokens = bool(reported_truncated) + else: + hit_max_tokens = total_tokens >= max_total_tokens + if max_new_tokens is not None: + hit_max_tokens = hit_max_tokens or max_gen_tokens_per_turn >= max_new_tokens return { "total_reward": result["full_result"]["reward"], diff --git a/tests/unit/experience/test_rollouts.py b/tests/unit/experience/test_rollouts.py index 3ac1b344efe..1f1ebf49ecc 100644 --- a/tests/unit/experience/test_rollouts.py +++ b/tests/unit/experience/test_rollouts.py @@ -1299,7 +1299,7 @@ def test_postprocess_nemo_gym_group_returns_task_index(log_full_result_tables): ) is log_full_result_tables -def _gym_result(prompt_tokens, assistant_turns, reward=1.0): +def _gym_result(prompt_tokens, assistant_turns, reward=1.0, is_truncated=None): message_log = [ { "role": "user", @@ -1316,10 +1316,13 @@ def _gym_result(prompt_tokens, assistant_turns, reward=1.0): "generation_logprobs": torch.zeros(turn_tokens), } ) + full_result = {"reward": reward} + if is_truncated is not None: + full_result["is_truncated"] = is_truncated return { "input_message_log": message_log[:1], "message_log": message_log, - "full_result": {"reward": reward}, + "full_result": full_result, } @@ -1353,6 +1356,28 @@ def test_nemo_gym_sample_metrics_detects_both_truncation_budgets( assert metrics["hit_max_tokens"] is expected +@pytest.mark.parametrize( + "reported,assistant_turns,max_new_tokens,expected", + [ + # Reported flag wins over the length fallback in BOTH directions. The second case + # is the one lengths cannot get right: a turn that emitted EOS on its final + # allowed token looks identical to one that ran out of budget. + (True, [4], 64, True), + (False, [64], 64, False), + ], +) +def test_nemo_gym_sample_metrics_prefers_the_reported_flag( + reported, assistant_turns, max_new_tokens, expected +): + metrics = rollouts_mod._nemo_gym_sample_metrics( + _gym_result(10, assistant_turns, is_truncated=reported), + max_total_tokens=1024, + max_new_tokens=max_new_tokens, + ) + + assert metrics["hit_max_tokens"] is expected + + def test_nemo_gym_sample_metrics_reports_lengths(): metrics = rollouts_mod._nemo_gym_sample_metrics( _gym_result(10, [30, 20], reward=2.5),