feat: Cherry pick efficiency PRs for Lag16 scaling

examples/nemo_gym/run_grpo_nemo_gym.py

@@ -36,7 +36,6 @@
     grpo_train,
     refit_policy_generation,
     setup,
-    shutdown_environments,
 )
 from nemo_rl.algorithms.utils import get_tokenizer
 from nemo_rl.data.utils import setup_response_data
@@ -45,6 +44,7 @@
     setup_nemo_gym_config,
     should_use_nemo_gym,
 )
+from nemo_rl.environments.utils import shutdown_environments
 from nemo_rl.experience.rollouts import run_nemo_gym_rollout_sync
 from nemo_rl.models.generation import configure_generation_config
 from nemo_rl.models.generation.vllm.config import materialize_vllm_video_config
@@ -271,6 +271,7 @@ def main() -> None:
     task_to_env = {"nemo_gym": nemo_gym}
     val_task_to_env = task_to_env
 
+    trainer_owns_environment_shutdown = False
     try:
         if is_trajectory_collection:
             collect_trajectories(
@@ -331,6 +332,7 @@ def main() -> None:
             print("🚀 Running synchronous GRPO training")
 
             # Run standard GRPO training
+            trainer_owns_environment_shutdown = True
             grpo_train(
                 policy,
                 policy_generation,
@@ -347,7 +349,8 @@ def main() -> None:
                 processor=processor,
             )
     finally:
-        shutdown_environments(task_to_env, val_task_to_env)
+        if not trainer_owns_environment_shutdown:
+            shutdown_environments(task_to_env, val_task_to_env)
         try:
             policy_generation.shutdown()
         except Exception as error:

examples/nemo_gym/run_nemo_gym_single_node_sanity_tests.sh

@@ -41,7 +41,16 @@ uv run python -c "import ray; ray.shutdown()"
 ./tests/run_unit.sh unit/environments/test_math_environment.py::test_math_env_step_basic
 
 # NeMo Gym integrates directly into NeMo RL as an Environment since that is the cleanest way. This tests the NeMo Gym integration logic and correctness.
-./tests/run_unit.sh unit/environments/test_nemo_gym.py::test_nemo_gym_sanity
+./tests/run_unit.sh \
+    unit/environments/test_nemo_gym.py::test_nemo_gym_sanity \
+    --nemo-gym-only
 
 # NeMo Gym uses a separate rollout loop inside grpo_train in NeMo RL. This tests the e2e rollout functionality and correctness.
-./tests/run_unit.sh unit/experience/test_rollouts.py::test_run_async_nemo_gym_rollout
+./tests/run_unit.sh \
+    unit/experience/test_rollouts.py::test_run_async_nemo_gym_rollout \
+    --nemo-gym-only
+
+# Exercise one pinned example from each L0 Gym environment through the real NeMo RL rollout boundary.
+./tests/run_unit.sh \
+    unit/environments/test_nemo_gym_rollout.py::test_l0_gym_environments_roll_out_through_nemo_rl \
+    --nemo-gym-only

examples/run_grpo.py

@@ -23,11 +23,11 @@
     MasterConfig,
     grpo_train,
     setup,
-    shutdown_environments,
 )
 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.utils import shutdown_environments
 from nemo_rl.models.generation import configure_generation_config
 from nemo_rl.utils.config import (
     load_config,
@@ -180,6 +180,7 @@ def _make_policy(**kwargs):
     print("=" * 60 + "\n", flush=True)
 
     try:
+        trainer_owns_environment_shutdown = False
         # Check if async mode is enabled
         if config.grpo.async_grpo.enabled:
             # Async GRPO does not support dynamic sampling, reward scaling, or reward shaping (DAPO features)
@@ -228,6 +229,7 @@ def _make_policy(**kwargs):
             # Two parallel synchronous trainers (verl-style — main_ppo.py vs
             # main_ppo_sync.py). data_plane.enabled selects which one runs.
             trainer = _select_trainer(master_config)
+            trainer_owns_environment_shutdown = trainer is grpo_train
             # grpo_train_sync defers checkpoint finalization to the checkpointer's
             # background threads; the context manager guarantees they are flushed on
             # exit. (grpo_train also flushes internally; shutdown() is idempotent.)
@@ -247,7 +249,8 @@ def _make_policy(**kwargs):
                     master_config,
                 )
     finally:
