diff --git a/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml b/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml index 08b728bef4..0eecae6111 100644 --- a/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml +++ b/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml @@ -299,6 +299,7 @@ env: port_range_high: 5999 rollout_max_attempts_to_avoid_lp_nan: 1 is_trajectory_collection: false # Set this to true to enable trajectory collection (no training). You may also want to increase `policy.generation.vllm_cfg.gpu_memory_utilization` + trajectory_collection_batch_size: null # Optional positive integer; null collects the validation set in one batch config_paths: - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml # Required! And it must be *for_training - resources_servers/workplace_assistant/configs/workplace_assistant.yaml diff --git a/examples/nemo_gym/run_grpo_nemo_gym.py b/examples/nemo_gym/run_grpo_nemo_gym.py index 25f4646162..a7c2e178f8 100644 --- a/examples/nemo_gym/run_grpo_nemo_gym.py +++ b/examples/nemo_gym/run_grpo_nemo_gym.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# 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. @@ -13,6 +13,7 @@ # limitations under the License. import argparse +import json import os import pprint import time @@ -38,7 +39,7 @@ refit_policy_generation, setup, ) -from nemo_rl.algorithms.utils import get_tokenizer +from nemo_rl.algorithms.utils import get_tokenizer, log_generation_metrics_to_wandb from nemo_rl.data.utils import setup_response_data from nemo_rl.distributed.virtual_cluster import init_ray from nemo_rl.environments.nemo_gym import ( @@ -68,6 +69,34 @@ def parse_args() -> tuple[argparse.Namespace, list[str]]: return args, overrides +def _pop_trajectory_collection_settings( + nemo_gym_config: dict[str, object], +) -> tuple[bool, int | None]: + """Remove and validate NeMo-RL trajectory-collection settings.""" + is_trajectory_collection = bool( + nemo_gym_config.pop("is_trajectory_collection", False) + ) + batch_size = nemo_gym_config.pop("trajectory_collection_batch_size", None) + if batch_size is None: + return is_trajectory_collection, None + + if not is_trajectory_collection: + raise ValueError( + "env.nemo_gym.trajectory_collection_batch_size requires " + "env.nemo_gym.is_trajectory_collection=true" + ) + if ( + isinstance(batch_size, bool) + or not isinstance(batch_size, int) + or batch_size <= 0 + ): + raise ValueError( + "env.nemo_gym.trajectory_collection_batch_size must be a positive integer" + ) + + return is_trajectory_collection, batch_size + + # These types are directly imported from grpo_train since if something about the architecture changes we want to immediately fail. def collect_trajectories( policy: ColocatablePolicyInterface, @@ -78,7 +107,13 @@ def collect_trajectories( logger: Logger, master_config: MasterConfig, ) -> None: - """Run trajectory collection.""" + """Run trajectory collection and persist every completed batch.""" + expected_trajectories = master_config.grpo["max_val_samples"] + if expected_trajectories is None or expected_trajectories <= 0: + raise ValueError( + "Trajectory collection requires a non-empty validation dataset" + ) + # common config/state items colocated_inference = master_config.policy["generation"]["colocated"]["enabled"] refit_policy_generation(policy, policy_generation, colocated_inference) @@ -87,33 +122,108 @@ def collect_trajectories( print("\nšŸ” Running trajectory collection...", flush=True) generation_config = master_config.policy["generation"] - for val_batch in val_dataloader: - nemo_gym_rollout_result = run_async_nemo_gym_rollout( - policy_generation=policy_generation, - input_batch=val_batch, - tokenizer=tokenizer, - task_to_env=val_task_to_env, - max_seq_len=master_config.policy["max_total_sequence_length"], - generation_config=generation_config, - max_rollout_turns=None, - greedy=False, - ) - - rows_to_log: list[str] = [] - for key, value in nemo_gym_rollout_result.rollout_metrics.items(): - if "full_result" not in key: - continue - - value: Table - data: list[list[str]] = value.data # (n, 1) - rows_to_log.extend(v[0] for v in data) - - logger.log_string_list_as_jsonl(rows_to_log, log_filename) + vllm_config = generation_config.get("vllm_cfg", {}) + should_log_generation_metrics = ( + vllm_config.get("enable_vllm_metrics_logger", False) + and vllm_config.get("async_engine", False) + and master_config.logger["wandb_enabled"] + ) + collected_trajectories = 0 + total_reward = 0.0 + + try: + for batch_idx, val_batch in enumerate(val_dataloader): + batch_step = batch_idx + 1 + if should_log_generation_metrics: + policy_generation.clear_logger_metrics() + + nemo_gym_rollout_result = run_async_nemo_gym_rollout( + policy_generation=policy_generation, + input_batch=val_batch, + tokenizer=tokenizer, + task_to_env=val_task_to_env, + max_seq_len=master_config.policy["max_total_sequence_length"], + generation_config=generation_config, + max_rollout_turns=None, + greedy=False, + ) + if should_log_generation_metrics: + generation_logger_metrics = policy_generation.get_logger_metrics() + + rows_to_log: list[str] = [] + for key, value in nemo_gym_rollout_result.rollout_metrics.items(): + if "full_result" not in key: + continue + + value: Table + data: list[list[str]] = value.data # (n, 1) + rows_to_log.extend(v[0] for v in data) + + if not rows_to_log: + raise RuntimeError( + f"Trajectory batch {batch_idx} did not contain any full Gym results" + ) + + attributed_rows: list[str] = [] + batch_size = len(rows_to_log) + batch_reward = 0.0 + for batch_position, serialized_result in enumerate(rows_to_log): + result = json.loads(serialized_result) + result["trajectory_collection_batch_index"] = batch_idx + result["trajectory_collection_batch_position"] = batch_position + result["trajectory_collection_batch_size"] = batch_size + batch_reward += float(result["reward"]) + attributed_rows.append(json.dumps(result, separators=(",", ":"))) + + # Append after every completed batch so earlier trajectories survive a later + # batch or worker failure during a long collection run. + logger.log_string_list_as_jsonl(attributed_rows, log_filename) + collected_trajectories += batch_size + total_reward += batch_reward + + batch_rollout_metrics = { + key: value + for key, value in nemo_gym_rollout_result.rollout_metrics.items() + if "full_result" not in key + } + # Match the training prefix so rollout-only and GRPO runs expose the same + # timing and rollout metric names for direct comparison. + logger.log_metrics(batch_rollout_metrics, batch_step, prefix="train") + if should_log_generation_metrics: + log_generation_metrics_to_wandb( + generation_logger_metrics, + batch_step, + vllm_config["vllm_metrics_logger_interval"], + logger, + ) + logger.log_metrics( + { + "mean_reward": total_reward / collected_trajectories, + "num_trajectories": collected_trajectories, + }, + batch_step, + prefix="trajectory_collection", + step_finished=True, + ) + print( + f"Collected {collected_trajectories}/{expected_trajectories} " + f"trajectories after batch {batch_idx + 1}", + flush=True, + ) + finally: + policy_generation.finish_generation() - # TODO: eventually as trajectory collection use cases exceed 4 hours, we can leverage the dataloader save functionality to resume - # And also leverage the TimeoutChecker functionality as well + if collected_trajectories != expected_trajectories: + raise RuntimeError( + "Trajectory collection was incomplete: " + f"expected {expected_trajectories}, got {collected_trajectories}" + ) - policy_generation.finish_generation() + print( + f"Trajectory collection complete: {collected_trajectories} trajectories, " + f"mean reward {total_reward / collected_trajectories:.6f}", + flush=True, + ) def main() -> None: @@ -163,6 +273,12 @@ def main() -> None: # NeMo-Gym specific config setup. setup_nemo_gym_config(config, tokenizer) + # These are NeMo-RL control-flow settings, not NeMo-Gym global config. + ( + is_trajectory_collection, + trajectory_collection_batch_size, + ) = _pop_trajectory_collection_settings(config.env["nemo_gym"]) + # We assert here since this is right after the final config has been materialized. assert _should_use_nemo_gym(config) @@ -184,11 +300,15 @@ def main() -> None: ) if val_dataset is not None: + val_batch_size = len(val_dataset) + if trajectory_collection_batch_size is not None: + val_batch_size = min(trajectory_collection_batch_size, len(val_dataset)) print( - f"Setting `grpo.max_val_samples` and `grpo.val_batch_size` to the length of the validation dataset, which is {len(val_dataset)}" + f"Setting `grpo.max_val_samples` to {len(val_dataset)} and " + f"`grpo.val_batch_size` to {val_batch_size}" ) config.grpo["max_val_samples"] = len(val_dataset) - config.grpo["val_batch_size"] = config.grpo["max_val_samples"] + config.grpo["val_batch_size"] = val_batch_size # Print config print("Final config:") @@ -197,13 +317,6 @@ def main() -> None: with rl_init_timer.time("ray_connect"): init_ray() - # `is_trajectory_collection` is a NeMo-RL-side control-flow knob; pop it - # before setup() so it is not forwarded into NeMo-Gym's global config (the - # gym actor is now created inside setup()). - is_trajectory_collection = ( - config.env["nemo_gym"].pop("is_trajectory_collection", False) or False - ) - with rl_init_timer.time("setup"): ( policy, diff --git a/tests/unit/test_run_grpo_nemo_gym.py b/tests/unit/test_run_grpo_nemo_gym.py new file mode 100644 index 0000000000..0461a03b6d --- /dev/null +++ b/tests/unit/test_run_grpo_nemo_gym.py @@ -0,0 +1,284 @@ +import json +from types import SimpleNamespace +from unittest.mock import MagicMock, call + +import pytest + +from examples.nemo_gym import run_grpo_nemo_gym + + +def _rollout_result( + rewards: list[float], + rollout_time: float, +) -> SimpleNamespace: + full_results = [json.dumps({"reward": reward}) for reward in rewards] + return SimpleNamespace( + rollout_metrics={ + "timing/rollout/total": rollout_time, + "timing/rollout/run_rollouts": rollout_time - 1, + "mean_gen_tokens_per_sample": 4.0, + "test_agent/full_result": SimpleNamespace( + data=[[result] for result in full_results] + ), + } + ) + + +def _master_config( + *, expected_trajectories: int, wandb_enabled: bool = True +) -> SimpleNamespace: + return SimpleNamespace( + policy={ + "generation": { + "colocated": {"enabled": False}, + "vllm_cfg": { + "async_engine": True, + "enable_vllm_metrics_logger": True, + "vllm_metrics_logger_interval": 0.5, + }, + }, + "max_total_sequence_length": 4096, + }, + grpo={"max_val_samples": expected_trajectories}, + logger={"wandb_enabled": wandb_enabled}, + ) + + +def test_pop_trajectory_collection_settings() -> None: + nemo_gym_config = { + "is_trajectory_collection": True, + "trajectory_collection_batch_size": 16, + "config_paths": ["gym.yaml"], + } + + settings = run_grpo_nemo_gym._pop_trajectory_collection_settings(nemo_gym_config) + + assert settings == (True, 16) + assert nemo_gym_config == {"config_paths": ["gym.yaml"]} + + +@pytest.mark.parametrize("batch_size", [True, 0, -1, 1.5, "16"]) +def test_pop_trajectory_collection_settings_rejects_invalid_batch_size( + batch_size: object, +) -> None: + with pytest.raises(ValueError, match="must be a positive integer"): + run_grpo_nemo_gym._pop_trajectory_collection_settings( + { + "is_trajectory_collection": True, + "trajectory_collection_batch_size": batch_size, + } + ) + + +def test_pop_trajectory_collection_settings_requires_collection_mode() -> None: + with pytest.raises(ValueError, match="requires"): + run_grpo_nemo_gym._pop_trajectory_collection_settings( + { + "is_trajectory_collection": False, + "trajectory_collection_batch_size": 16, + } + ) + + +def test_collect_trajectories_logs_each_batch_and_generation_metrics( + monkeypatch: pytest.MonkeyPatch, +) -> None: + refit_policy_generation = MagicMock() + run_rollout = MagicMock( + side_effect=[ + _rollout_result([1.0, 0.0], rollout_time=12.0), + _rollout_result([0.5], rollout_time=18.0), + ] + ) + log_generation_metrics = MagicMock() + monkeypatch.setattr( + run_grpo_nemo_gym, "refit_policy_generation", refit_policy_generation + ) + monkeypatch.setattr(run_grpo_nemo_gym, "run_async_nemo_gym_rollout", run_rollout) + monkeypatch.setattr( + run_grpo_nemo_gym, + "log_generation_metrics_to_wandb", + log_generation_metrics, + ) + + generation_metrics = [ + {"inflight_batch_sizes": {0: [1, 0]}}, + {"inflight_batch_sizes": {0: [2, 0]}}, + ] + policy_generation = MagicMock() + policy_generation.get_logger_metrics.side_effect = generation_metrics + logger = MagicMock() + + run_grpo_nemo_gym.collect_trajectories( + policy=MagicMock(), + policy_generation=policy_generation, + val_dataloader=[object(), object()], + tokenizer=MagicMock(), + val_task_to_env={"nemo_gym": MagicMock()}, + logger=logger, + master_config=_master_config(expected_trajectories=3), + ) + + refit_policy_generation.assert_called_once() + assert policy_generation.clear_logger_metrics.call_count == 2 + assert policy_generation.get_logger_metrics.call_count == 2 + policy_generation.finish_generation.assert_called_once_with() + + assert logger.log_string_list_as_jsonl.call_count == 2 + logged_batches = [ + [json.loads(row) for row in log_call.args[0]] + for log_call in logger.log_string_list_as_jsonl.call_args_list + ] + assert [ + result["trajectory_collection_batch_index"] + for batch in logged_batches + for result in batch + ] == [0, 0, 1] + assert [ + result["trajectory_collection_batch_position"] + for batch in logged_batches + for result in batch + ] == [0, 1, 0] + assert [ + result["trajectory_collection_batch_size"] + for batch in logged_batches + for result in batch + ] == [2, 2, 1] + + rollout_log_calls = [ + log_call + for log_call in logger.log_metrics.call_args_list + if log_call.kwargs.get("prefix") == "train" + ] + assert [ + log_call.args[0]["timing/rollout/total"] for log_call in rollout_log_calls + ] == [12.0, 18.0] + assert [log_call.args[1] for log_call in rollout_log_calls] == [1, 2] + assert all( + "full_result" not in key + for log_call in rollout_log_calls + for key in log_call.args[0] + ) + + assert log_generation_metrics.call_args_list == [ + call(generation_metrics[0], 1, 0.5, logger), + call(generation_metrics[1], 2, 0.5, logger), + ] + collection_log_calls = [ + log_call + for log_call in logger.log_metrics.call_args_list + if log_call.kwargs.get("prefix") == "trajectory_collection" + ] + assert collection_log_calls[-1].args[0] == { + "mean_reward": 0.5, + "num_trajectories": 3, + } + assert collection_log_calls[-1].args[1] == 2 + assert collection_log_calls[-1].kwargs["step_finished"] is True + + +def test_collect_trajectories_rejects_empty_validation_dataset( + monkeypatch: pytest.MonkeyPatch, +) -> None: + refit_policy_generation = MagicMock() + monkeypatch.setattr( + run_grpo_nemo_gym, "refit_policy_generation", refit_policy_generation + ) + + with pytest.raises(ValueError, match="non-empty validation dataset"): + run_grpo_nemo_gym.collect_trajectories( + policy=MagicMock(), + policy_generation=MagicMock(), + val_dataloader=[], + tokenizer=MagicMock(), + val_task_to_env={"nemo_gym": MagicMock()}, + logger=MagicMock(), + master_config=_master_config(expected_trajectories=0), + ) + + refit_policy_generation.assert_not_called() + + +def test_collect_trajectories_skips_generation_metrics_without_wandb( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(run_grpo_nemo_gym, "refit_policy_generation", MagicMock()) + monkeypatch.setattr( + run_grpo_nemo_gym, + "run_async_nemo_gym_rollout", + MagicMock(return_value=_rollout_result([1.0], rollout_time=12.0)), + ) + log_generation_metrics = MagicMock() + monkeypatch.setattr( + run_grpo_nemo_gym, + "log_generation_metrics_to_wandb", + log_generation_metrics, + ) + policy_generation = MagicMock() + + run_grpo_nemo_gym.collect_trajectories( + policy=MagicMock(), + policy_generation=policy_generation, + val_dataloader=[object()], + tokenizer=MagicMock(), + val_task_to_env={"nemo_gym": MagicMock()}, + logger=MagicMock(), + master_config=_master_config(expected_trajectories=1, wandb_enabled=False), + ) + + policy_generation.clear_logger_metrics.assert_not_called() + policy_generation.get_logger_metrics.assert_not_called() + log_generation_metrics.assert_not_called() + policy_generation.finish_generation.assert_called_once_with() + + +def test_collect_trajectories_preserves_completed_batches_before_incomplete_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(run_grpo_nemo_gym, "refit_policy_generation", MagicMock()) + monkeypatch.setattr( + run_grpo_nemo_gym, + "run_async_nemo_gym_rollout", + MagicMock(return_value=_rollout_result([1.0], rollout_time=12.0)), + ) + policy_generation = MagicMock() + logger = MagicMock() + + with pytest.raises(RuntimeError, match="expected 2, got 1"): + run_grpo_nemo_gym.collect_trajectories( + policy=MagicMock(), + policy_generation=policy_generation, + val_dataloader=[object()], + tokenizer=MagicMock(), + val_task_to_env={"nemo_gym": MagicMock()}, + logger=logger, + master_config=_master_config(expected_trajectories=2), + ) + + logger.log_string_list_as_jsonl.assert_called_once() + policy_generation.finish_generation.assert_called_once_with() + + +def test_collect_trajectories_finishes_generation_after_rollout_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(run_grpo_nemo_gym, "refit_policy_generation", MagicMock()) + monkeypatch.setattr( + run_grpo_nemo_gym, + "run_async_nemo_gym_rollout", + MagicMock(side_effect=RuntimeError("rollout failed")), + ) + policy_generation = MagicMock() + + with pytest.raises(RuntimeError, match="rollout failed"): + run_grpo_nemo_gym.collect_trajectories( + policy=MagicMock(), + policy_generation=policy_generation, + val_dataloader=[object()], + tokenizer=MagicMock(), + val_task_to_env={"nemo_gym": MagicMock()}, + logger=MagicMock(), + master_config=_master_config(expected_trajectories=1), + ) + + policy_generation.finish_generation.assert_called_once_with()