diff --git a/3rdparty/Gym-workspace/Gym b/3rdparty/Gym-workspace/Gym index 74a67f6d104..c4fc3d9d6ca 160000 --- a/3rdparty/Gym-workspace/Gym +++ b/3rdparty/Gym-workspace/Gym @@ -1 +1 @@ -Subproject commit 74a67f6d104ff116862a041473f71726ee1697e5 +Subproject commit c4fc3d9d6cac8c73beea1f09042e96ab39e5103d diff --git a/nemo_rl/experience/rollouts.py b/nemo_rl/experience/rollouts.py index 1a6145160b8..78118ca2f18 100644 --- a/nemo_rl/experience/rollouts.py +++ b/nemo_rl/experience/rollouts.py @@ -2313,6 +2313,63 @@ 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. + + 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 = [ + 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) + + 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"], + "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 +2416,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..1f1ebf49ecc 100644 --- a/tests/unit/experience/test_rollouts.py +++ b/tests/unit/experience/test_rollouts.py @@ -1299,6 +1299,124 @@ 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, is_truncated=None): + 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), + } + ) + 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": full_result, + } + + +@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 + + +@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), + 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(