feat(sc): NeMo-Gym path (#3267)

Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: ruit <ruit@nvidia.com>
Signed-off-by: Terry Kong <terryk@nvidia.com>
Co-authored-by: ruit <ruit@nvidia.com>
Co-authored-by: Terry Kong <terryk@nvidia.com>

Author: 3 people

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,

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")

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 ────────────────────────────────────
 

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,
     )
 

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

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

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'

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},