diff --git a/examples/nemo_gym/run_distillation_nemo_gym.py b/examples/nemo_gym/run_distillation_nemo_gym.py index 56fe23eb6b..ca2fad0d95 100644 --- a/examples/nemo_gym/run_distillation_nemo_gym.py +++ b/examples/nemo_gym/run_distillation_nemo_gym.py @@ -32,9 +32,7 @@ from nemo_rl.algorithms.utils import get_tokenizer 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 ( - setup_nemo_gym_config, -) +from nemo_rl.environments.nemo_gym import setup_nemo_gym_config from nemo_rl.models.generation import configure_generation_config from nemo_rl.utils.config import ( load_config, diff --git a/examples/run_grpo_single_controller.py b/examples/run_grpo_single_controller.py index 3b776ff773..544124c9fa 100644 --- a/examples/run_grpo_single_controller.py +++ b/examples/run_grpo_single_controller.py @@ -34,6 +34,7 @@ ) from nemo_rl.algorithms.utils import get_tokenizer from nemo_rl.distributed.virtual_cluster import init_ray +from nemo_rl.environments.nemo_gym import setup_nemo_gym_config from nemo_rl.models.generation import configure_generation_config from nemo_rl.utils.config import ( load_config, @@ -117,6 +118,10 @@ def main() -> None: trains_mtp=trains_mtp, ) + # NeMo-Gym specific config setup. + if bool(config.env.get("should_use_nemo_gym")): + setup_nemo_gym_config(config, tokenizer) + actor_args = setup_single_controller(config, tokenizer) print("🚀 Launching SingleControllerActor") diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index b9234ae739..d83f07ecf7 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -118,6 +118,26 @@ def validate_single_controller_config(master_config: MasterConfig) -> None: sampler_name=async_config.sampler.name, ) + # A non-zero reference-policy KL penalty makes the loss read + # ``reference_policy_logprobs``, but the SC train pump only computes them + # when ``skip_reference_policy_logprobs_calculation`` is false (see + # SingleControllerActor._reference_logprobs_required). Catch the + # inconsistent pair at setup instead of a mid-training KeyError. + reference_policy_kl_penalty = getattr( + master_config.loss_fn, "reference_policy_kl_penalty", 0 + ) + if reference_policy_kl_penalty and master_config.grpo.get( + "skip_reference_policy_logprobs_calculation" + ): + raise ValueError( + "loss_fn.reference_policy_kl_penalty=" + f"{reference_policy_kl_penalty} requires reference_policy_logprobs, " + "but grpo.skip_reference_policy_logprobs_calculation=true skips " + "computing them on the SingleController path. Set " + "grpo.skip_reference_policy_logprobs_calculation=false, or set " + "loss_fn.reference_policy_kl_penalty=0." + ) + # ── Internal SingleController configs ──────────────────────────────────── diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index e9db9d6690..821c50cfd9 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -49,11 +49,16 @@ from nemo_rl.data_plane import DataPlaneClient, build_data_plane_client from nemo_rl.distributed.virtual_cluster import RayVirtualCluster from nemo_rl.environments.interfaces import EnvironmentInterface +from nemo_rl.environments.nemo_gym import spinup_nemo_gym_actor from nemo_rl.experience.rollout_manager import RolloutManager from nemo_rl.models.generation.sglang.config import SGLangConfig from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration from nemo_rl.models.generation.vllm import VllmGeneration from nemo_rl.models.generation.vllm.config import VllmConfig +from nemo_rl.models.megatron.router_replay import ( + configure_vllm_for_router_replay, + router_replay_enabled, +) from nemo_rl.models.policy.tq_policy import TQPolicy from nemo_rl.weight_sync import WeightSynchronizer, create_weight_synchronizer @@ -167,6 +172,7 @@ def _build_generation( vllm_config.setdefault("vllm_kwargs", {})["hf_overrides"] = ( master_config.policy.get("hf_config_overrides", {}) ) + configure_vllm_for_router_replay(master_config.policy) gen = VllmGeneration(cluster=inference_cluster, config=vllm_config) elif backend == "sglang": sglang_config = cast(SGLangConfig, generation_config) @@ -318,16 +324,25 @@ def setup_single_controller( # Setup Dataset & Environments # ========================== # TODO: add validate dataset wiring. - if _should_use_nemo_gym(cast(GrpoMasterConfig, master_config)): + use_nemo_gym = _should_use_nemo_gym(cast(GrpoMasterConfig, master_config)) + if use_nemo_gym and generation_config["backend"] != "vllm": raise NotImplementedError( - "NeMo-Gym integration for SingleController is not supported yet; " - "it will land in https://github.com/NVIDIA-NeMo/RL/pull/3267." + "SC NeMo-Gym integration currently supports the vllm backend " + f"only; got {generation_config['backend']!r}" ) - response_data = setup_response_data( - tokenizer, data_config, env_configs=master_config.env - ) - assert len(response_data) == 4 - dataset, _val_dataset, env_handles, _val_env_handles = response_data + if use_nemo_gym: + # NeMo-Gym creates the env actor outside setup_response_data; we wire + # it in after generation is up (it needs the OpenAI server URLs). + response_data = setup_response_data(tokenizer, data_config, env_configs=None) + assert len(response_data) == 2 + dataset, _val_dataset = response_data + env_handles: dict[str, EnvironmentInterface] = {} + else: + response_data = setup_response_data( + tokenizer, data_config, env_configs=master_config.env + ) + assert len(response_data) == 4 + dataset, _val_dataset, env_handles, _val_env_handles = response_data dataloader = StatefulDataLoader( dataset, batch_size=grpo_config["num_prompts_per_step"], @@ -363,6 +378,20 @@ def setup_single_controller( generation = gen_future.result() policy = policy_future.result() + # ========================== + # NeMo-Gym actor (after generation is up so OpenAI URLs are available) + # ========================== + if use_nemo_gym: + # TODO(#2625): Mirror GRPO's deferred vLLM load so NeMo-Gym spinup + # overlaps model loading instead of running serially afterward. + enable_router_replay = router_replay_enabled(master_config.policy) + env_handles["nemo_gym"] = spinup_nemo_gym_actor( + env_configs=master_config.env, + base_urls=generation.dp_openai_server_base_urls, + model_name=generation_config["model_name"], + enable_router_replay=enable_router_replay, + ) + # ========================== # Setup Data Plane Client & Weight Sync # ========================== @@ -403,6 +432,7 @@ def setup_single_controller( max_rollout_turns=grpo_config.get("max_rollout_turns"), policy_generation=generation, generation_config=generation_config, + use_nemo_gym=use_nemo_gym, tq_buffer=tq_buffer, ) diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 513cb00b18..37f4f13150 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -18,12 +18,14 @@ from collections import Counter from collections.abc import AsyncGenerator from pathlib import Path -from typing import Any, Dict, List, NotRequired, TypedDict +from typing import Any, Dict, List, NotRequired, Optional, TypedDict import ray import torch +from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy from transformers import PreTrainedTokenizerBase +from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env from nemo_rl.distributed.virtual_cluster import ( DEFAULT_GYM_PORT_RANGE_HIGH, DEFAULT_GYM_PORT_RANGE_LOW, @@ -32,6 +34,7 @@ ) from nemo_rl.environments.interfaces import EnvironmentInterface from nemo_rl.utils.timer import Timer +from nemo_rl.utils.venvs import create_local_venv_on_each_node # Kept local (not imported from models.generation) so the gym actor stays free of # generation-module imports. Must cover every name resolve_routed_experts_dtype @@ -600,3 +603,79 @@ def setup_nemo_gym_config(config, tokenizer) -> None: # Stop strings or token ids are not supported generation_config["stop_strings"] = None generation_config["stop_token_ids"] = None + + +def spinup_nemo_gym_actor( + env_configs: dict[str, Any], + base_urls: list[Optional[str]], + model_name: str, + enable_router_replay: bool, +) -> Any: + """Spin up the NeMo-Gym actor against the given generation server URLs. + + When ``env_configs["nemo_gym"]["num_gpu_nodes"] > 0``, the actor is + scheduled with soft NodeAffinity to the current Ray node so its colocated + GPU resources land where the caller expects. + + Args: + env_configs: The ``master_config.env`` mapping; ``env_configs["nemo_gym"]`` + supplies the Gym global config plus NeMo-RL detection knobs + (``invalid_tool_call_patterns``, ``thinking_tags``, ``num_gpu_nodes``). + base_urls: Per-DP-rank OpenAI-compatible server base URLs from the + generation backend. + model_name: Served model name the Gym rollouts should target. + enable_router_replay: Sets ``require_routed_experts`` on the + ``NemoGymConfig``. + + Returns: + The spun-up ``NemoGym`` Ray actor handle (``_spinup`` already awaited). + """ + nemo_gym_dict = dict(env_configs["nemo_gym"]) + + # NeMo-RL-side detection knobs are top-level NemoGymConfig fields + # (where the detector reads them), not part of Gym's global config. + invalid_tool_call_patterns = nemo_gym_dict.pop("invalid_tool_call_patterns", None) + thinking_tags = nemo_gym_dict.pop("thinking_tags", None) + + # Pass prebuilt cache + venv dirs through the global config so the gym reuses + # image-baked venvs instead of rebuilding them. + uv_cache_dir = get_nemo_gym_uv_cache_dir() + if uv_cache_dir is not None: + nemo_gym_dict.setdefault("uv_cache_dir", uv_cache_dir) + uv_venv_dir = get_nemo_gym_venv_dir() + if uv_venv_dir is not None: + nemo_gym_dict.setdefault("uv_venv_dir", uv_venv_dir) + + nemo_gym_cfg = NemoGymConfig( + model_name=model_name, + base_urls=base_urls, + invalid_tool_call_patterns=invalid_tool_call_patterns, + thinking_tags=thinking_tags, + require_routed_experts=enable_router_replay, + initial_global_config_dict=nemo_gym_dict, + ) + + nemo_gym_py_exec = get_actor_python_env("nemo_rl.environments.nemo_gym.NemoGym") + if nemo_gym_py_exec.startswith("uv"): + nemo_gym_py_exec = create_local_venv_on_each_node( + nemo_gym_py_exec, "nemo_rl.environments.nemo_gym.NemoGym" + ) + + nemo_gym_opts: dict[str, Any] = {} + if nemo_gym_dict.get("num_gpu_nodes", 0): + nemo_gym_opts["scheduling_strategy"] = NodeAffinitySchedulingStrategy( + node_id=ray.get_runtime_context().get_node_id(), + soft=True, + ) + nemo_gym_opts["runtime_env"] = { + "py_executable": nemo_gym_py_exec, + "env_vars": { + **os.environ, + "VIRTUAL_ENV": nemo_gym_py_exec, + "UV_PROJECT_ENVIRONMENT": nemo_gym_py_exec, + }, + } + + actor = NemoGym.options(**nemo_gym_opts).remote(nemo_gym_cfg) + ray.get(actor._spinup.remote()) + return actor diff --git a/tests/functional/L1_Functional_Tests_SingleController.sh b/tests/functional/L1_Functional_Tests_SingleController.sh index ea818aaf8c..2a6fad63ee 100755 --- a/tests/functional/L1_Functional_Tests_SingleController.sh +++ b/tests/functional/L1_Functional_Tests_SingleController.sh @@ -35,6 +35,7 @@ run_test() { } run_test fast uv run --no-sync bash ./tests/functional/grpo_dp_single_controller.sh +run_test fast uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller.sh cd ${PROJECT_ROOT}/tests if compgen -G ".coverage*" > /dev/null; then diff --git a/tests/functional/grpo_async_gym_single_controller.sh b/tests/functional/grpo_async_gym_single_controller.sh new file mode 100755 index 0000000000..54449013fd --- /dev/null +++ b/tests/functional/grpo_async_gym_single_controller.sh @@ -0,0 +1,114 @@ +#!/bin/bash +# SingleController + NeMo-Gym e2e smoke. Mirrors grpo_async_gym.sh but +# routes everything through the SC path (TransferQueue data plane + +# SingleControllerActor) instead of async_grpo_train. + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..) +# Mark the current repo as safe, since wandb fetches metadata about the repo +git config --global --add safe.directory $PROJECT_ROOT + +set -eou pipefail + +EXP_NAME=$(basename $0 .sh) +EXP_DIR=$SCRIPT_DIR/$EXP_NAME +LOG_DIR=$EXP_DIR/logs +JSON_METRICS=$EXP_DIR/metrics.json +RUN_LOG=$EXP_DIR/run.log +CHECKPOINT_DIR=$EXP_DIR/checkpoints +DATA_DIR=$EXP_DIR/data +export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} + +rm -rf $EXP_DIR $LOG_DIR +mkdir -p $EXP_DIR $LOG_DIR $CHECKPOINT_DIR $DATA_DIR + +# clean up checkpoint directory on exit +trap "rm -rf $CHECKPOINT_DIR" EXIT + +cd $PROJECT_ROOT + +# Follow nemo-gym instructions here to get this data: +# https://docs.nvidia.com/nemo/gym/0.1.0/tutorials/nemo-rl-grpo/setup.html#training-nemo-rl-grpo-setup +cd 3rdparty/Gym-workspace/Gym + +# We need HF_TOKEN to download the data from huggingface +if [[ ! -f env.yaml ]]; then + if [[ -z "${HF_TOKEN:-}" ]]; then + echo "[ERROR] HF_TOKEN is not set" + exit 1 + fi + echo "hf_token: $HF_TOKEN" >> env.yaml +fi + +uv run ng_prepare_data "+config_paths=[resources_servers/workplace_assistant/configs/workplace_assistant.yaml]" \ + +output_dirpath=data/workplace_assistant \ + +mode=train_preparation \ + +should_download=true \ + +data_source=huggingface +cd - + +# This trimming of the workplace assistant dataset is necessary b/c with all the tools the first prompt is >4000 tokens +# which will cause vllm to return nothing on the first prompt and crash RL. Since we want to keep this test short to +# smoke test, we trim all but the first tool +TRAIN_PATH=$DATA_DIR/workplace_assistant_train.jsonl +VALIDATION_PATH=$DATA_DIR/workplace_assistant_validation.jsonl +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/train.jsonl > $TRAIN_PATH +jq -c '.responses_create_params.tools |= (.[0:1])' 3rdparty/Gym-workspace/Gym/data/workplace_assistant/validation.jsonl > $VALIDATION_PATH + +uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ + $PROJECT_ROOT/examples/run_grpo_single_controller.py \ + --config $PROJECT_ROOT/examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml \ + policy.model_name=Qwen/Qwen3-0.6B \ + policy.dtensor_cfg.enabled=false \ + policy.megatron_cfg.enabled=true \ + policy.megatron_cfg.tensor_model_parallel_size=1 \ + policy.megatron_cfg.pipeline_model_parallel_size=1 \ + policy.megatron_cfg.expert_model_parallel_size=1 \ + policy.megatron_cfg.context_parallel_size=1 \ + policy.megatron_cfg.sequence_parallel=false \ + policy.generation.vllm_cfg.tensor_parallel_size=1 \ + policy.generation.vllm_cfg.async_engine=true \ + policy.max_total_sequence_length=512 \ + policy.generation.colocated.enabled=false \ + policy.generation.colocated.resources.num_nodes=1 \ + policy.generation.colocated.resources.gpus_per_node=1 \ + grpo.num_prompts_per_step=4 \ + grpo.num_generations_per_prompt=2 \ + grpo.max_num_steps=10 \ + grpo.val_period=-1 \ + grpo.val_at_start=false \ + policy.train_global_batch_size=8 \ + policy.train_micro_batch_size=1 \ + cluster.gpus_per_node=2 \ + loss_fn.reference_policy_kl_penalty=0.01 \ + grpo.skip_reference_policy_logprobs_calculation=false \ + loss_fn.use_importance_sampling_correction=true \ + logger.tensorboard_enabled=true \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=false \ + logger.monitor_gpus=true \ + checkpointing.enabled=false \ + data.train.data_path=$TRAIN_PATH \ + data.validation.data_path=$VALIDATION_PATH \ + ++data_plane.enabled=true \ + ++data_plane.impl=transfer_queue \ + ++data_plane.backend=simple \ + ++data_plane.storage_capacity=1000000 \ + ++data_plane.num_storage_units=2 \ + ++data_plane.claim_meta_poll_interval_s=0.5 \ + ++data_plane.global_segment_size=549755813888 \ + ++data_plane.local_buffer_size=68719476736 \ + ++async_rl.sampler.name=in_order \ + ++async_rl.sampler.max_lookahead_versions=0 \ + ++async_rl.min_groups_for_streaming_train=4 \ + ++async_rl.max_inflight_prompts=4 \ + ++async_rl.max_buffered_rollouts=4 \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# Observed to be between 0.8-1.3 +uv run tests/check_metrics.py $JSON_METRICS \ + 'median(data["train/gen_kl_error"]) < 1.3' \ + 'max(data["train/reward"]) > 0' diff --git a/tests/unit/single_controller/test_run_grpo_single_controller.py b/tests/unit/single_controller/test_run_grpo_single_controller.py index c1331a66f4..3a22bd44c8 100644 --- a/tests/unit/single_controller/test_run_grpo_single_controller.py +++ b/tests/unit/single_controller/test_run_grpo_single_controller.py @@ -31,6 +31,7 @@ def main_context(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: "draft": {"enabled": False}, "megatron_cfg": {"mtp_num_layers": 2}, }, + env={}, data_plane={"enabled": True}, logger={"log_dir": "/tmp/logs"}, checkpointing={"enabled": False}, diff --git a/tests/unit/single_controller/test_single_controller_setup.py b/tests/unit/single_controller/test_single_controller_setup.py index 3d6ba75aa5..4defb58607 100644 --- a/tests/unit/single_controller/test_single_controller_setup.py +++ b/tests/unit/single_controller/test_single_controller_setup.py @@ -282,15 +282,6 @@ def test_returns_actor_args(self, patched_factories): assert actor_args.partition_id == "rollout_data" assert actor_args.tq_buffer._partition_id == "rollout_data" - def test_nemo_gym_not_supported(self): - """SC path trips the nemo-gym guard until PR #3267 lands.""" - mc = _make_master_config() - with ( - patch.object(sc_setup_mod, "_should_use_nemo_gym", return_value=True), - pytest.raises(NotImplementedError, match="NeMo-Gym"), - ): - setup_single_controller(mc, MagicMock(pad_token_id=0)) - def test_env_handles_sourced_from_setup_response_data(self, patched_factories): """setup_response_data receives master_config.env and supplies env handles.""" math_env_cfg = {"some": "value"} @@ -400,3 +391,52 @@ def test_megatron_train_iters_not_set_when_disabled(self, patched_factories): setup_single_controller(mc, MagicMock(pad_token_id=0)) assert "train_iters" not in mc.policy.get("megatron_cfg", {}) + + def test_nemo_gym_wires_env_handle(self, patched_factories): + """When _should_use_nemo_gym is True the nemo-gym actor is spun up and stored.""" + mc = _make_master_config(colocated=True, backend="vllm") + mc.policy["generation"]["model_name"] = "test-model" + mc.policy["generation"]["stop_strings"] = None + mc.policy["generation"]["stop_token_ids"] = None + mc.policy["generation"]["top_k"] = None + patched_factories["setup_response_data"].return_value = ( + list(range(8)), + None, + ) + fake_gym_actor = MagicMock(name="nemo_gym_actor") + + with ( + patch.object(sc_setup_mod, "_should_use_nemo_gym", return_value=True), + patch.object( + sc_setup_mod, "spinup_nemo_gym_actor", return_value=fake_gym_actor + ) as mock_spinup, + patch.object(sc_setup_mod, "router_replay_enabled", return_value=False), + ): + actor_args = setup_single_controller(mc, MagicMock(pad_token_id=0)) + + mock_spinup.assert_called_once_with( + env_configs=mc.env, + base_urls=patched_factories[ + "_build_generation" + ].return_value.dp_openai_server_base_urls, + model_name="test-model", + enable_router_replay=False, + ) + assert actor_args.env_handles["nemo_gym"] is fake_gym_actor + + @pytest.mark.parametrize("backend", ["sglang", "megatron"]) + def test_nemo_gym_rejects_non_vllm_backend(self, patched_factories, backend): + """SC nemo-gym wiring only supports vLLM; every other backend must raise.""" + mc = _make_master_config(colocated=True, backend=backend) + patched_factories["setup_response_data"].return_value = ( + list(range(8)), + None, + ) + + with ( + patch.object(sc_setup_mod, "_should_use_nemo_gym", return_value=True), + patch.object(sc_setup_mod, "spinup_nemo_gym_actor") as mock_spinup, + pytest.raises(NotImplementedError, match="vllm"), + ): + setup_single_controller(mc, MagicMock(pad_token_id=0)) + mock_spinup.assert_not_called()