-        shutdown_environments(task_to_env, val_task_to_env)
+        if not trainer_owns_environment_shutdown:
+            shutdown_environments(task_to_env, val_task_to_env)
         try:
             policy_generation.shutdown()
         except Exception as error:

nemo_rl/algorithms/distillation.py

@@ -56,6 +56,7 @@
     get_nemo_gym_venv_dir,
     should_use_nemo_gym,
 )
+from nemo_rl.environments.utils import shutdown_environments
 from nemo_rl.experience.rollouts import (
     run_async_multi_turn_rollout,
     run_multi_turn_rollout,
@@ -674,7 +675,7 @@ def init_nemo_gym():
 # ===============================================================================
 
 
-def distillation_train(
+def _distillation_train_impl(
     student_policy: ColocatablePolicyInterface,
     teacher_policy: ColocatablePolicyInterface,
     student_generation: Optional[GenerationInterface],
@@ -1172,6 +1173,43 @@ def distillation_train(
     checkpointer.shutdown()
 
 
+@trace_fn(RLSpanGroup.JOB, "rl.distillation.job")
+def distillation_train(
+    student_policy: ColocatablePolicyInterface,
+    teacher_policy: ColocatablePolicyInterface,
+    student_generation: Optional[GenerationInterface],
+    dataloader: StatefulDataLoader,
+    val_dataloader: Optional[StatefulDataLoader],
+    tokenizer: TokenizerType,
+    loss_fn: DistillationLossFn,
+    task_to_env: dict[str, EnvironmentInterface],
+    val_task_to_env: Optional[dict[str, EnvironmentInterface]],
+    logger: Logger,
+    checkpointer: CheckpointManager,
+    distillation_save_state: DistillationSaveState,
+    master_config: MasterConfig,
+) -> None:
+    """Run distillation training and always tear down its environments."""
+    try:
+        _distillation_train_impl(
+            student_policy=student_policy,
+            teacher_policy=teacher_policy,
+            student_generation=student_generation,
+            dataloader=dataloader,
+            val_dataloader=val_dataloader,
+            tokenizer=tokenizer,
+            loss_fn=loss_fn,
+            task_to_env=task_to_env,
+            val_task_to_env=val_task_to_env,
+            logger=logger,
+            checkpointer=checkpointer,
+            distillation_save_state=distillation_save_state,
+            master_config=master_config,
+        )
+    finally:
+        shutdown_environments(task_to_env, val_task_to_env)
+
+
 def validate(
     policy_generation: GenerationInterface,
     val_dataloader: Optional[StatefulDataLoader],

nemo_rl/algorithms/grpo.py

@@ -86,6 +86,7 @@
 )
 from nemo_rl.environments.interfaces import EnvironmentInterface
 from nemo_rl.environments.nemo_gym import should_use_nemo_gym, spinup_nemo_gym_actor
+from nemo_rl.environments.utils import shutdown_environments
 from nemo_rl.experience.interfaces import (
     NEXT_NEMO_GYM_TASK_INDEX_KEY,
 )
@@ -402,35 +403,6 @@ def _needs_hf_refit_handshake(
     return not (nccl_reshard_refit_enabled and not colocated_inference)
 
 
-def shutdown_environments(
-    task_to_env: dict[str, EnvironmentInterface] | None,
-    val_task_to_env: dict[str, EnvironmentInterface] | None,
-) -> None:
-    """Shut down each unique environment actor before generation stops."""
-    seen_environment_handles: set[int] = set()
-    for environment_map in (task_to_env, val_task_to_env):
-        if environment_map is None:
-            continue
-        for task_name, environment in environment_map.items():
-            handle_id = id(environment)
-            if handle_id in seen_environment_handles:
-                continue
-            seen_environment_handles.add(handle_id)
-
-            print(f"🛑 Shutting down environment {task_name}...")
-            try:
-                ray.get(environment.shutdown.remote(), timeout=10)
-            except Exception as shutdown_error:
-                print(
-                    f"Environment {task_name} graceful shutdown failed: "
-                    f"{shutdown_error}"
-                )
-                try:
-                    ray.kill(environment)
-                except Exception as kill_error:
-                    print(f"Error stopping environment {task_name}: {kill_error}")
-
-
 def setup(
     master_config: MasterConfig,
     tokenizer: TokenizerType,
@@ -2698,7 +2670,7 @@ def _validation_early_stop_message(
     )
 
 
-def grpo_train(
+def _grpo_train_impl(
     policy: ColocatablePolicyInterface,
     policy_generation: Optional[GenerationInterface],
     wrapped_dataloader: StatefulDataLoader | MultipleDataloaderWrapper,
@@ -3838,6 +3810,43 @@ def grpo_train(
     checkpointer.shutdown()
 
 
+@trace_fn(RLSpanGroup.JOB, "rl.grpo.job")
+def grpo_train(
+    policy: ColocatablePolicyInterface,
+    policy_generation: Optional[GenerationInterface],
+    wrapped_dataloader: StatefulDataLoader | MultipleDataloaderWrapper,
+    val_dataloader: Optional[StatefulDataLoader],
+    tokenizer: TokenizerType,
+    loss_fn: LossFunction,
+    task_to_env: dict[str, EnvironmentInterface],
+    val_task_to_env: Optional[dict[str, EnvironmentInterface]],
+    logger: Logger,
+    checkpointer: CheckpointManager,
+    grpo_save_state: GRPOSaveState,
+    master_config: MasterConfig,
+    processor: Optional[AutoProcessor] = None,
+) -> None:
+    """Run GRPO training and always tear down its environments."""
+    try:
+        _grpo_train_impl(
+            policy=policy,
+            policy_generation=policy_generation,
+            wrapped_dataloader=wrapped_dataloader,
+            val_dataloader=val_dataloader,
+            tokenizer=tokenizer,
+            loss_fn=loss_fn,
+            task_to_env=task_to_env,
+            val_task_to_env=val_task_to_env,
+            logger=logger,
+            checkpointer=checkpointer,
+            grpo_save_state=grpo_save_state,
+            master_config=master_config,
+            processor=processor,
+        )
+    finally:
+        shutdown_environments(task_to_env, val_task_to_env)
+
+
 def validate(
     policy_generation: GenerationInterface,
     val_dataloader: Optional[StatefulDataLoader],

nemo_rl/environments/nemo_gym.py

@@ -42,6 +42,11 @@
     _get_node_ip_local,
 )
 from nemo_rl.environments.interfaces import EnvironmentInterface
+from nemo_rl.environments.nemo_gym_shards import (
+    SHARDING_CONFIG_KEYS,
+    ShardConfigError,
+    parse_shard_plan,
+)
 from nemo_rl.experience.failures import (
     GymTransportError,
     RolloutDataFailure,
@@ -53,6 +58,8 @@
 from nemo_rl.utils.timer import Timer
 from nemo_rl.utils.venvs import create_local_venv_on_each_node
 
+NEMO_GYM_GRACEFUL_SHUTDOWN_TIMEOUT_S = 120
+
 # Kept local so the Gym actor does not depend on model-config dtype resolution.
 # Must cover every name resolve_routed_experts_dtype can produce.
 _ROUTED_EXPERTS_DTYPES = {
@@ -485,7 +492,11 @@ def _attach_multimodal_data_to_user_message(
     )
 
 
-@ray.remote(max_restarts=-1, max_task_retries=-1)  # pragma: no cover
+# Fail fast rather than restart. The servers this actor owns are started in
+# _spinup, which Ray does not re-run after a restart, so a restarted actor is
+# permanently broken: _require_spinup() rejects every later rollout call, and
+# the caller never sees the RayActorError it is waiting for.
+@ray.remote(max_restarts=0, max_task_retries=0)  # pragma: no cover
 class NemoGym(EnvironmentInterface):
     """This environment class isn't really used for training. It's really meant as an integration wrapper around NeMo-Gym that hooks into the existing NeMo RL resource management via ray. So there is still one source of truth for resource management in NeMo RL."""
 
@@ -1029,13 +1040,16 @@ def _postprocess_nemo_gym_to_nemo_rl_result(
         return result
 
     def shutdown(self) -> None:
-        # Teardown runs in a finally block, so it must not turn a real training error
-        # into a confusing AttributeError from a never-spun-up (e.g. restarted) actor.
-        if self.rh is None:
-            return
-        run_helper = self.rh
-        self.rh = None
-        run_helper.shutdown()
+        """Stop the Gym servers. Safe to call more than once, and before spinup.
+
+        Teardown runs in a finally block and may be requested more than once.
+        RunHelper.shutdown() is not idempotent, so the handle is cleared before
+        it is used. A failure therefore cannot leave a live handle that a later
+        cleanup attempt invokes again.
+        """
+        rh, self.rh = self.rh, None
+        if rh is not None:
+            rh.shutdown()
 
     def step(self, message_log_batch, metadata):
         # This is not used since NeMo-Gym will handle the rollouts entirely.
@@ -1301,11 +1315,34 @@ def spinup_nemo_gym_actor(
     """
     nemo_gym_dict = dict(env_configs["nemo_gym"])
 
+    # Validate the shards block even though only single-actor creation is wired
+    # up so far, so a malformed or premature sharded config fails at setup with
+    # a precise message instead of silently running unsharded.
+    shard_plan = parse_shard_plan(nemo_gym_dict)
+    if shard_plan is not None:
+        raise ShardConfigError(
+            f"env.nemo_gym.shards defines {len(shard_plan.shards)} shards "
+            f"({', '.join(s.name for s in shard_plan.shards)}), but multi-actor "
+            f"creation is not wired up yet. Remove 'shards' and use "
+            f"'config_paths' to run this job on a single actor."
+        )
+
+    # NeMo-RL-only keys are consumed here and must never reach Gym: the merged
+    # config is serialized into every Gym child process, and unrecognized
+    # dict-shaped top-level keys are parsed as server instance configs.
+    for key in SHARDING_CONFIG_KEYS:
+        nemo_gym_dict.pop(key, None)
+
     # 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)
     tokenizer_config = nemo_gym_dict.pop("tokenizer_config", None)
+    port_range = {
+        key: value
+        for key in ("port_range_low", "port_range_high")
+        if (value := nemo_gym_dict.pop(key, None)) is not None
+    }
     # Same treatment for the multimodal knobs: NemoGymConfig declares them as
     # top-level fields, so populate them here instead of leaving the actor to
     # read them back out of Gym's global config dict.
@@ -1340,6 +1377,7 @@ def spinup_nemo_gym_actor(
         routed_experts_dtype=routed_experts_dtype,
         use_fastokens=use_fastokens,
         initial_global_config_dict=nemo_gym_dict,
+        **port_range,
         **multimodal_flags,
     )
 
@@ -1386,6 +1424,27 @@ def spinup_nemo_gym_actor(
         flush=True,
     )
     actor = NemoGym.options(**nemo_gym_opts).remote(nemo_gym_cfg)
-    ray.get(actor._spinup.remote())
-    ray.get(actor.set_tokenizer.remote(tokenizer))
+    try:
+        ray.get(actor._spinup.remote())
+        ray.get(actor.set_tokenizer.remote(tokenizer))
+    except Exception:
+        # _spinup can fail after RunHelper has started some Gym subprocesses.
+        # Ask the actor to reap anything it owns, then force-stop the actor as a
+        # final safety net. Cleanup errors must not hide the startup failure.
+        try:
+            ray.get(
+                actor.shutdown.remote(),
+                timeout=NEMO_GYM_GRACEFUL_SHUTDOWN_TIMEOUT_S,
+            )
+        except Exception as cleanup_error:
+            print(
+                f"Warning: NeMo-Gym actor cleanup after startup failure failed: {cleanup_error}"
+            )
+        try:
+            ray.kill(actor)
+        except Exception as kill_error:
+            print(
+                f"Warning: NeMo-Gym actor kill after startup failure failed: {kill_error}"
+            )
+        raise
     return actor