From 15074e4ccfe626a71df6972d247e6bf394f77eef Mon Sep 17 00:00:00 2001 From: Youngeun Kwon Date: Mon, 31 Aug 2026 01:28:15 -0700 Subject: [PATCH 01/13] feat(single-controller): pause generation during refit Signed-off-by: Youngeun Kwon --- docs/guides/single-controller.md | 4 +- nemo_rl/algorithms/single_controller.py | 52 +++++++-- .../models/generation/vllm/vllm_generation.py | 20 ++-- .../models/generation/test_vllm_generation.py | 38 +++--- .../single_controller/test_checkpointing.py | 16 ++- .../single_controller/test_refit_recovery.py | 11 ++ .../test_single_controller_actor.py | 108 ++++++++++++++++-- .../single_controller/test_train_pump_e2e.py | 13 ++- 8 files changed, 211 insertions(+), 51 deletions(-) diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index edce5196e93..ec2e91b0739 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -197,7 +197,7 @@ The SC path splits the async-GRPO loop across a rollout pump and a train pump th 2. **Actor startup**: `SingleControllerActor` launches `_rollout_pump` and `_train_pump` concurrently as asyncio tasks; both share the same `TQReplayBuffer` and `StalenessSampler`. 3. **Rollout pump loop**: `sampler.admit` gates dispatch against the current trainer version (returning a `target_step` for `in_order`); the pump then reserves a buffer slot, drives `RolloutManager.generate_and_push`, and commits with the observed `start_weight` / `end_weight`. 4. **Train pump loop**: `sampler.evict` drops out-of-window groups, `sampler.select` picks the next batch, `_value_stage` and `_value_train` run the critic forward and its optimizer step on a PPO run, `_advantage_stage` computes advantages, and the TQPolicy split API runs one optimizer step per RL step on GRPO, or `ppo.ppo_epochs` of them on PPO. -5. **Weight sync**: after each optimizer step the pump bumps the trainer version, clears rollout permission, calls the weight synchronizer, and re-opens the rollout pump for the next version. +5. **Weight sync**: after each optimizer step the pump bumps the trainer version, clears rollout permission, asks the generation backend to pause in-flight work, calls the weight synchronizer, resumes generation, and re-opens the rollout pump for the next version. vLLM's native pause preserves active request state; backends without pause support warn and retain their existing refit behavior. ## Relation to Legacy Async GRPO @@ -224,7 +224,7 @@ Do not carry `max_num_epochs: -1` across either. [ppo.md](./ppo.md#asynchronous- | `enabled: true` | Implicit — SC is always async; use `sampler.max_lookahead_versions: 0` for sync semantics, `>= 1` for async | | `max_trajectory_age_steps: N` | `sampler.name: in_order` with `sampler.max_lookahead_versions: N` | | `warmup_generation_lead_steps` (PPO) | `sampler.warmup_lookahead_versions` — the lookahead to use while critic warmup is in progress | -| `recompute_kv_cache_after_weight_updates` | `recompute_kv_cache_after_weight_updates` (same) | +| `recompute_kv_cache_after_weight_updates` | `recompute_kv_cache_after_weight_updates` (same); when vLLM pauses natively, the preserved requests recompute their KV cache after refit | | `in_flight_weight_updates` | Always effectively true; `false`-equivalent behavior is not yet supported (drain-gate tracked in [issue #2625](https://github.com/NVIDIA-NeMo/RL/issues/2625)) | | *(no legacy equivalent — matches legacy full-batch train semantics)* | `min_groups_for_streaming_train: ${grpo.num_prompts_per_step}`, or `${ppo.num_prompts_per_step}` on a PPO run | | *(no legacy equivalent — matches legacy `max_trajectory_age + 1` batches in flight)* | `max_inflight_prompts: num_prompts_per_step × (max_lookahead_versions + 1)` | diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 405c9ec436a..92a3f662b22 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -2983,17 +2983,20 @@ async def _sync_weights( *, calibration_data: Optional[BatchedDataDict[Any]] = None, ) -> int: - """Pause new rollout dispatches, synchronize weights, resume. + """Pause rollout dispatch and generation, synchronize weights, resume. - SC owns the pause gate; in-flight generations continue through the - refit — vLLM V1 async engine supports weight updates during pending - requests. + SC owns the rollout-dispatch gate and asks every generation backend to + pause in-flight work through the common refit lifecycle hooks. Backends + without native pause support warn and retain their existing in-flight + update behavior. vLLM preserves request state while paused. Flow: 1. _rollout_permitted.clear() — no new dispatches - 2. Optionally calibrate FP8 KV-cache scales. - 3. weight_synchronizer.sync_weights(kv_scales=...) - 4. _rollout_permitted.set() — resume + 2. Reconcile refit membership and optionally calibrate FP8 KV-cache scales. + 3. generation.pause_generation_for_refit(clear_cache=...) + 4. weight_synchronizer.sync_weights(kv_scales=...) + 5. generation.resume_generation_after_refit() + 6. _rollout_permitted.set() — resume dispatch Args: calibration_data: Optional data used to calibrate FP8 KV-cache @@ -3013,8 +3016,6 @@ async def _sync_weights( else await self._abort_stale_inflight() ) - # TODO(#2625): Add drain-gate support during refit. - # Reconcile before the refit, not on a death event. The refit group is provably # idle here and every rank is synchronized, which is required because the # operations that change membership are themselves collectives. Doing it every @@ -3047,6 +3048,24 @@ async def _sync_weights( # once a shard was gone, because absent_shards() never empties again. await self._reconcile_refit_membership() + clear_cache = self._async_cfg.recompute_kv_cache_after_weight_updates + print("⏸️ Requesting generation pause before refit", flush=True) + generation_paused_for_refit = await asyncio.to_thread( + self._gen.pause_generation_for_refit, + clear_cache=clear_cache, + ) + if generation_paused_for_refit: + print( + f" {len(self._inflight_by_group_id)} in-flight rollout group(s) paused", + flush=True, + ) + else: + print( + " generation backend has no native pause support; " + "in-flight requests retain their existing refit behavior", + flush=True, + ) + try: await self._sync_weights_within(kv_scales, "first") except (RefitAborted, RayActorError) as failure: @@ -3088,11 +3107,22 @@ async def _sync_weights( # promoted inside its window, and everything below this must still run on # both paths. self._promote_refit_shards() - if self._async_cfg.recompute_kv_cache_after_weight_updates: + + print("▶️ Requesting generation resume after refit", flush=True) + generation_resumed_after_refit = await asyncio.to_thread( + self._gen.resume_generation_after_refit + ) + if generation_paused_for_refit and not generation_resumed_after_refit: + raise RuntimeError( + "Failed to resume generation after a successful refit pause" + ) + + if clear_cache and not generation_paused_for_refit: # to_thread, like every other call into the workers here. Run directly on # the loop this is a blocking Ray call, and a wedged generation worker would # freeze the event loop itself -- taking the watchdog, which is an asyncio - # task on that same loop, down with it. + # task on that same loop, down with it. A backend that paused natively already + # received clear_cache=True at pause time, so invalidating again is redundant. await asyncio.to_thread(self._gen.invalidate_kv_cache) elapsed = time.monotonic() - t0 diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index f15c2bf081e..9d6f58dc9f8 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -1460,11 +1460,13 @@ def pause_generation_for_refit(self, *, clear_cache: bool) -> bool: if not self.worker_group or not self.worker_group.workers: raise RuntimeError("Worker group is not initialized") - futures = self.worker_group.run_all_workers_single_data( - "pause_generation_async", - clear_cache=clear_cache, - run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"], - ) + # Address surviving DP leaders directly. SingleController may have rebuilt the + # refit communicator after losing a shard; walking the whole worker group here + # would call the dead Ray actor and undo that recovery before the refit starts. + futures = [ + worker.pause_generation_async.remote(clear_cache=clear_cache) + for worker in self._refit_leader_workers() + ] if not all(ray.get(futures)): raise RuntimeError("Failed to pause every async vLLM engine") return True @@ -1478,10 +1480,10 @@ def resume_generation_after_refit(self) -> bool: if not self.worker_group or not self.worker_group.workers: raise RuntimeError("Worker group is not initialized") - futures = self.worker_group.run_all_workers_single_data( - "resume_generation_async", - run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"], - ) + futures = [ + worker.resume_generation_async.remote() + for worker in self._refit_leader_workers() + ] if not all(ray.get(futures)): raise RuntimeError("Failed to resume every async vLLM engine") return True diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index 72a1ae68da7..35bba7e7e32 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -19,7 +19,7 @@ import types from copy import deepcopy from pathlib import Path -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, call import pytest import ray @@ -189,25 +189,26 @@ def test_vllm_generation_broadcasts_native_refit_pause_and_resume( ) -> None: generation = VllmGeneration.__new__(VllmGeneration) generation.cfg = {"vllm_cfg": {"async_engine": True}} - generation.worker_group = MagicMock() - generation.worker_group.workers = [object(), object()] - futures = [object(), object()] - generation.worker_group.run_all_workers_single_data.return_value = futures + generation.worker_group = MagicMock(workers=[object(), object()]) + leaders = [MagicMock(), MagicMock()] + pause_futures = [object(), object()] + resume_futures = [object(), object()] + for worker, pause_future, resume_future in zip( + leaders, pause_futures, resume_futures + ): + worker.pause_generation_async.remote.return_value = pause_future + worker.resume_generation_async.remote.return_value = resume_future + generation._refit_leader_workers = MagicMock(return_value=leaders) ray_get = MagicMock(side_effect=[[True, True], [True, True]]) monkeypatch.setattr(ray, "get", ray_get) assert generation.pause_generation_for_refit(clear_cache=True) assert generation.resume_generation_after_refit() - generation.worker_group.run_all_workers_single_data.assert_any_call( - "pause_generation_async", - clear_cache=True, - run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"], - ) - generation.worker_group.run_all_workers_single_data.assert_any_call( - "resume_generation_async", - run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"], - ) + for worker in leaders: + worker.pause_generation_async.remote.assert_called_once_with(clear_cache=True) + worker.resume_generation_async.remote.assert_called_once_with() + assert ray_get.call_args_list == [call(pause_futures), call(resume_futures)] def test_vllm_generation_rejects_partial_refit_pause_and_resume( @@ -215,12 +216,9 @@ def test_vllm_generation_rejects_partial_refit_pause_and_resume( ) -> None: generation = VllmGeneration.__new__(VllmGeneration) generation.cfg = {"vllm_cfg": {"async_engine": True}} - generation.worker_group = MagicMock() - generation.worker_group.workers = [object(), object()] - generation.worker_group.run_all_workers_single_data.return_value = [ - object(), - object(), - ] + generation.worker_group = MagicMock(workers=[object(), object()]) + leaders = [MagicMock(), MagicMock()] + generation._refit_leader_workers = MagicMock(return_value=leaders) monkeypatch.setattr( ray, "get", diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 69b5caada71..a04b5b03901 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -103,6 +103,20 @@ # ── fakes ──────────────────────────────────────────────────────────────────── +class _FakeGeneration: + """Unsupported generation backend for checkpointing-only actor tests.""" + + def pause_generation_for_refit(self, *, clear_cache: bool) -> bool: + del clear_cache + return False + + def resume_generation_after_refit(self) -> bool: + return False + + def invalidate_kv_cache(self) -> None: + pass + + class _FakeTrainer: """TQPolicy stand-in: train methods are no-ops, save_checkpoint records calls.""" @@ -541,7 +555,7 @@ def _make_actor_args( data_plane_checkpoint_metadata: Optional[DataPlaneCheckpointMetadata] = None, ) -> SingleControllerActorArgs: return SingleControllerActorArgs( - gen_handle=object(), + gen_handle=_FakeGeneration(), trainer_handle=trainer if trainer is not None else _FakeTrainer(), env_handles={}, train_cluster=None, # type: ignore[arg-type] diff --git a/tests/unit/single_controller/test_refit_recovery.py b/tests/unit/single_controller/test_refit_recovery.py index 703adf9390f..24d7eb00d98 100644 --- a/tests/unit/single_controller/test_refit_recovery.py +++ b/tests/unit/single_controller/test_refit_recovery.py @@ -132,6 +132,8 @@ def _make_controller( ctrl._gen = SimpleNamespace( requires_kv_scale_sync=False, invalidate_kv_cache=MagicMock(), + pause_generation_for_refit=MagicMock(return_value=True), + resume_generation_after_refit=MagicMock(return_value=True), worker_group=SimpleNamespace( get_dp_leader_worker_idx=lambda shard: shard, workers=[ @@ -193,6 +195,15 @@ def test_the_retry_runs_against_the_smaller_fleet(self): assert sync.sync_calls == 2 assert sync.absent_at_retry == [0] + def test_generation_stays_paused_across_the_recovery_retry(self): + ctrl, _, sync = _make_controller(ABORTED) + + asyncio.run(ctrl._sync_weights()) + + assert sync.sync_calls == 2 + ctrl._gen.pause_generation_for_refit.assert_called_once_with(clear_cache=False) + ctrl._gen.resume_generation_after_refit.assert_called_once_with() + def test_survivors_are_pulled_from_service_then_given_back(self): """Partial weights must not serve -- and must not be stranded either. diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index d9c7d359b21..fe9e0fee221 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -432,13 +432,26 @@ def test_retune_is_a_noop_off_the_ppo_path(self): @pytest.mark.parametrize( - ("recompute_kv_cache", "expected_invalidation_calls"), - [(False, 0), (True, 1)], + ( + "generation_pause_supported", + "recompute_kv_cache", + "expected_invalidation_calls", + ), + [ + (True, False, 0), + (True, True, 0), + (False, False, 0), + (False, True, 1), + ], ) +@pytest.mark.parametrize("use_nemo_gym", [False, True]) def test_sync_weights_honors_recompute_kv_cache_config( + generation_pause_supported: bool, recompute_kv_cache: bool, expected_invalidation_calls: int, + use_nemo_gym: bool, ) -> None: + events: list[str] = [] controller_cls = SingleControllerActor.__ray_metadata__.modified_class ctrl = object.__new__(controller_cls) ctrl._async_cfg = AsyncRLConfig( @@ -449,21 +462,50 @@ def test_sync_weights_honors_recompute_kv_cache_config( # No fleet health: _sync_weights reconciles refit membership first, and with no # monitor there is nothing to reconcile. ctrl._gen_fleet = None - ctrl._weight_synchronizer = SimpleNamespace(sync_weights=MagicMock()) + ctrl._weight_synchronizer = SimpleNamespace( + sync_weights=MagicMock(side_effect=lambda **_: events.append("sync")) + ) ctrl._gen = SimpleNamespace( - invalidate_kv_cache=MagicMock(), + invalidate_kv_cache=MagicMock(side_effect=lambda: events.append("invalidate")), + pause_generation_for_refit=MagicMock( + side_effect=lambda **_: ( + events.append("pause") or generation_pause_supported + ) + ), requires_kv_scale_sync=False, + resume_generation_after_refit=MagicMock( + side_effect=lambda: (events.append("resume") or generation_pause_supported) + ), ) ctrl._inflight_by_group_id = {} ctrl._rollout_recovery_enabled = False - # env={} -> should_use_nemo_gym is False, so _sync_weights takes the native - # abort path (empty registry -> no-op) instead of the gym gate. - ctrl._master_config = SimpleNamespace(env={}) + ctrl._master_config = SimpleNamespace( + env={"should_use_nemo_gym": use_nemo_gym}, + policy={ + "generation": { + "backend": "vllm", + "vllm_cfg": { + "async_engine": True, + "expose_http_server": True, + }, + } + }, + ) asyncio.run(ctrl._sync_weights()) ctrl._weight_synchronizer.sync_weights.assert_called_once_with(kv_scales=None) + ctrl._gen.pause_generation_for_refit.assert_called_once_with( + clear_cache=recompute_kv_cache + ) + ctrl._gen.resume_generation_after_refit.assert_called_once_with() assert ctrl._gen.invalidate_kv_cache.call_count == expected_invalidation_calls + assert events == [ + "pause", + "sync", + "resume", + *(["invalidate"] if expected_invalidation_calls else []), + ] assert ctrl._rollout_permitted.is_set() @@ -479,7 +521,9 @@ def test_sync_weights_calibrates_and_forwards_fp8_kv_scales() -> None: ctrl._weight_synchronizer = SimpleNamespace(sync_weights=MagicMock()) ctrl._gen = SimpleNamespace( invalidate_kv_cache=MagicMock(), + pause_generation_for_refit=MagicMock(return_value=True), requires_kv_scale_sync=True, + resume_generation_after_refit=MagicMock(return_value=True), ) ctrl._trainer = SimpleNamespace( calibrate_qkv_fp8_scales=MagicMock(return_value={"layers": {"layer.0": 0.5}}) @@ -507,6 +551,56 @@ def test_sync_weights_calibrates_and_forwards_fp8_kv_scales() -> None: ) +def test_sync_weights_keeps_dispatch_paused_when_generation_resume_fails() -> None: + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + ctrl = object.__new__(controller_cls) + ctrl._async_cfg = AsyncRLConfig() + ctrl._rollout_permitted = asyncio.Event() + ctrl._rollout_permitted.set() + ctrl._gen_fleet = None + ctrl._weight_synchronizer = SimpleNamespace(sync_weights=MagicMock()) + ctrl._gen = SimpleNamespace( + pause_generation_for_refit=MagicMock(return_value=True), + requires_kv_scale_sync=False, + resume_generation_after_refit=MagicMock(return_value=False), + ) + ctrl._inflight_by_group_id = {} + ctrl._rollout_recovery_enabled = False + ctrl._master_config = SimpleNamespace(env={}) + + with pytest.raises(RuntimeError, match="successful refit pause"): + asyncio.run(ctrl._sync_weights()) + + ctrl._weight_synchronizer.sync_weights.assert_called_once_with(kv_scales=None) + assert not ctrl._rollout_permitted.is_set() + + +def test_sync_weights_does_not_resume_generation_after_failed_refit() -> None: + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + ctrl = object.__new__(controller_cls) + ctrl._async_cfg = AsyncRLConfig() + ctrl._rollout_permitted = asyncio.Event() + ctrl._rollout_permitted.set() + ctrl._gen_fleet = None + ctrl._weight_synchronizer = SimpleNamespace( + sync_weights=MagicMock(side_effect=RuntimeError("refit failed")) + ) + ctrl._gen = SimpleNamespace( + pause_generation_for_refit=MagicMock(return_value=True), + requires_kv_scale_sync=False, + resume_generation_after_refit=MagicMock(return_value=True), + ) + ctrl._inflight_by_group_id = {} + ctrl._rollout_recovery_enabled = False + ctrl._master_config = SimpleNamespace(env={}) + + with pytest.raises(RuntimeError, match="refit failed"): + asyncio.run(ctrl._sync_weights()) + + ctrl._gen.resume_generation_after_refit.assert_not_called() + assert not ctrl._rollout_permitted.is_set() + + class _AdvantageDataPlane: def __init__(self, data: TensorDict) -> None: self._data = data diff --git a/tests/unit/single_controller/test_train_pump_e2e.py b/tests/unit/single_controller/test_train_pump_e2e.py index f489d722e59..9570d00893f 100644 --- a/tests/unit/single_controller/test_train_pump_e2e.py +++ b/tests/unit/single_controller/test_train_pump_e2e.py @@ -173,6 +173,17 @@ def compute_advantage( return rewards.detach().unsqueeze(-1).expand_as(mask).clone() +class _FakeGeneration: + """Unsupported generation backend for this training-pump-only test.""" + + def pause_generation_for_refit(self, *, clear_cache: bool) -> bool: + del clear_cache + return False + + def resume_generation_after_refit(self) -> bool: + return False + + @ray.remote(num_cpus=0) # pragma: no cover class _CallLog: """Ordered append-only log the fakes below write to.""" @@ -348,7 +359,7 @@ def test_train_pump_drives_mcore_training_step( ) actor_args = SingleControllerActorArgs( - gen_handle=None, + gen_handle=_FakeGeneration(), trainer_handle=trainer, env_handles={}, train_cluster=None, # type: ignore[arg-type] From 1f68285998a58fa59fb994b847b79ec8fed2b54f Mon Sep 17 00:00:00 2001 From: Youngeun Kwon Date: Mon, 31 Aug 2026 14:15:43 -0700 Subject: [PATCH 02/13] fix(single-controller): enforce refit cache correctness Signed-off-by: Youngeun Kwon --- docs/guides/single-controller.md | 2 + nemo_rl/algorithms/single_controller.py | 104 +++++++++++------- nemo_rl/models/generation/interfaces.py | 12 +- .../models/generation/vllm/vllm_generation.py | 4 +- nemo_rl/models/generation/vllm/vllm_worker.py | 5 +- .../generation/vllm/vllm_worker_async.py | 5 +- .../models/generation/test_vllm_generation.py | 64 +++++++++++ .../single_controller/test_refit_recovery.py | 21 +++- .../test_single_controller_actor.py | 56 +++++++++- 9 files changed, 221 insertions(+), 52 deletions(-) diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index ec2e91b0739..171ea039190 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -199,6 +199,8 @@ The SC path splits the async-GRPO loop across a rollout pump and a train pump th 4. **Train pump loop**: `sampler.evict` drops out-of-window groups, `sampler.select` picks the next batch, `_value_stage` and `_value_train` run the critic forward and its optimizer step on a PPO run, `_advantage_stage` computes advantages, and the TQPolicy split API runs one optimizer step per RL step on GRPO, or `ppo.ppo_epochs` of them on PPO. 5. **Weight sync**: after each optimizer step the pump bumps the trainer version, clears rollout permission, asks the generation backend to pause in-flight work, calls the weight synchronizer, resumes generation, and re-opens the rollout pump for the next version. vLLM's native pause preserves active request state; backends without pause support warn and retain their existing refit behavior. +Freezing a vLLM request does not stop its existing end-to-end generation timeout. Size `NRL_VLLM_ASYNC_TIMEOUT_SECONDS` and any generation-router `backend_timeout_s` to cover the longest expected generation plus time spent paused for refit. + ## Relation to Legacy Async GRPO The [legacy async GRPO](./async-grpo.md) (`grpo.async_grpo.enabled: true` under `run_grpo.py`) and the SC path both target the same async training problem but partition responsibilities differently: diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 92a3f662b22..4ae58d6adb7 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -2910,25 +2910,27 @@ async def _save_checkpoint( _REFIT_UNWIND_GRACE_S = 60.0 def _refit_await_budget_s(self) -> Optional[float]: - """How long to wait for the refit before giving up, or None to wait forever.""" + """How long to wait for a refit lifecycle call, or None to wait forever.""" deadline = self._async_cfg.generation_fleet_health.refit_timeout_s return None if deadline is None else deadline + self._REFIT_UNWIND_GRACE_S - async def _sync_weights_within(self, kv_scales, what: str) -> None: - """Run the refit off-loop, and stop waiting if it outlives the deadline. + async def _run_refit_call_within( + self, call: Callable[[], Any], *, what: str + ) -> Any: + """Run a blocking refit lifecycle call off-loop with the refit deadline. - WHY THIS IS NEEDED ON TOP OF EVERY WORKER-SIDE BOUND. A frozen-but-alive rank is a - Ray actor that never answers, and Ray puts no timeout on an actor call. So the - controller's await never resolves however well the workers behave. Job 6508251 - measured that end state: the deadline fired, the workers aborted, the trainers - returned, every actor was idle -- and the run still sat for 1800s because this - await had no bound. It is the last unbounded wait on the refit path, and unlike the - others it is not in NCCL or CUDA. + WHY THIS IS NEEDED ON TOP OF EVERY WORKER-SIDE BOUND. A frozen-but-alive rank is + a Ray actor that never answers, and Ray puts no timeout on an actor call. Pause, + weight sync, resume, and fallback cache invalidation all cross that boundary, so + every one must use this wrapper. Job 6508251 measured the old end state: the + deadline fired, the workers aborted, the trainers returned, every actor was idle + -- and the run still sat for 1800s because the controller await had no bound. - Timing out raises RefitAborted so this joins the existing recovery path rather than - inventing a second one: the caller reconciles membership and retries once. By then - the fleet probe has usually condemned the silent shard, so the rebuild can exclude - it and the retry may genuinely succeed. + Timing out raises RefitAborted. The weight-sync caller routes that through its + existing recovery path; pause/resume deliberately remain fail-stop because a + partial multi-engine lifecycle call leaves the controller unable to prove which + preserved requests are paused or cache-cleared. In every case the dispatch gate + stays closed rather than allowing requests into uncertain engine state. A DEDICATED DAEMON THREAD, not asyncio.to_thread. to_thread runs on the default ThreadPoolExecutor, whose workers are non-daemon and are joined at interpreter @@ -2936,18 +2938,14 @@ async def _sync_weights_within(self, kv_scales, what: str) -> None: a wedge in the refit for a wedge on the way out. asyncio.wait_for cannot cancel a running thread either way; this only controls whether the orphan can block exit. - The orphan is a real consequence, not a free win. If the retry succeeds the run - continues with a thread still parked inside the old sync_weights, and if that rank - were ever resumed it would wake up holding a communicator that has since been - replaced. Acceptable against a guaranteed 30-minute stall, and it is why this path - is bounded-failure-first rather than resume-and-forget. + The orphan is a real consequence, not a free win. A timed-out call can leave a + thread parked on the old Ray future. The thread is therefore a daemon: bounded + failure is preferable to trading a runtime wedge for an interpreter-shutdown + wedge. """ budget_s = self._refit_await_budget_s() if budget_s is None: - await asyncio.to_thread( - self._weight_synchronizer.sync_weights, kv_scales=kv_scales - ) - return + return await asyncio.to_thread(call) loop = asyncio.get_running_loop() settled: asyncio.Future = loop.create_future() @@ -2960,24 +2958,34 @@ def _settle(setter, value) -> None: def _run() -> None: try: - self._weight_synchronizer.sync_weights(kv_scales=kv_scales) + result = call() except BaseException as exc: # noqa: BLE001 - re-raised on the loop below loop.call_soon_threadsafe(_settle, settled.set_exception, exc) else: - loop.call_soon_threadsafe(_settle, settled.set_result, None) + loop.call_soon_threadsafe(_settle, settled.set_result, result) - threading.Thread(target=_run, name=f"sc-refit-{what}", daemon=True).start() + thread_label = what.replace(" ", "-") + threading.Thread( + target=_run, name=f"sc-refit-{thread_label}", daemon=True + ).start() try: - await asyncio.wait_for(settled, budget_s) + return await asyncio.wait_for(settled, budget_s) except asyncio.TimeoutError: raise RefitAborted( - f"the {what} refit did not return within {budget_s}s. Every worker-side " - "bound has passed, so a generation rank is most likely alive but not " - "answering -- a Ray actor call has no timeout of its own, so waiting here " - "is unbounded. Giving up so the fleet can be reconciled and retried." + f"{what} did not return within {budget_s}s. A generation rank is most " + "likely alive but not answering -- a Ray actor call has no timeout of " + "its own, so waiting here is unbounded. Giving up with rollout dispatch " + "still paused." ) from None + async def _sync_weights_within(self, kv_scales, what: str) -> None: + """Run one weight sync through the bounded refit lifecycle wrapper.""" + await self._run_refit_call_within( + partial(self._weight_synchronizer.sync_weights, kv_scales=kv_scales), + what=f"the {what} refit", + ) + async def _sync_weights( self, *, @@ -3050,9 +3058,12 @@ async def _sync_weights( clear_cache = self._async_cfg.recompute_kv_cache_after_weight_updates print("⏸️ Requesting generation pause before refit", flush=True) - generation_paused_for_refit = await asyncio.to_thread( - self._gen.pause_generation_for_refit, - clear_cache=clear_cache, + generation_paused_for_refit = await self._run_refit_call_within( + partial( + self._gen.pause_generation_for_refit, + clear_cache=clear_cache, + ), + what="generation pause before refit", ) if generation_paused_for_refit: print( @@ -3109,8 +3120,9 @@ async def _sync_weights( self._promote_refit_shards() print("▶️ Requesting generation resume after refit", flush=True) - generation_resumed_after_refit = await asyncio.to_thread( - self._gen.resume_generation_after_refit + generation_resumed_after_refit = await self._run_refit_call_within( + self._gen.resume_generation_after_refit, + what="generation resume after refit", ) if generation_paused_for_refit and not generation_resumed_after_refit: raise RuntimeError( @@ -3118,12 +3130,20 @@ async def _sync_weights( ) if clear_cache and not generation_paused_for_refit: - # to_thread, like every other call into the workers here. Run directly on - # the loop this is a blocking Ray call, and a wedged generation worker would - # freeze the event loop itself -- taking the watchdog, which is an asyncio - # task on that same loop, down with it. A backend that paused natively already - # received clear_cache=True at pause time, so invalidating again is redundant. - await asyncio.to_thread(self._gen.invalidate_kv_cache) + # A backend that paused natively already received clear_cache=True at pause + # time. Other backends must explicitly confirm equivalent engine-side cache + # handling; merely clearing reusable prefix entries does not invalidate KV + # blocks still owned by running requests. + invalidated = await self._run_refit_call_within( + self._gen.invalidate_kv_cache, + what="generation KV-cache invalidation after refit", + ) + if not invalidated: + raise NotImplementedError( + "recompute_kv_cache_after_weight_updates=True requires the " + "generation backend to clear active-request KV state during a " + "native refit pause or confirm equivalent engine-side handling" + ) elapsed = time.monotonic() - t0 print(f" _sync_weights: sync done in {elapsed:.3f}s", flush=True) diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index bfd62367364..fa1dfe610d9 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -557,9 +557,17 @@ def attach_fleet_health(self, monitor: Any, selector: Any) -> None: f"{type(self).__name__} generation backend" ) - # Optional hook; backends may override to invalidate any reusable caches - # (e.g., vLLM prefix/KV caches) after weight updates. def invalidate_kv_cache(self) -> bool: + """Confirm stale KV state cannot be reused after a weight update. + + Backends may implement this by invalidating active-request KV state directly or + by performing equivalent engine-side handling atomically with refit. Clearing + only unreferenced prefix-cache entries is insufficient while requests are active. + + Returns: + True only when the backend guarantees the recompute contract; False when it + is unsupported or could not be completed. + """ return False def pause_generation_for_refit(self, *, clear_cache: bool) -> bool: diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index 2e253df2411..63e31271565 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -1450,8 +1450,8 @@ def invalidate_kv_cache(self) -> bool: method_name, run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"], ) - results = ray.get(futures) - return all(result for result in results if result is not None) + results = [result for result in ray.get(futures) if result is not None] + return bool(results) and all(results) except Exception as e: print(f"Error invalidating vLLM caches: {e}") return False diff --git a/nemo_rl/models/generation/vllm/vllm_worker.py b/nemo_rl/models/generation/vllm/vllm_worker.py index ba9ad5fb837..f1e1ded7480 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker.py +++ b/nemo_rl/models/generation/vllm/vllm_worker.py @@ -1412,7 +1412,7 @@ def nccl_reshard_refit(self, refit_timeout_s: Optional[float] = None) -> bool: traceback.print_exc() return False - def reset_prefix_cache(self): + def reset_prefix_cache(self) -> bool: """Reset the prefix cache of vLLM engine.""" assert self.llm is not None, ( "Attempting to reset prefix cache with either an uninitialized vLLM or non-model-owner" @@ -1423,9 +1423,10 @@ def reset_prefix_cache(self): "reset_prefix_cache can only be used with async_engine=False. Use reset_prefix_cache_async instead." ) - self.llm.llm_engine.reset_prefix_cache() + reset_successful = self.llm.llm_engine.reset_prefix_cache() gc.collect() torch.cuda.empty_cache() + return reset_successful def sleep(self): """Put the vLLM engine to sleep.""" diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index 01076d19481..b30c900bc6c 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -1587,7 +1587,7 @@ async def nccl_reshard_refit_async( traceback.print_exc() return False - async def reset_prefix_cache_async(self): + async def reset_prefix_cache_async(self) -> bool: """Async version of reset_prefix_cache.""" assert self.llm is not None, ( "Attempting to reset prefix cache with either an uninitialized vLLM or non-model-owner" @@ -1598,9 +1598,10 @@ async def reset_prefix_cache_async(self): "reset_prefix_cache_async can only be used with async_engine=True. Use reset_prefix_cache instead." ) - await self.llm.reset_prefix_cache() + reset_successful = await self.llm.reset_prefix_cache() gc.collect() torch.cuda.empty_cache() + return reset_successful async def pause_generation_async(self, *, clear_cache: bool) -> bool: """Pause vLLM generation for an in-flight weight update.""" diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index f6bc2d5829b..4faa82cdf7c 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -184,6 +184,70 @@ async def test_async_vllm_worker_uses_native_keep_pause_and_resume() -> None: worker.llm.resume_generation.assert_awaited_once_with() +@pytest.mark.asyncio +async def test_async_vllm_worker_propagates_prefix_reset_result( + monkeypatch: pytest.MonkeyPatch, +) -> None: + worker = VllmAsyncGenerationWorkerImpl.__new__(VllmAsyncGenerationWorkerImpl) + worker.cfg = {"vllm_cfg": {"async_engine": True}} + worker.llm = MagicMock(reset_prefix_cache=AsyncMock(return_value=False)) + empty_cache = MagicMock() + monkeypatch.setattr(torch.cuda, "empty_cache", empty_cache) + + assert await worker.reset_prefix_cache_async() is False + + worker.llm.reset_prefix_cache.assert_awaited_once_with() + empty_cache.assert_called_once_with() + + +def test_sync_vllm_worker_propagates_prefix_reset_result( + monkeypatch: pytest.MonkeyPatch, +) -> None: + worker = VllmGenerationWorkerImpl.__new__(VllmGenerationWorkerImpl) + worker.cfg = {"vllm_cfg": {"async_engine": False}} + reset_prefix_cache = MagicMock(return_value=False) + worker.llm = types.SimpleNamespace( + llm_engine=types.SimpleNamespace(reset_prefix_cache=reset_prefix_cache) + ) + empty_cache = MagicMock() + monkeypatch.setattr(torch.cuda, "empty_cache", empty_cache) + + assert worker.reset_prefix_cache() is False + + reset_prefix_cache.assert_called_once_with() + empty_cache.assert_called_once_with() + + +@pytest.mark.parametrize( + ("worker_results", "expected"), + [ + ([True, None], True), + ([False, None], False), + ([None, None], False), + ], +) +def test_vllm_generation_reports_prefix_reset_result( + monkeypatch: pytest.MonkeyPatch, + worker_results: list[bool | None], + expected: bool, +) -> None: + generation = VllmGeneration.__new__(VllmGeneration) + generation.cfg = {"vllm_cfg": {"async_engine": True}} + generation.worker_group = MagicMock() + futures = [object(), object()] + generation.worker_group.run_all_workers_single_data.return_value = futures + ray_get = MagicMock(return_value=worker_results) + monkeypatch.setattr(ray, "get", ray_get) + + assert generation.invalidate_kv_cache() is expected + + generation.worker_group.run_all_workers_single_data.assert_called_once_with( + "reset_prefix_cache_async", + run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"], + ) + ray_get.assert_called_once_with(futures) + + def test_vllm_generation_broadcasts_native_refit_pause_and_resume( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/unit/single_controller/test_refit_recovery.py b/tests/unit/single_controller/test_refit_recovery.py index 24d7eb00d98..1d1e1a50cd6 100644 --- a/tests/unit/single_controller/test_refit_recovery.py +++ b/tests/unit/single_controller/test_refit_recovery.py @@ -197,12 +197,16 @@ def test_the_retry_runs_against_the_smaller_fleet(self): def test_generation_stays_paused_across_the_recovery_retry(self): ctrl, _, sync = _make_controller(ABORTED) + sync_calls_when_resumed: list[int] = [] + ctrl._gen.resume_generation_after_refit = MagicMock( + side_effect=lambda: sync_calls_when_resumed.append(sync.sync_calls) or True + ) asyncio.run(ctrl._sync_weights()) assert sync.sync_calls == 2 ctrl._gen.pause_generation_for_refit.assert_called_once_with(clear_cache=False) - ctrl._gen.resume_generation_after_refit.assert_called_once_with() + assert sync_calls_when_resumed == [2] def test_survivors_are_pulled_from_service_then_given_back(self): """Partial weights must not serve -- and must not be stranded either. @@ -479,6 +483,21 @@ def test_a_refit_that_never_returns_gives_up(self): never.set() assert elapsed < 10.0, "it must give up, not wait on a rank that never answers" + def test_a_generation_lifecycle_call_uses_the_same_bound(self): + never = threading.Event() + ctrl = self._ctrl(0.1, lambda **_: None) + ctrl._REFIT_UNWIND_GRACE_S = 0.1 + + with pytest.raises(RefitAborted, match="generation pause before refit"): + asyncio.run( + ctrl._run_refit_call_within( + never.wait, + what="generation pause before refit", + ) + ) + + never.set() + def test_giving_up_is_recognised_as_an_abort(self): """So it lands in the existing `except (RefitAborted, RayActorError)` recovery.""" never = threading.Event() diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index fe9e0fee221..21c0d72c8b9 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -466,7 +466,9 @@ def test_sync_weights_honors_recompute_kv_cache_config( sync_weights=MagicMock(side_effect=lambda **_: events.append("sync")) ) ctrl._gen = SimpleNamespace( - invalidate_kv_cache=MagicMock(side_effect=lambda: events.append("invalidate")), + invalidate_kv_cache=MagicMock( + side_effect=lambda: (events.append("invalidate") or True) + ), pause_generation_for_refit=MagicMock( side_effect=lambda **_: ( events.append("pause") or generation_pause_supported @@ -509,6 +511,31 @@ def test_sync_weights_honors_recompute_kv_cache_config( assert ctrl._rollout_permitted.is_set() +def test_sync_weights_rejects_unconfirmed_inflight_kv_recomputation() -> None: + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + ctrl = object.__new__(controller_cls) + ctrl._async_cfg = AsyncRLConfig(recompute_kv_cache_after_weight_updates=True) + ctrl._rollout_permitted = asyncio.Event() + ctrl._rollout_permitted.set() + ctrl._gen_fleet = None + ctrl._weight_synchronizer = SimpleNamespace(sync_weights=MagicMock()) + ctrl._gen = SimpleNamespace( + invalidate_kv_cache=MagicMock(return_value=False), + pause_generation_for_refit=MagicMock(return_value=False), + requires_kv_scale_sync=False, + resume_generation_after_refit=MagicMock(return_value=False), + ) + ctrl._inflight_by_group_id = {} + ctrl._rollout_recovery_enabled = False + ctrl._master_config = SimpleNamespace(env={}) + + with pytest.raises(NotImplementedError, match="active-request KV state"): + asyncio.run(ctrl._sync_weights()) + + ctrl._gen.invalidate_kv_cache.assert_called_once_with() + assert not ctrl._rollout_permitted.is_set() + + def test_sync_weights_calibrates_and_forwards_fp8_kv_scales() -> None: controller_cls = SingleControllerActor.__ray_metadata__.modified_class ctrl = object.__new__(controller_cls) @@ -575,6 +602,33 @@ def test_sync_weights_keeps_dispatch_paused_when_generation_resume_fails() -> No assert not ctrl._rollout_permitted.is_set() +def test_sync_weights_keeps_dispatch_paused_when_generation_pause_fails() -> None: + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + ctrl = object.__new__(controller_cls) + ctrl._async_cfg = AsyncRLConfig() + ctrl._rollout_permitted = asyncio.Event() + ctrl._rollout_permitted.set() + ctrl._gen_fleet = None + ctrl._weight_synchronizer = SimpleNamespace(sync_weights=MagicMock()) + ctrl._gen = SimpleNamespace( + pause_generation_for_refit=MagicMock( + side_effect=RuntimeError("Failed to pause every async vLLM engine") + ), + requires_kv_scale_sync=False, + resume_generation_after_refit=MagicMock(return_value=True), + ) + ctrl._inflight_by_group_id = {} + ctrl._rollout_recovery_enabled = False + ctrl._master_config = SimpleNamespace(env={}) + + with pytest.raises(RuntimeError, match="pause every async vLLM engine"): + asyncio.run(ctrl._sync_weights()) + + ctrl._weight_synchronizer.sync_weights.assert_not_called() + ctrl._gen.resume_generation_after_refit.assert_not_called() + assert not ctrl._rollout_permitted.is_set() + + def test_sync_weights_does_not_resume_generation_after_failed_refit() -> None: controller_cls = SingleControllerActor.__ray_metadata__.modified_class ctrl = object.__new__(controller_cls) From 463f247737b65a7257a5f3dfaec94b1b3881cdbf Mon Sep 17 00:00:00 2001 From: Youngeun Kwon Date: Mon, 31 Aug 2026 15:55:58 -0700 Subject: [PATCH 03/13] refactor(single-controller): simplify refit pause lifecycle Signed-off-by: Youngeun Kwon --- nemo_rl/algorithms/single_controller.py | 89 ++++++++----------- .../models/generation/vllm/vllm_generation.py | 4 +- nemo_rl/models/generation/vllm/vllm_worker.py | 5 +- .../generation/vllm/vllm_worker_async.py | 5 +- .../models/generation/test_vllm_generation.py | 64 ------------- .../single_controller/test_refit_recovery.py | 15 ---- 6 files changed, 43 insertions(+), 139 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 4ae58d6adb7..ed5e4708776 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -2910,27 +2910,25 @@ async def _save_checkpoint( _REFIT_UNWIND_GRACE_S = 60.0 def _refit_await_budget_s(self) -> Optional[float]: - """How long to wait for a refit lifecycle call, or None to wait forever.""" + """How long to wait for the refit before giving up, or None to wait forever.""" deadline = self._async_cfg.generation_fleet_health.refit_timeout_s return None if deadline is None else deadline + self._REFIT_UNWIND_GRACE_S - async def _run_refit_call_within( - self, call: Callable[[], Any], *, what: str - ) -> Any: - """Run a blocking refit lifecycle call off-loop with the refit deadline. + async def _sync_weights_within(self, kv_scales, what: str) -> None: + """Run the refit off-loop, and stop waiting if it outlives the deadline. - WHY THIS IS NEEDED ON TOP OF EVERY WORKER-SIDE BOUND. A frozen-but-alive rank is - a Ray actor that never answers, and Ray puts no timeout on an actor call. Pause, - weight sync, resume, and fallback cache invalidation all cross that boundary, so - every one must use this wrapper. Job 6508251 measured the old end state: the - deadline fired, the workers aborted, the trainers returned, every actor was idle - -- and the run still sat for 1800s because the controller await had no bound. + WHY THIS IS NEEDED ON TOP OF EVERY WORKER-SIDE BOUND. A frozen-but-alive rank is a + Ray actor that never answers, and Ray puts no timeout on an actor call. So the + controller's await never resolves however well the workers behave. Job 6508251 + measured that end state: the deadline fired, the workers aborted, the trainers + returned, every actor was idle -- and the run still sat for 1800s because this + await had no bound. It is the last unbounded wait on the refit path, and unlike the + others it is not in NCCL or CUDA. - Timing out raises RefitAborted. The weight-sync caller routes that through its - existing recovery path; pause/resume deliberately remain fail-stop because a - partial multi-engine lifecycle call leaves the controller unable to prove which - preserved requests are paused or cache-cleared. In every case the dispatch gate - stays closed rather than allowing requests into uncertain engine state. + Timing out raises RefitAborted so this joins the existing recovery path rather than + inventing a second one: the caller reconciles membership and retries once. By then + the fleet probe has usually condemned the silent shard, so the rebuild can exclude + it and the retry may genuinely succeed. A DEDICATED DAEMON THREAD, not asyncio.to_thread. to_thread runs on the default ThreadPoolExecutor, whose workers are non-daemon and are joined at interpreter @@ -2938,14 +2936,18 @@ async def _run_refit_call_within( a wedge in the refit for a wedge on the way out. asyncio.wait_for cannot cancel a running thread either way; this only controls whether the orphan can block exit. - The orphan is a real consequence, not a free win. A timed-out call can leave a - thread parked on the old Ray future. The thread is therefore a daemon: bounded - failure is preferable to trading a runtime wedge for an interpreter-shutdown - wedge. + The orphan is a real consequence, not a free win. If the retry succeeds the run + continues with a thread still parked inside the old sync_weights, and if that rank + were ever resumed it would wake up holding a communicator that has since been + replaced. Acceptable against a guaranteed 30-minute stall, and it is why this path + is bounded-failure-first rather than resume-and-forget. """ budget_s = self._refit_await_budget_s() if budget_s is None: - return await asyncio.to_thread(call) + await asyncio.to_thread( + self._weight_synchronizer.sync_weights, kv_scales=kv_scales + ) + return loop = asyncio.get_running_loop() settled: asyncio.Future = loop.create_future() @@ -2958,34 +2960,24 @@ def _settle(setter, value) -> None: def _run() -> None: try: - result = call() + self._weight_synchronizer.sync_weights(kv_scales=kv_scales) except BaseException as exc: # noqa: BLE001 - re-raised on the loop below loop.call_soon_threadsafe(_settle, settled.set_exception, exc) else: - loop.call_soon_threadsafe(_settle, settled.set_result, result) + loop.call_soon_threadsafe(_settle, settled.set_result, None) - thread_label = what.replace(" ", "-") - threading.Thread( - target=_run, name=f"sc-refit-{thread_label}", daemon=True - ).start() + threading.Thread(target=_run, name=f"sc-refit-{what}", daemon=True).start() try: - return await asyncio.wait_for(settled, budget_s) + await asyncio.wait_for(settled, budget_s) except asyncio.TimeoutError: raise RefitAborted( - f"{what} did not return within {budget_s}s. A generation rank is most " - "likely alive but not answering -- a Ray actor call has no timeout of " - "its own, so waiting here is unbounded. Giving up with rollout dispatch " - "still paused." + f"the {what} refit did not return within {budget_s}s. Every worker-side " + "bound has passed, so a generation rank is most likely alive but not " + "answering -- a Ray actor call has no timeout of its own, so waiting here " + "is unbounded. Giving up so the fleet can be reconciled and retried." ) from None - async def _sync_weights_within(self, kv_scales, what: str) -> None: - """Run one weight sync through the bounded refit lifecycle wrapper.""" - await self._run_refit_call_within( - partial(self._weight_synchronizer.sync_weights, kv_scales=kv_scales), - what=f"the {what} refit", - ) - async def _sync_weights( self, *, @@ -3058,12 +3050,9 @@ async def _sync_weights( clear_cache = self._async_cfg.recompute_kv_cache_after_weight_updates print("⏸️ Requesting generation pause before refit", flush=True) - generation_paused_for_refit = await self._run_refit_call_within( - partial( - self._gen.pause_generation_for_refit, - clear_cache=clear_cache, - ), - what="generation pause before refit", + generation_paused_for_refit = await asyncio.to_thread( + self._gen.pause_generation_for_refit, + clear_cache=clear_cache, ) if generation_paused_for_refit: print( @@ -3120,9 +3109,8 @@ async def _sync_weights( self._promote_refit_shards() print("▶️ Requesting generation resume after refit", flush=True) - generation_resumed_after_refit = await self._run_refit_call_within( - self._gen.resume_generation_after_refit, - what="generation resume after refit", + generation_resumed_after_refit = await asyncio.to_thread( + self._gen.resume_generation_after_refit ) if generation_paused_for_refit and not generation_resumed_after_refit: raise RuntimeError( @@ -3134,10 +3122,7 @@ async def _sync_weights( # time. Other backends must explicitly confirm equivalent engine-side cache # handling; merely clearing reusable prefix entries does not invalidate KV # blocks still owned by running requests. - invalidated = await self._run_refit_call_within( - self._gen.invalidate_kv_cache, - what="generation KV-cache invalidation after refit", - ) + invalidated = await asyncio.to_thread(self._gen.invalidate_kv_cache) if not invalidated: raise NotImplementedError( "recompute_kv_cache_after_weight_updates=True requires the " diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index 63e31271565..2e253df2411 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -1450,8 +1450,8 @@ def invalidate_kv_cache(self) -> bool: method_name, run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"], ) - results = [result for result in ray.get(futures) if result is not None] - return bool(results) and all(results) + results = ray.get(futures) + return all(result for result in results if result is not None) except Exception as e: print(f"Error invalidating vLLM caches: {e}") return False diff --git a/nemo_rl/models/generation/vllm/vllm_worker.py b/nemo_rl/models/generation/vllm/vllm_worker.py index f1e1ded7480..ba9ad5fb837 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker.py +++ b/nemo_rl/models/generation/vllm/vllm_worker.py @@ -1412,7 +1412,7 @@ def nccl_reshard_refit(self, refit_timeout_s: Optional[float] = None) -> bool: traceback.print_exc() return False - def reset_prefix_cache(self) -> bool: + def reset_prefix_cache(self): """Reset the prefix cache of vLLM engine.""" assert self.llm is not None, ( "Attempting to reset prefix cache with either an uninitialized vLLM or non-model-owner" @@ -1423,10 +1423,9 @@ def reset_prefix_cache(self) -> bool: "reset_prefix_cache can only be used with async_engine=False. Use reset_prefix_cache_async instead." ) - reset_successful = self.llm.llm_engine.reset_prefix_cache() + self.llm.llm_engine.reset_prefix_cache() gc.collect() torch.cuda.empty_cache() - return reset_successful def sleep(self): """Put the vLLM engine to sleep.""" diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index b30c900bc6c..01076d19481 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -1587,7 +1587,7 @@ async def nccl_reshard_refit_async( traceback.print_exc() return False - async def reset_prefix_cache_async(self) -> bool: + async def reset_prefix_cache_async(self): """Async version of reset_prefix_cache.""" assert self.llm is not None, ( "Attempting to reset prefix cache with either an uninitialized vLLM or non-model-owner" @@ -1598,10 +1598,9 @@ async def reset_prefix_cache_async(self) -> bool: "reset_prefix_cache_async can only be used with async_engine=True. Use reset_prefix_cache instead." ) - reset_successful = await self.llm.reset_prefix_cache() + await self.llm.reset_prefix_cache() gc.collect() torch.cuda.empty_cache() - return reset_successful async def pause_generation_async(self, *, clear_cache: bool) -> bool: """Pause vLLM generation for an in-flight weight update.""" diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index 4faa82cdf7c..f6bc2d5829b 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -184,70 +184,6 @@ async def test_async_vllm_worker_uses_native_keep_pause_and_resume() -> None: worker.llm.resume_generation.assert_awaited_once_with() -@pytest.mark.asyncio -async def test_async_vllm_worker_propagates_prefix_reset_result( - monkeypatch: pytest.MonkeyPatch, -) -> None: - worker = VllmAsyncGenerationWorkerImpl.__new__(VllmAsyncGenerationWorkerImpl) - worker.cfg = {"vllm_cfg": {"async_engine": True}} - worker.llm = MagicMock(reset_prefix_cache=AsyncMock(return_value=False)) - empty_cache = MagicMock() - monkeypatch.setattr(torch.cuda, "empty_cache", empty_cache) - - assert await worker.reset_prefix_cache_async() is False - - worker.llm.reset_prefix_cache.assert_awaited_once_with() - empty_cache.assert_called_once_with() - - -def test_sync_vllm_worker_propagates_prefix_reset_result( - monkeypatch: pytest.MonkeyPatch, -) -> None: - worker = VllmGenerationWorkerImpl.__new__(VllmGenerationWorkerImpl) - worker.cfg = {"vllm_cfg": {"async_engine": False}} - reset_prefix_cache = MagicMock(return_value=False) - worker.llm = types.SimpleNamespace( - llm_engine=types.SimpleNamespace(reset_prefix_cache=reset_prefix_cache) - ) - empty_cache = MagicMock() - monkeypatch.setattr(torch.cuda, "empty_cache", empty_cache) - - assert worker.reset_prefix_cache() is False - - reset_prefix_cache.assert_called_once_with() - empty_cache.assert_called_once_with() - - -@pytest.mark.parametrize( - ("worker_results", "expected"), - [ - ([True, None], True), - ([False, None], False), - ([None, None], False), - ], -) -def test_vllm_generation_reports_prefix_reset_result( - monkeypatch: pytest.MonkeyPatch, - worker_results: list[bool | None], - expected: bool, -) -> None: - generation = VllmGeneration.__new__(VllmGeneration) - generation.cfg = {"vllm_cfg": {"async_engine": True}} - generation.worker_group = MagicMock() - futures = [object(), object()] - generation.worker_group.run_all_workers_single_data.return_value = futures - ray_get = MagicMock(return_value=worker_results) - monkeypatch.setattr(ray, "get", ray_get) - - assert generation.invalidate_kv_cache() is expected - - generation.worker_group.run_all_workers_single_data.assert_called_once_with( - "reset_prefix_cache_async", - run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"], - ) - ray_get.assert_called_once_with(futures) - - def test_vllm_generation_broadcasts_native_refit_pause_and_resume( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/unit/single_controller/test_refit_recovery.py b/tests/unit/single_controller/test_refit_recovery.py index 1d1e1a50cd6..5c1b05104d6 100644 --- a/tests/unit/single_controller/test_refit_recovery.py +++ b/tests/unit/single_controller/test_refit_recovery.py @@ -483,21 +483,6 @@ def test_a_refit_that_never_returns_gives_up(self): never.set() assert elapsed < 10.0, "it must give up, not wait on a rank that never answers" - def test_a_generation_lifecycle_call_uses_the_same_bound(self): - never = threading.Event() - ctrl = self._ctrl(0.1, lambda **_: None) - ctrl._REFIT_UNWIND_GRACE_S = 0.1 - - with pytest.raises(RefitAborted, match="generation pause before refit"): - asyncio.run( - ctrl._run_refit_call_within( - never.wait, - what="generation pause before refit", - ) - ) - - never.set() - def test_giving_up_is_recognised_as_an_abort(self): """So it lands in the existing `except (RefitAborted, RayActorError)` recovery.""" never = threading.Event() From c6a082b0f760da80f8013cddc9035d2916298863 Mon Sep 17 00:00:00 2001 From: Youngeun Kwon Date: Mon, 31 Aug 2026 23:32:00 -0700 Subject: [PATCH 04/13] docs(sc): remove refit timeout guidance Signed-off-by: Youngeun Kwon --- docs/guides/single-controller.md | 2 -- 1 file changed, 2 deletions(-) diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index 53048bfa1f2..704e47e3e1d 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -199,8 +199,6 @@ The SC path splits the async-GRPO loop across a rollout pump and a train pump th 4. **Train pump loop**: `sampler.evict` drops out-of-window groups and `sampler.select` picks the next batch. On PPO, `_value_stage` runs the critic forward, `_advantage_stage` computes advantages, and `_value_train_epochs` runs `ppo.critic_ppo_epochs` critic updates. The TQPolicy split API then runs one optimizer step per RL step on GRPO, or `ppo.ppo_epochs` policy updates on PPO. 5. **Weight sync**: after each optimizer step the pump bumps the trainer version, clears rollout permission, asks the generation backend to pause in-flight work, calls the weight synchronizer, resumes generation, and re-opens the rollout pump for the next version. vLLM's native pause preserves active request state; backends without pause support warn and retain their existing refit behavior. -Freezing a vLLM request does not stop its existing end-to-end generation timeout. Size `NRL_VLLM_ASYNC_TIMEOUT_SECONDS` and any generation-router `backend_timeout_s` to cover the longest expected generation plus time spent paused for refit. - ## Relation to Legacy Async GRPO The [legacy async GRPO](./async-grpo.md) (`grpo.async_grpo.enabled: true` under `run_grpo.py`) and the SC path both target the same async training problem but partition responsibilities differently: From c877a8e22970ac2dd79436cf4541637528bff7df Mon Sep 17 00:00:00 2001 From: Youngeun Kwon Date: Mon, 14 Sep 2026 09:37:01 -0700 Subject: [PATCH 05/13] refactor(single-controller): remove duplicate refit warning Signed-off-by: Youngeun Kwon --- nemo_rl/algorithms/single_controller.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index e52a74e343e..dd317ec5d87 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -4595,12 +4595,6 @@ async def _sync_weights( f" {len(self._inflight_by_group_id)} in-flight rollout group(s) paused", flush=True, ) - else: - print( - " generation backend has no native pause support; " - "in-flight requests retain their existing refit behavior", - flush=True, - ) try: await self._sync_weights_within(kv_scales, "first") From 558ceabab7d82c9f1392e8b8236876bce9f50481 Mon Sep 17 00:00:00 2001 From: Youngeun Kwon Date: Mon, 14 Sep 2026 15:45:20 -0700 Subject: [PATCH 06/13] Bound single-controller generation pause and resume waits Signed-off-by: Youngeun Kwon --- nemo_rl/algorithms/single_controller.py | 34 +++++++--- nemo_rl/models/generation/interfaces.py | 12 +++- .../models/generation/vllm/vllm_generation.py | 25 +++++-- .../models/generation/test_vllm_generation.py | 35 +++++++++- .../single_controller/test_checkpointing.py | 8 ++- .../single_controller/test_refit_recovery.py | 68 ++++++++++++++++++- .../test_single_controller_actor.py | 32 +++++++-- .../single_controller/test_train_pump_e2e.py | 8 ++- 8 files changed, 188 insertions(+), 34 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index dd317ec5d87..ed4f45e28c8 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -4552,6 +4552,8 @@ async def _sync_weights( else await self._abort_stale_inflight() ) + # TODO(#2625): Add drain-gate support during refit. + # Reconcile before the refit, not on a death event. The refit group is provably # idle here and every rank is synchronized, which is required because the # operations that change membership are themselves collectives. Doing it every @@ -4585,18 +4587,20 @@ async def _sync_weights( await self._reconcile_refit_membership() clear_cache = self._async_cfg.recompute_kv_cache_after_weight_updates + timeout_s = self._refit_await_budget_s() + generation_paused_for_refit = False print("⏸️ Requesting generation pause before refit", flush=True) - generation_paused_for_refit = await asyncio.to_thread( - self._gen.pause_generation_for_refit, - clear_cache=clear_cache, - ) - if generation_paused_for_refit: - print( - f" {len(self._inflight_by_group_id)} in-flight rollout group(s) paused", - flush=True, - ) - try: + generation_paused_for_refit = await asyncio.to_thread( + self._gen.pause_generation_for_refit, + clear_cache=clear_cache, + timeout_s=timeout_s, + ) + if generation_paused_for_refit: + print( + f" {len(self._inflight_by_group_id)} in-flight rollout group(s) paused", + flush=True, + ) await self._sync_weights_within(kv_scales, "first") except (RefitAborted, RayActorError) as failure: # DETECT AND FAIL FAST, because this one cannot be recovered from. @@ -4624,6 +4628,14 @@ async def _sync_weights( raise with self._recovery_window(): await self._recover_from_failed_refit(failure) + if not generation_paused_for_refit: + # The failed pause may have reached only some engines. Pause all + # surviving engines before retrying the refit. + generation_paused_for_refit = await asyncio.to_thread( + self._gen.pause_generation_for_refit, + clear_cache=clear_cache, + timeout_s=timeout_s, + ) # Once only: a second failure is a real fault, not a membership problem, # and retrying forever would recreate the wedge this exists to remove. await self._sync_weights_within(kv_scales, "retry") @@ -4640,7 +4652,7 @@ async def _sync_weights( print("▶️ Requesting generation resume after refit", flush=True) generation_resumed_after_refit = await asyncio.to_thread( - self._gen.resume_generation_after_refit + self._gen.resume_generation_after_refit, timeout_s=timeout_s ) if generation_paused_for_refit and not generation_resumed_after_refit: raise RuntimeError( diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index 15c2f44c7db..8d063d52276 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -581,7 +581,9 @@ def invalidate_kv_cache(self) -> bool: """ return False - def pause_generation_for_refit(self, *, clear_cache: bool) -> bool: + def pause_generation_for_refit( + self, *, clear_cache: bool, timeout_s: Optional[float] = None + ) -> bool: """Pause in-flight generation while preserving request state. Backends with native in-flight refit support override this hook. The default @@ -593,6 +595,7 @@ def pause_generation_for_refit(self, *, clear_cache: bool) -> bool: Args: clear_cache: Also clear the engine's reusable caches at pause time so preserved requests recompute their KV after the weight update. + timeout_s: Maximum time to wait for engines to pause; None waits indefinitely. Returns: True if every engine paused; False when the backend has no native pause @@ -601,13 +604,18 @@ def pause_generation_for_refit(self, *, clear_cache: bool) -> bool: _warn_unsupported_in_flight_refit_pause_once(type(self).__name__) return False - def resume_generation_after_refit(self) -> bool: + def resume_generation_after_refit( + self, *, timeout_s: Optional[float] = None + ) -> bool: """Resume generation paused by :meth:`pause_generation_for_refit`. The default implementation shares the once-per-backend warning emitted by :meth:`pause_generation_for_refit` and lets the refit continue for backends without native pause/resume support. + Args: + timeout_s: Maximum time to wait for engines to resume; None waits indefinitely. + Returns: True if every engine resumed; False when the backend has no native resume support. Backends with native support raise when resuming fails. diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index 0f69b2c6816..767610dc577 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -33,6 +33,7 @@ from nemo_rl.distributed.batched_data_dict import BatchedDataDict, SlicedDataDict from nemo_rl.distributed.named_sharding import NamedSharding from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env +from nemo_rl.distributed.refit_watchdog import RefitAborted from nemo_rl.distributed.virtual_cluster import NVLINK_DOMAIN_UNKNOWN, RayVirtualCluster from nemo_rl.distributed.worker_groups import RayWorkerBuilder, RayWorkerGroup from nemo_rl.models.generation.fleet_health import ( @@ -1532,7 +1533,9 @@ def invalidate_kv_cache(self) -> bool: print(f"Error invalidating vLLM caches: {e}") return False - def pause_generation_for_refit(self, *, clear_cache: bool) -> bool: + def pause_generation_for_refit( + self, *, clear_cache: bool, timeout_s: Optional[float] = None + ) -> bool: """Pause every async vLLM engine while preserving in-flight requests.""" if not self.cfg["vllm_cfg"]["async_engine"]: raise RuntimeError("pause_generation_for_refit requires async_engine=True") @@ -1546,11 +1549,19 @@ def pause_generation_for_refit(self, *, clear_cache: bool) -> bool: worker.pause_generation_async.remote(clear_cache=clear_cache) for worker in self._refit_leader_workers() ] - if not all(ray.get(futures)): + try: + results = ray.get(futures, timeout=timeout_s) + except ray.exceptions.GetTimeoutError as exc: + raise RefitAborted( + f"vLLM generation pause did not return within {timeout_s}s" + ) from exc + if not all(results): raise RuntimeError("Failed to pause every async vLLM engine") return True - def resume_generation_after_refit(self) -> bool: + def resume_generation_after_refit( + self, *, timeout_s: Optional[float] = None + ) -> bool: """Resume every async vLLM engine paused for refit.""" if not self.cfg["vllm_cfg"]["async_engine"]: raise RuntimeError( @@ -1563,7 +1574,13 @@ def resume_generation_after_refit(self) -> bool: worker.resume_generation_async.remote() for worker in self._refit_leader_workers() ] - if not all(ray.get(futures)): + try: + results = ray.get(futures, timeout=timeout_s) + except ray.exceptions.GetTimeoutError as exc: + raise RefitAborted( + f"vLLM generation resume did not return within {timeout_s}s" + ) from exc + if not all(results): raise RuntimeError("Failed to resume every async vLLM engine") return True diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index 62c36148cc7..4168003a6a7 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -32,6 +32,7 @@ from nemo_rl.algorithms.loss import NLLLossFn from nemo_rl.algorithms.utils import get_tokenizer from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.distributed.refit_watchdog import RefitAborted from nemo_rl.distributed.virtual_cluster import RayVirtualCluster from nemo_rl.models.generation import configure_generation_config from nemo_rl.models.generation.interfaces import ( @@ -491,8 +492,10 @@ def shutdown(self) -> None: assert worker.llm is None +@pytest.mark.parametrize("timeout_s", [None, 0.25]) def test_vllm_generation_broadcasts_native_refit_pause_and_resume( monkeypatch: pytest.MonkeyPatch, + timeout_s: float | None, ) -> None: generation = VllmGeneration.__new__(VllmGeneration) generation.cfg = {"vllm_cfg": {"async_engine": True}} @@ -509,13 +512,39 @@ def test_vllm_generation_broadcasts_native_refit_pause_and_resume( ray_get = MagicMock(side_effect=[[True, True], [True, True]]) monkeypatch.setattr(ray, "get", ray_get) - assert generation.pause_generation_for_refit(clear_cache=True) - assert generation.resume_generation_after_refit() + assert generation.pause_generation_for_refit(clear_cache=True, timeout_s=timeout_s) + assert generation.resume_generation_after_refit(timeout_s=timeout_s) for worker in leaders: worker.pause_generation_async.remote.assert_called_once_with(clear_cache=True) worker.resume_generation_async.remote.assert_called_once_with() - assert ray_get.call_args_list == [call(pause_futures), call(resume_futures)] + assert ray_get.call_args_list == [ + call(pause_futures, timeout=timeout_s), + call(resume_futures, timeout=timeout_s), + ] + + +@pytest.mark.parametrize("operation", ["pause", "resume"]) +def test_vllm_generation_refit_pause_and_resume_raise_on_timeout( + monkeypatch: pytest.MonkeyPatch, + operation: str, +) -> None: + generation = VllmGeneration.__new__(VllmGeneration) + generation.cfg = {"vllm_cfg": {"async_engine": True}} + generation.worker_group = MagicMock(workers=[object()]) + worker = MagicMock() + generation._refit_leader_workers = MagicMock(return_value=[worker]) + ray_get = MagicMock(side_effect=ray.exceptions.GetTimeoutError("timed out")) + monkeypatch.setattr(ray, "get", ray_get) + + with pytest.raises(RefitAborted, match=rf"{operation} did not return within 0.25s"): + if operation == "pause": + generation.pause_generation_for_refit(clear_cache=True, timeout_s=0.25) + else: + generation.resume_generation_after_refit(timeout_s=0.25) + + future = getattr(worker, f"{operation}_generation_async").remote.return_value + ray_get.assert_called_once_with([future], timeout=0.25) def test_vllm_generation_rejects_partial_refit_pause_and_resume( diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index a9a1d91ad62..e5287778f27 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -437,11 +437,13 @@ def snapshot_step_metrics(self) -> None: def get_step_metrics(self) -> dict[str, float]: return {} - def pause_generation_for_refit(self, *, clear_cache: bool) -> bool: - del clear_cache + def pause_generation_for_refit( + self, *, clear_cache: bool, timeout_s: float | None = None + ) -> bool: + del clear_cache, timeout_s return False - def resume_generation_after_refit(self) -> bool: + def resume_generation_after_refit(self, *, timeout_s: float | None = None) -> bool: return False def invalidate_kv_cache(self) -> bool: diff --git a/tests/unit/single_controller/test_refit_recovery.py b/tests/unit/single_controller/test_refit_recovery.py index eb848b2aea3..a97348da703 100644 --- a/tests/unit/single_controller/test_refit_recovery.py +++ b/tests/unit/single_controller/test_refit_recovery.py @@ -171,6 +171,67 @@ def _make_controller( ABORTED = RefitAborted("refit broadcast exceeded 60.0s and was aborted") +@pytest.mark.parametrize("failure", [ABORTED, ray.exceptions.ActorDiedError()]) +def test_failed_pause_recovers_and_pauses_survivors_before_refit(failure): + ctrl, monitor, sync = _make_controller(failure) + ctrl._async_cfg.generation_fleet_health.refit_timeout_s = 1.0 + events = [] + + def pause(**kwargs): + events.append("pause") + if events == ["pause"]: + raise failure + assert monitor.absent_shards() == [0] + return True + + ctrl._gen.pause_generation_for_refit.side_effect = pause + sync.sync_weights = MagicMock(side_effect=lambda **_: events.append("sync")) + + asyncio.run(ctrl._sync_weights()) + + assert events == ["pause", "pause", "sync"] + assert ctrl._gen.pause_generation_for_refit.call_args_list == [ + mock.call(clear_cache=False, timeout_s=61.0), + mock.call(clear_cache=False, timeout_s=61.0), + ] + assert sync.reconciled_with[-1] == [0] + ctrl._gen.resume_generation_after_refit.assert_called_once_with(timeout_s=61.0) + assert ctrl._rollout_permitted.is_set() + + +@pytest.mark.parametrize("failure", [ABORTED, ray.exceptions.ActorDiedError()]) +@pytest.mark.parametrize("with_monitor", [False, True]) +def test_unrecoverable_pause_stops_before_weight_transfer(failure, with_monitor): + ctrl, _, sync = _make_controller(failure, with_monitor=with_monitor) + ctrl._gen.pause_generation_for_refit.side_effect = failure + + with pytest.raises(type(failure)): + asyncio.run(ctrl._sync_weights()) + + assert ctrl._gen.pause_generation_for_refit.call_count == (2 if with_monitor else 1) + assert sync.sync_calls == 0 + ctrl._gen.resume_generation_after_refit.assert_not_called() + assert not ctrl._rollout_permitted.is_set() + + +@pytest.mark.parametrize("failure", [ABORTED, ray.exceptions.ActorDiedError()]) +def test_failed_resume_keeps_dispatch_closed_without_retrying_refit(failure): + ctrl, _, sync = _make_controller(failure) + sync.sync_weights = MagicMock() + ctrl._gen.resume_generation_after_refit.side_effect = failure + + with pytest.raises(type(failure)): + asyncio.run(ctrl._sync_weights()) + + sync.sync_weights.assert_called_once_with(kv_scales=None) + ctrl._gen.pause_generation_for_refit.assert_called_once_with( + clear_cache=False, timeout_s=None + ) + ctrl._gen.resume_generation_after_refit.assert_called_once_with(timeout_s=None) + ctrl._rollout_manager.resume_request_deadlines.assert_not_called() + assert not ctrl._rollout_permitted.is_set() + + class TestDeathInsideTheCollective: """RefitAborted: the watchdog broke a collective a dead peer had stalled.""" @@ -205,13 +266,16 @@ def test_generation_stays_paused_across_the_recovery_retry(self): ctrl, _, sync = _make_controller(ABORTED) sync_calls_when_resumed: list[int] = [] ctrl._gen.resume_generation_after_refit = MagicMock( - side_effect=lambda: sync_calls_when_resumed.append(sync.sync_calls) or True + side_effect=lambda **_: sync_calls_when_resumed.append(sync.sync_calls) + or True ) asyncio.run(ctrl._sync_weights()) assert sync.sync_calls == 2 - ctrl._gen.pause_generation_for_refit.assert_called_once_with(clear_cache=False) + ctrl._gen.pause_generation_for_refit.assert_called_once_with( + clear_cache=False, timeout_s=None + ) assert sync_calls_when_resumed == [2] def test_survivors_are_pulled_from_service_then_given_back(self): diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index ac6e852a7d9..a594874893c 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -432,7 +432,8 @@ def test_retune_is_a_noop_off_the_ppo_path(self): ], ) @pytest.mark.parametrize("use_nemo_gym", [False, True]) -def test_sync_weights_honors_recompute_kv_cache_config( +@pytest.mark.asyncio +async def test_sync_weights_honors_recompute_kv_cache_config( generation_pause_supported: bool, recompute_kv_cache: bool, expected_invalidation_calls: int, @@ -464,10 +465,11 @@ def test_sync_weights_honors_recompute_kv_cache_config( ), requires_kv_scale_sync=False, resume_generation_after_refit=MagicMock( - side_effect=lambda: (events.append("resume") or generation_pause_supported) + side_effect=lambda **_: ( + events.append("resume") or generation_pause_supported + ) ), ) - ctrl._inflight_by_group_id = {} ctrl._rollout_recovery_enabled = False ctrl._master_config = SimpleNamespace( env={"should_use_nemo_gym": use_nemo_gym}, @@ -483,13 +485,31 @@ def test_sync_weights_honors_recompute_kv_cache_config( token_capture=SimpleNamespace(enabled=False), ) - asyncio.run(ctrl._sync_weights()) + stale_task = asyncio.create_task(asyncio.Event().wait()) + ctrl._inflight_by_group_id = {"stale-group": (stale_task, 0)} + ctrl._trainer_version = 1 + ctrl._sampler = SimpleNamespace(should_abort_inflight=MagicMock(return_value=True)) + try: + aborted = await ctrl._sync_weights() + assert aborted == (0 if use_nemo_gym else 1) + assert stale_task.cancelled() is (not use_nemo_gym) + if use_nemo_gym: + ctrl._sampler.should_abort_inflight.assert_not_called() + else: + ctrl._sampler.should_abort_inflight.assert_called_once_with( + start_weight_version=0, current_train_weight=1 + ) + finally: + stale_task.cancel() + await asyncio.gather(stale_task, return_exceptions=True) ctrl._weight_synchronizer.sync_weights.assert_called_once_with(kv_scales=None) ctrl._gen.pause_generation_for_refit.assert_called_once_with( - clear_cache=recompute_kv_cache + clear_cache=recompute_kv_cache, timeout_s=ctrl._refit_await_budget_s() + ) + ctrl._gen.resume_generation_after_refit.assert_called_once_with( + timeout_s=ctrl._refit_await_budget_s() ) - ctrl._gen.resume_generation_after_refit.assert_called_once_with() assert ctrl._gen.invalidate_kv_cache.call_count == expected_invalidation_calls assert events == [ "pause", diff --git a/tests/unit/single_controller/test_train_pump_e2e.py b/tests/unit/single_controller/test_train_pump_e2e.py index 1a41622955f..ab85655e16a 100644 --- a/tests/unit/single_controller/test_train_pump_e2e.py +++ b/tests/unit/single_controller/test_train_pump_e2e.py @@ -194,11 +194,13 @@ def snapshot_step_metrics(self) -> None: def get_step_metrics(self) -> dict[str, float]: return {} - def pause_generation_for_refit(self, *, clear_cache: bool) -> bool: - del clear_cache + def pause_generation_for_refit( + self, *, clear_cache: bool, timeout_s: float | None = None + ) -> bool: + del clear_cache, timeout_s return False - def resume_generation_after_refit(self) -> bool: + def resume_generation_after_refit(self, *, timeout_s: float | None = None) -> bool: return False From bad0cb1691e9e20f0130219ddaee22b76a7a07b4 Mon Sep 17 00:00:00 2001 From: Youngeun Kwon Date: Tue, 15 Sep 2026 09:35:58 -0700 Subject: [PATCH 07/13] fix: preserve fleet health across refit pause and resume failures Signed-off-by: Youngeun Kwon --- nemo_rl/algorithms/single_controller.py | 32 ++-- .../single_controller_utils/config.py | 5 +- nemo_rl/models/generation/interfaces.py | 5 +- .../models/generation/vllm/vllm_generation.py | 52 ++++-- .../models/generation/test_refit_resume.py | 174 ++++++++++++++++++ .../models/generation/test_vllm_generation.py | 27 ++- .../single_controller/test_refit_recovery.py | 34 +++- 7 files changed, 287 insertions(+), 42 deletions(-) create mode 100644 tests/unit/models/generation/test_refit_resume.py diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index e21bd098800..57c453d4e74 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -3489,18 +3489,19 @@ def _recovery_window(self) -> Iterator[None]: # exhaustion check for the rest of the run. self._recovering_from_refit = False - async def _recover_from_failed_refit(self, failure: BaseException) -> None: + async def _recover_from_failed_refit( + self, failure: BaseException, *, refit_started: bool + ) -> None: """Drop whatever stopped participating, rebuild the communicator, allow a retry. Two failures arrive here and they are not the same event: - ``RefitAborted`` -- a rank went silent *inside* the collective and a worker's - watchdog broke it. Every engine that was receiving is left holding a mix of old - and new weights, so none of them may serve until a refit completes. + ``RefitAborted`` -- pausing timed out, or a rank went silent inside the + collective and a watchdog broke it. Only the latter leaves receiving engines + holding partial weights; those cannot serve until a refit completes. - ``RayActorError`` -- the collective finished and a shard died in the epilogue, - before its RPC returned. Nothing is partial; the survivors have complete weights. - Left uncaught this killed a run whose data transfer had *already succeeded*. + ``RayActorError`` -- a shard died before transfer, or in the epilogue after + the collective finished. Nothing is partial; survivors have complete weights. Both need the same repair, because both leave a communicator that no longer matches the fleet, and in the abort case no communicator at all. @@ -3519,10 +3520,9 @@ async def _recover_from_failed_refit(self, failure: BaseException) -> None: # 1. Establish who is actually gone, now, rather than on the probe's clock. await self._probe_generation_fleet() - # 2. Only an abort leaves partial weights behind. Marking survivors stale after - # a completed broadcast would pull a healthy fleet out of service over a - # transfer that succeeded. - if isinstance(failure, RefitAborted): + # 2. Only an interrupted transfer leaves partial weights. A pause failure + # before the transfer, like a death after it completed, does not. + if refit_started and isinstance(failure, RefitAborted): for shard_idx in self._gen_fleet.serving_shards(): self._gen_fleet.mark_weights_partial(shard_idx) @@ -4560,6 +4560,8 @@ async def _save_checkpoint_impl( # and have Ray deliver it -- seconds, not minutes. If this fires first we lose their # diagnosis and report only "the refit never came back", which is true but less useful. _REFIT_UNWIND_GRACE_S = 60.0 + # Pause/resume remain bounded even when the collective watchdog is disabled. + _GENERATION_CONTROL_TIMEOUT_S = 360.0 def _refit_await_budget_s(self) -> Optional[float]: """How long to wait for the refit before giving up, or None to wait forever.""" @@ -4711,7 +4713,10 @@ async def _sync_weights( clear_cache = self._async_cfg.recompute_kv_cache_after_weight_updates timeout_s = self._refit_await_budget_s() + if timeout_s is None: + timeout_s = self._GENERATION_CONTROL_TIMEOUT_S generation_paused_for_refit = False + refit_started = False print("⏸️ Requesting generation pause before refit", flush=True) try: generation_paused_for_refit = await asyncio.to_thread( @@ -4724,6 +4729,7 @@ async def _sync_weights( f" {len(self._inflight_by_group_id)} in-flight rollout group(s) paused", flush=True, ) + refit_started = True await self._sync_weights_within(kv_scales, "first") except (RefitAborted, RayActorError) as failure: # DETECT AND FAIL FAST, because this one cannot be recovered from. @@ -4750,7 +4756,9 @@ async def _sync_weights( ) raise with self._recovery_window(): - await self._recover_from_failed_refit(failure) + await self._recover_from_failed_refit( + failure, refit_started=refit_started + ) participants = self._refit_participants() # A restarted engine may have joined the rebuilt membership without # receiving the initial pause, even if it reuses the same shard index. diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 417ebf18007..3e7dbf7dcca 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -313,8 +313,9 @@ class FleetHealthConfig(BaseModel, extra="allow"): # with RefitAborted -- still far better than hanging forever inside NCCL, but choose # the deadline knowing there is no second chance. # - # None disarms it: no watchdog thread is started and the refit path is byte-identical - # to before. Set it well above a healthy refit, because the cost of firing early is + # None disables the collective watchdog and controller's transfer deadline. + # Generation pause/resume still have a 360s timeout. Set this well above a healthy + # refit, because the cost of firing early is # aborting a run that was merely slow, while the cost of firing late is only that a # wedge lasts longer before it is broken. # diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index 2dc7708f84a..e3e5e1f8c04 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -660,8 +660,9 @@ def resume_generation_after_refit( timeout_s: Maximum time to wait for engines to resume; None waits indefinitely. Returns: - True if every engine resumed; False when the backend has no native resume - support. Backends with native support raise when resuming fails. + True if every surviving engine resumed; False when the backend has no + native resume support. With fleet health, confirmed dead engines may be + excluded if enough survivors remain. Other resume failures raise. """ _warn_unsupported_in_flight_refit_pause_once(type(self).__name__) return False diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index 5aeafde7aee..39dbaf45640 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -1629,7 +1629,9 @@ def pause_generation_for_refit( ) -> bool: """Pause every async vLLM engine while preserving in-flight requests.""" if not self.cfg["vllm_cfg"]["async_engine"]: - raise RuntimeError("pause_generation_for_refit requires async_engine=True") + return super().pause_generation_for_refit( + clear_cache=clear_cache, timeout_s=timeout_s + ) if not self.worker_group or not self.worker_group.workers: raise RuntimeError("Worker group is not initialized") @@ -1653,26 +1655,46 @@ def pause_generation_for_refit( def resume_generation_after_refit( self, *, timeout_s: Optional[float] = None ) -> bool: - """Resume every async vLLM engine paused for refit.""" + """Resume surviving engines; fleet health excludes confirmed actor deaths.""" if not self.cfg["vllm_cfg"]["async_engine"]: - raise RuntimeError( - "resume_generation_after_refit requires async_engine=True" - ) + return super().resume_generation_after_refit(timeout_s=timeout_s) if not self.worker_group or not self.worker_group.workers: raise RuntimeError("Worker group is not initialized") - futures = [ - worker.resume_generation_async.remote() - for worker in self._refit_leader_workers() - ] - try: - results = ray.get(futures, timeout=timeout_s) - except ray.exceptions.GetTimeoutError as exc: + membership = self._refit_membership + shard_indices = ( + list(membership.shard_prefixes) + if membership is not None + else list(range(self.dp_size)) + ) + per_shard = ( + membership.workers_per_shard + if membership is not None + else len(self.worker_group.workers) // self.dp_size + ) + leaders = self._refit_leader_workers() + futures = [worker.resume_generation_async.remote() for worker in leaders] + # Settle all calls within one budget before inspecting failures. A dead leader + # must not prevent us from waiting for the surviving engines to resume. + _, pending = ray.wait(futures, num_returns=len(futures), timeout=timeout_s) + if pending: raise RefitAborted( f"vLLM generation resume did not return within {timeout_s}s" - ) from exc - if not all(results): - raise RuntimeError("Failed to resume every async vLLM engine") + ) + for shard_idx, worker, future in zip(shard_indices, leaders, futures): + try: + resumed = ray.get(future) + except ray.exceptions.ActorDiedError as exc: + if self.fleet_monitor is None: + raise + # A late error from an old actor must not condemn its replacement. + if self.worker_group.workers[shard_idx * per_shard] is worker: + self.fleet_monitor.record_actor_death(shard_idx, error=str(exc)) + continue + if not resumed: + raise RuntimeError("Failed to resume every async vLLM engine") + if self.fleet_monitor is not None: + self.fleet_monitor.raise_if_exhausted() return True @property diff --git a/tests/unit/models/generation/test_refit_resume.py b/tests/unit/models/generation/test_refit_resume.py new file mode 100644 index 00000000000..b19b6f6cf71 --- /dev/null +++ b/tests/unit/models/generation/test_refit_resume.py @@ -0,0 +1,174 @@ +# 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. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import ray + +from nemo_rl.models.generation.fleet_health import ( + FleetHealthPolicy, + GenerationFleetExhausted, + GenerationFleetHealth, + ShardState, +) +from nemo_rl.models.generation.interfaces import ( + _warn_unsupported_in_flight_refit_pause_once, +) +from nemo_rl.models.generation.vllm import VllmGeneration + + +@pytest.fixture +def generation(monkeypatch: pytest.MonkeyPatch) -> VllmGeneration: + gen = VllmGeneration.__new__(VllmGeneration) + gen.cfg = {"vllm_cfg": {"async_engine": True}} + gen.dp_size = 3 + gen.worker_group = SimpleNamespace(workers=[MagicMock() for _ in range(6)]) + gen._refit_membership = SimpleNamespace( + shard_prefixes={0: 0, 2: 2}, workers_per_shard=2 + ) + gen.fleet_monitor = GenerationFleetHealth( + shard_count=3, policy=FleetHealthPolicy(min_healthy_shards=1) + ) + gen.fleet_monitor.record_actor_death(1) + monkeypatch.setattr( + ray, "wait", MagicMock(side_effect=lambda refs, **_: (refs, [])) + ) + return gen + + +@pytest.mark.parametrize("dead_shard", [0, 2]) +def test_resume_records_death_and_waits_for_survivors( + generation: VllmGeneration, monkeypatch: pytest.MonkeyPatch, dead_shard: int +) -> None: + leaders = generation._refit_leader_workers() + futures = [worker.resume_generation_async.remote.return_value for worker in leaders] + dead_future = futures[dead_shard // 2] + + def result(ref: object) -> bool: + # Even when the dead leader comes first, every RPC has settled already. + ray.wait.assert_called_once_with(futures, num_returns=2, timeout=0.25) + if ref is dead_future: + raise ray.exceptions.ActorDiedError() + return True + + ray_get = MagicMock(side_effect=result) + monkeypatch.setattr(ray, "get", ray_get) + + assert generation.resume_generation_after_refit(timeout_s=0.25) + + assert ray_get.call_count == 2 + assert generation.fleet_monitor.state_of(dead_shard) is ShardState.DEAD + assert generation.fleet_monitor.serving_shards() == [2 - dead_shard] + + +def test_late_resume_error_does_not_condemn_replacement( + generation: VllmGeneration, monkeypatch: pytest.MonkeyPatch +) -> None: + monitor = generation.fleet_monitor + + def replace_actor( + refs: list[object], **kwargs: object + ) -> tuple[list[object], list]: + monitor.record_actor_death(0) + monitor.mark_restarting(0) + generation.worker_group.workers[0] = MagicMock() + monitor.mark_loaded(0) + return refs, [] + + monkeypatch.setattr(ray, "wait", MagicMock(side_effect=replace_actor)) + monkeypatch.setattr( + ray, "get", MagicMock(side_effect=[ray.exceptions.ActorDiedError(), True]) + ) + + assert generation.resume_generation_after_refit(timeout_s=0.25) + assert monitor.state_of(0) is ShardState.STALE + assert monitor.serving_shards() == [2] + + +def test_resume_fails_if_the_remaining_fleet_is_exhausted( + generation: VllmGeneration, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + ray, "get", MagicMock(side_effect=ray.exceptions.ActorDiedError()) + ) + with pytest.raises(GenerationFleetExhausted): + generation.resume_generation_after_refit(timeout_s=0.25) + assert generation.fleet_monitor.serving_shards() == [] + + +def test_resume_does_not_suppress_death_without_fleet_health( + generation: VllmGeneration, monkeypatch: pytest.MonkeyPatch +) -> None: + generation.fleet_monitor = None + monkeypatch.setattr( + ray, "get", MagicMock(side_effect=ray.exceptions.ActorDiedError()) + ) + with pytest.raises(ray.exceptions.ActorDiedError): + generation.resume_generation_after_refit(timeout_s=0.25) + + +def test_resume_does_not_treat_unavailability_as_death( + generation: VllmGeneration, monkeypatch: pytest.MonkeyPatch +) -> None: + failure = ray.exceptions.ActorUnavailableError("temporarily unavailable", None) + monkeypatch.setattr(ray, "get", MagicMock(side_effect=failure)) + with pytest.raises(ray.exceptions.ActorUnavailableError): + generation.resume_generation_after_refit(timeout_s=0.25) + assert generation.fleet_monitor.serving_shards() == [0, 2] + + +def test_controller_continues_after_a_confirmed_resume_death( + generation: VllmGeneration, monkeypatch: pytest.MonkeyPatch +) -> None: + from tests.unit.single_controller.test_refit_recovery import _make_controller + + ctrl, monitor, sync = _make_controller(None, shard_count=3, dead_shards=()) + monitor.record_actor_death(1) + generation.fleet_monitor = monitor + ctrl._gen.resume_generation_after_refit = generation.resume_generation_after_refit + sync.sync_weights = MagicMock() + monkeypatch.setattr( + ray, "get", MagicMock(side_effect=[ray.exceptions.ActorDiedError(), True]) + ) + + asyncio.run(ctrl._sync_weights()) + + sync.sync_weights.assert_called_once() + assert monitor.state_of(0) is ShardState.DEAD + assert monitor.state_of(2) is ShardState.HEALTHY + assert monitor.snapshot()[2].weight_version == ctrl._trainer_version + assert ctrl._rollout_permitted.is_set() + ctrl._rollout_manager.resume_request_deadlines.assert_called_once() + + +def test_sync_engine_uses_unsupported_pause_contract( + generation: VllmGeneration, capsys: pytest.CaptureFixture[str] +) -> None: + generation.cfg["vllm_cfg"]["async_engine"] = False + _warn_unsupported_in_flight_refit_pause_once.cache_clear() + try: + assert generation.pause_generation_for_refit(clear_cache=False) is False + assert generation.resume_generation_after_refit() is False + assert ( + capsys.readouterr().out.count("no native generation pause/resume support") + == 1 + ) + for worker in generation.worker_group.workers: + worker.pause_generation_async.remote.assert_not_called() + worker.resume_generation_async.remote.assert_not_called() + finally: + _warn_unsupported_in_flight_refit_pause_once.cache_clear() diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index 4168003a6a7..6ca89d36ce4 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -509,8 +509,13 @@ def test_vllm_generation_broadcasts_native_refit_pause_and_resume( worker.pause_generation_async.remote.return_value = pause_future worker.resume_generation_async.remote.return_value = resume_future generation._refit_leader_workers = MagicMock(return_value=leaders) - ray_get = MagicMock(side_effect=[[True, True], [True, True]]) + generation._refit_membership = None + generation.dp_size = 2 + generation.fleet_monitor = None + ray_get = MagicMock(side_effect=[[True, True], True, True]) monkeypatch.setattr(ray, "get", ray_get) + ray_wait = MagicMock(return_value=(resume_futures, [])) + monkeypatch.setattr(ray, "wait", ray_wait) assert generation.pause_generation_for_refit(clear_cache=True, timeout_s=timeout_s) assert generation.resume_generation_after_refit(timeout_s=timeout_s) @@ -520,8 +525,10 @@ def test_vllm_generation_broadcasts_native_refit_pause_and_resume( worker.resume_generation_async.remote.assert_called_once_with() assert ray_get.call_args_list == [ call(pause_futures, timeout=timeout_s), - call(resume_futures, timeout=timeout_s), + call(resume_futures[0]), + call(resume_futures[1]), ] + ray_wait.assert_called_once_with(resume_futures, num_returns=2, timeout=timeout_s) @pytest.mark.parametrize("operation", ["pause", "resume"]) @@ -534,8 +541,13 @@ def test_vllm_generation_refit_pause_and_resume_raise_on_timeout( generation.worker_group = MagicMock(workers=[object()]) worker = MagicMock() generation._refit_leader_workers = MagicMock(return_value=[worker]) + generation._refit_membership = None + generation.dp_size = 1 ray_get = MagicMock(side_effect=ray.exceptions.GetTimeoutError("timed out")) monkeypatch.setattr(ray, "get", ray_get) + resume_future = worker.resume_generation_async.remote.return_value + ray_wait = MagicMock(return_value=([], [resume_future])) + monkeypatch.setattr(ray, "wait", ray_wait) with pytest.raises(RefitAborted, match=rf"{operation} did not return within 0.25s"): if operation == "pause": @@ -544,7 +556,11 @@ def test_vllm_generation_refit_pause_and_resume_raise_on_timeout( generation.resume_generation_after_refit(timeout_s=0.25) future = getattr(worker, f"{operation}_generation_async").remote.return_value - ray_get.assert_called_once_with([future], timeout=0.25) + if operation == "pause": + ray_get.assert_called_once_with([future], timeout=0.25) + else: + ray_wait.assert_called_once_with([future], num_returns=1, timeout=0.25) + ray_get.assert_not_called() def test_vllm_generation_rejects_partial_refit_pause_and_resume( @@ -555,10 +571,13 @@ def test_vllm_generation_rejects_partial_refit_pause_and_resume( generation.worker_group = MagicMock(workers=[object(), object()]) leaders = [MagicMock(), MagicMock()] generation._refit_leader_workers = MagicMock(return_value=leaders) + generation._refit_membership = None + generation.dp_size = 2 + monkeypatch.setattr(ray, "wait", MagicMock(return_value=([], []))) monkeypatch.setattr( ray, "get", - MagicMock(side_effect=[[True, False], [False, True]]), + MagicMock(side_effect=[[True, False], False, True]), ) with pytest.raises(RuntimeError, match="pause every async vLLM engine"): diff --git a/tests/unit/single_controller/test_refit_recovery.py b/tests/unit/single_controller/test_refit_recovery.py index 5b18a036fcd..6f0de7f3495 100644 --- a/tests/unit/single_controller/test_refit_recovery.py +++ b/tests/unit/single_controller/test_refit_recovery.py @@ -168,9 +168,11 @@ def test_failed_pause_recovers_and_pauses_survivors_before_refit(failure): ctrl, monitor, sync = _make_controller(failure) ctrl._async_cfg.generation_fleet_health.refit_timeout_s = 1.0 events = [] + serving_at_pause = [] def pause(**kwargs): events.append("pause") + serving_at_pause.append(monitor.serving_shards()) if events == ["pause"]: raise failure assert monitor.absent_shards() == [0] @@ -182,6 +184,7 @@ def pause(**kwargs): asyncio.run(ctrl._sync_weights()) assert events == ["pause", "pause", "sync"] + assert serving_at_pause == [[0, 1], [1]] assert ctrl._gen.pause_generation_for_refit.call_args_list == [ mock.call(clear_cache=False, timeout_s=61.0), mock.call(clear_cache=False, timeout_s=61.0), @@ -194,7 +197,7 @@ def pause(**kwargs): @pytest.mark.parametrize("failure", [ABORTED, ray.exceptions.ActorDiedError()]) @pytest.mark.parametrize("with_monitor", [False, True]) def test_unrecoverable_pause_stops_before_weight_transfer(failure, with_monitor): - ctrl, _, sync = _make_controller(failure, with_monitor=with_monitor) + ctrl, monitor, sync = _make_controller(failure, with_monitor=with_monitor) ctrl._gen.pause_generation_for_refit.side_effect = failure with pytest.raises(type(failure)): @@ -202,11 +205,28 @@ def test_unrecoverable_pause_stops_before_weight_transfer(failure, with_monitor) assert ctrl._gen.pause_generation_for_refit.call_count == (2 if with_monitor else 1) assert sync.sync_calls == 0 + if monitor is not None: + assert monitor.serving_shards() == [1] + assert monitor.snapshot()[1].weight_version == 0 ctrl._gen.resume_generation_after_refit.assert_not_called() assert not ctrl._rollout_permitted.is_set() -@pytest.mark.parametrize("failure", [ABORTED, ray.exceptions.ActorDiedError()]) +def test_retry_pause_failure_preserves_partial_weights() -> None: + ctrl, monitor, sync = _make_controller(ABORTED) + ctrl._gen.pause_generation_for_refit.side_effect = [True, ABORTED] + + with pytest.raises(RefitAborted): + asyncio.run(ctrl._sync_weights()) + + assert sync.sync_calls == 1 + assert monitor.state_of(1) is ShardState.STALE + assert monitor.serving_shards() == [] + assert monitor.snapshot()[1].weight_version == 0 + assert not ctrl._rollout_permitted.is_set() + + +@pytest.mark.parametrize("failure", [ABORTED, RuntimeError("resume failed")]) def test_failed_resume_keeps_dispatch_closed_without_retrying_refit(failure): ctrl, _, sync = _make_controller(failure) sync.sync_weights = MagicMock() @@ -217,9 +237,9 @@ def test_failed_resume_keeps_dispatch_closed_without_retrying_refit(failure): sync.sync_weights.assert_called_once_with(kv_scales=None) ctrl._gen.pause_generation_for_refit.assert_called_once_with( - clear_cache=False, timeout_s=None + clear_cache=False, timeout_s=360.0 ) - ctrl._gen.resume_generation_after_refit.assert_called_once_with(timeout_s=None) + ctrl._gen.resume_generation_after_refit.assert_called_once_with(timeout_s=360.0) ctrl._rollout_manager.resume_request_deadlines.assert_not_called() assert not ctrl._rollout_permitted.is_set() @@ -266,8 +286,8 @@ def test_generation_stays_paused_across_the_recovery_retry(self): assert sync.sync_calls == 2 assert ctrl._gen.pause_generation_for_refit.call_args_list == [ - mock.call(clear_cache=False, timeout_s=None), - mock.call(clear_cache=False, timeout_s=None), + mock.call(clear_cache=False, timeout_s=360.0), + mock.call(clear_cache=False, timeout_s=360.0), ] assert sync_calls_when_resumed == [2] @@ -301,7 +321,7 @@ def sync_weights(*, kv_scales: dict[str, object] | None = None) -> None: assert sync.absent_at_retry == [0] assert monitor.state_of(2) is ShardState.HEALTHY assert monitor.snapshot()[2].weight_version == ctrl._trainer_version - ctrl._gen.resume_generation_after_refit.assert_called_once_with(timeout_s=None) + ctrl._gen.resume_generation_after_refit.assert_called_once_with(timeout_s=360.0) def test_survivors_are_pulled_from_service_then_given_back(self): """Partial weights must not serve -- and must not be stranded either. From 70f3b3967c8026863b4ffe4bd108a5803f7a1d45 Mon Sep 17 00:00:00 2001 From: Youngeun Kwon Date: Mon, 12 Oct 2026 14:55:50 -0700 Subject: [PATCH 08/13] test: cover pause recovery failures and restore timeout guidance Signed-off-by: Youngeun Kwon --- docs/guides/single-controller.md | 2 ++ .../single_controller/test_refit_recovery.py | 30 +++++++++++++++++-- 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index 59287629869..9cb37661316 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -424,6 +424,8 @@ The SC path splits the async-GRPO loop across a rollout pump and a train pump th 4. **Train pump loop**: `sampler.evict` removes out-of-window groups, and `sampler.select` chooses the next chunk. For PPO, `_value_stage` computes critic values and `_advantage_stage` computes advantages. TQPolicy then updates the policy: full-batch PPO uses `ppo.ppo_epochs` updates, while streaming PPO accumulates gradients across chunks for one update. Streaming PPO refits the updated policy before `_value_train_epochs` trains the critic on the assembled full batch for `ppo.critic_ppo_epochs`. GRPO uses one TQPolicy optimizer step per RL step. 5. **Weight sync**: after each optimizer step the pump bumps the trainer version, clears rollout permission, asks the generation backend to pause in-flight work, calls the weight synchronizer, resumes generation, and re-opens the rollout pump for the next version. vLLM's native pause preserves active request state; backends without pause support warn and retain their existing refit behavior. In streaming PPO, it does the weight sync as soon as the policy training is finished, which is before the critic model training starts. This allows new batch generation with the updated policy weight overlapping with the critic training epochs. + For non-colocated vLLM, freezing a request does not stop its end-to-end generation timeout. Size `NRL_VLLM_ASYNC_TIMEOUT_SECONDS` and, when the generation router is enabled, `async_rl.generation_router.backend_timeout_s` to cover the longest expected generation plus time spent paused for refit, including membership recovery and a possible retry. With a finite `async_rl.generation_fleet_health.refit_timeout_s`, pause, each transfer attempt, and resume each use that timeout plus a 60 s grace; these are separate call budgets, not a single total refit budget. When `refit_timeout_s` is `null`, pause and resume retain a 360 s budget each, but the collective transfer wait is unbounded, so no finite worst-case total pause duration is guaranteed. + ## Relation to Legacy Async GRPO The [legacy async GRPO](./async-grpo.md) (`grpo.async_grpo.enabled: true` under `run_grpo.py`) and the SC path both target the same async training problem but partition responsibilities differently: diff --git a/tests/unit/single_controller/test_refit_recovery.py b/tests/unit/single_controller/test_refit_recovery.py index ded83c74993..4e5bd954e50 100644 --- a/tests/unit/single_controller/test_refit_recovery.py +++ b/tests/unit/single_controller/test_refit_recovery.py @@ -217,17 +217,41 @@ def test_unrecoverable_pause_stops_before_weight_transfer(failure, with_monitor) assert not ctrl._rollout_permitted.is_set() -def test_retry_pause_failure_preserves_partial_weights() -> None: +def test_pause_timeout_with_no_absent_shard_stops_without_withholding_survivors() -> ( + None +): + # A frozen-but-alive leader times out before any transfer. Without an absent + # or suspect shard, recovery must stop without marking healthy weights partial. + ctrl, monitor, sync = _make_controller(ABORTED, dead_shards=()) + ctrl._gen.pause_generation_for_refit.side_effect = ABORTED + + with pytest.raises(RuntimeError, match="could be identified as absent"): + asyncio.run(ctrl._sync_weights()) + + ctrl._gen.pause_generation_for_refit.assert_called_once() + assert sync.sync_calls == 0 + assert monitor.serving_shards() == [0, 1] + assert [shard.state for shard in monitor.snapshot()] == [ShardState.HEALTHY] * 2 + ctrl._gen.resume_generation_after_refit.assert_not_called() + assert not ctrl._rollout_permitted.is_set() + + +@pytest.mark.parametrize("retry_failure", [ABORTED, ray.exceptions.ActorDiedError()]) +def test_retry_pause_failure_preserves_partial_weights( + retry_failure: BaseException, +) -> None: ctrl, monitor, sync = _make_controller(ABORTED) - ctrl._gen.pause_generation_for_refit.side_effect = [True, ABORTED] + ctrl._gen.pause_generation_for_refit.side_effect = [True, retry_failure] - with pytest.raises(RefitAborted): + with pytest.raises(type(retry_failure)): asyncio.run(ctrl._sync_weights()) + assert ctrl._gen.pause_generation_for_refit.call_count == 2 # no second recovery assert sync.sync_calls == 1 assert monitor.state_of(1) is ShardState.STALE assert monitor.serving_shards() == [] assert monitor.snapshot()[1].weight_version == 0 + ctrl._gen.resume_generation_after_refit.assert_not_called() assert not ctrl._rollout_permitted.is_set() From 7de9e6d85fa6d42688b2698db1f322789c00a996 Mon Sep 17 00:00:00 2001 From: Youngeun Kwon Date: Fri, 16 Oct 2026 10:52:27 -0700 Subject: [PATCH 09/13] fix: default async refit KV recomputation on Signed-off-by: Youngeun Kwon --- docs/guides/async-grpo.md | 6 +- docs/guides/ppo.md | 2 +- docs/guides/single-controller.md | 2 +- examples/configs/grpo_math_1B.yaml | 4 +- ...po_math_1B_megatron_single_controller.yaml | 2 +- examples/configs/ppo_math_1B.yaml | 2 +- ...po_math_1B_megatron_single_controller.yaml | 2 +- ...eneration-colocated-reshard-async-gym.yaml | 1 + ...ron_generation-noncolocated-async-gym.yaml | 1 + ...g-megatron-single-controller-async.v1.yaml | 6 +- ...megatron_generation_single_controller.yaml | 1 + nemo_rl/algorithms/grpo.py | 2 +- nemo_rl/algorithms/ppo.py | 2 +- .../single_controller_utils/config.py | 2 +- ...i_clevr_megatron_single_controller_1n2g.sh | 1 + tests/unit/algorithms/test_grpo.py | 2 +- .../unit/reference_configs/grpo_math_1B.yaml | 4 +- .../ppo_math_1B_megatron.yaml | 2 +- tests/unit/single_controller/test_setup.py | 6 +- tests/unit/test_config_v2.py | 58 +++++++++++++++++++ 20 files changed, 87 insertions(+), 21 deletions(-) diff --git a/docs/guides/async-grpo.md b/docs/guides/async-grpo.md index 68e229238f9..57f9c2ed53b 100644 --- a/docs/guides/async-grpo.md +++ b/docs/guides/async-grpo.md @@ -43,7 +43,7 @@ grpo: max_trajectory_age_steps: 1 # Maximum age, in training steps, for trajectories max_generation_failures: 0 # Consecutive worker failures to tolerate in_flight_weight_updates: false # Enable for faster weight synchronization - recompute_kv_cache_after_weight_updates: false # Invalidates kv cache after weight-updates + recompute_kv_cache_after_weight_updates: true # Recompute KV after weight updates (default) ``` ### Complete Example Config @@ -72,7 +72,7 @@ grpo: max_trajectory_age_steps: 1 max_generation_failures: 0 # Consecutive worker failures to tolerate in_flight_weight_updates: false # Enable for faster weight synchronization - recompute_kv_cache_after_weight_updates: false # Invalidates kv cache after weight-updates + recompute_kv_cache_after_weight_updates: true # Recompute KV after weight updates (default) cluster: num_nodes: 2 @@ -223,7 +223,7 @@ If no `replay_buffer.pt` file is found in the latest checkpoint directory, train 4. **In-Flight Weight Updates**: Enable `in_flight_weight_updates: true` to refit without waiting for the longest in-flight generation to finish. Except for managed Dynamo, the collector requests a generation pause and resume from every async backend around the weight transfer. Async vLLM implements this contract while preserving request state. A backend that does not implement the hook emits a warning once per backend type per process and refits without a collector-side pause or drain; SGLang is in this group today and instead relies on the pause its own weight synchronizer performs around the transfer. Managed Dynamo always drains active trajectories before refit. vLLM requires `async_engine: true`; the Megatron backend is always async-engine. -5. **Recompute KV Cache After Weight Updates**: Set `recompute_kv_cache_after_weight_updates: true` to invalidate reusable KV/prefix caches when weights change. On the native async vLLM in-flight path, caches are cleared while generation is paused, so preserved requests recompute their KV after resuming. Other refit paths keep their existing post-update invalidation behavior. When false, in-flight requests retain their pre-update KV cache. On the Megatron generation backend, this must agree with `policy.generation.mcore_generation_config.kv_cache_management_mode`; setup errors on a mismatch. +5. **Recompute KV Cache After Weight Updates**: `recompute_kv_cache_after_weight_updates` defaults to `true` to invalidate reusable KV/prefix caches when weights change. On the native async vLLM in-flight path, caches are cleared while generation is paused, so preserved requests recompute their KV after resuming. This repeats prefill for those requests and can reduce throughput. Other refit paths keep their existing post-update invalidation behavior. Set it to `false` to retain pre-update KV caches and avoid this recomputation cost. On the Megatron generation backend, the flag must agree with `policy.generation.mcore_generation_config.kv_cache_management_mode`: use `true` with `recompute`, or explicitly set `false` with `persist`/`offload`; setup errors on a mismatch. ## Why Importance Sampling Correction Is Required for Async diff --git a/docs/guides/ppo.md b/docs/guides/ppo.md index d2082baaf1f..1e7afa5e910 100644 --- a/docs/guides/ppo.md +++ b/docs/guides/ppo.md @@ -276,7 +276,7 @@ ppo: max_trajectory_age_steps: 1 warmup_generation_lead_steps: null in_flight_weight_updates: false - recompute_kv_cache_after_weight_updates: false + recompute_kv_cache_after_weight_updates: true # Default; false avoids repeated prefill but retains old KV. drop_incomplete_targets_on_restore: false adv_estimator: diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index 9cb37661316..a5226b1ddd4 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -451,7 +451,7 @@ Do not carry `max_num_epochs: -1` across either. [ppo.md](./ppo.md#asynchronous- | `enabled: true` | Implicit — SC is always async; use `sampler.max_lookahead_versions: 0` for sync semantics, `>= 1` for async | | `max_trajectory_age_steps: N` | `sampler.name: in_order` with `sampler.max_lookahead_versions: N` | | `warmup_generation_lead_steps` (PPO) | `sampler.warmup_lookahead_versions` — the lookahead to use while critic warmup is in progress | -| `recompute_kv_cache_after_weight_updates` | `recompute_kv_cache_after_weight_updates` (same); when vLLM pauses natively, the preserved requests recompute their KV cache after refit | +| `recompute_kv_cache_after_weight_updates` | `recompute_kv_cache_after_weight_updates` (same, defaults to `true` in both loops); when vLLM pauses natively, preserved requests recompute their KV cache after refit, repeating prefill. Set `false` to retain pre-update KV. For Megatron generation, pair `true` with `kv_cache_management_mode: recompute`, or `false` with `persist`/`offload`. | | `in_flight_weight_updates` | Always effectively true; `false`-equivalent behavior is not yet supported (drain-gate tracked in [issue #2625](https://github.com/NVIDIA-NeMo/RL/issues/2625)) | | *(no legacy equivalent — matches legacy full-batch train semantics)* | `min_groups_for_streaming_train: ${grpo.num_prompts_per_step}`, or `${ppo.num_prompts_per_step}` on a PPO run | | *(no legacy equivalent — matches legacy `max_trajectory_age + 1` batches in flight)* | `max_inflight_prompts: num_prompts_per_step × (max_lookahead_versions + 1)` | diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 28603f05d51..48cb18e6576 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -64,7 +64,7 @@ grpo: # when transient generation errors are expected and acceptable to drop. max_generation_failures: 0 in_flight_weight_updates: false # Set to true to enable in-flight weight updates - recompute_kv_cache_after_weight_updates: false # Set to true to recompute kv cache after weight updates. + recompute_kv_cache_after_weight_updates: true # Set false to retain pre-update KV caches (avoids repeated prefill). # Reward-zeroing penalties applied to NeMo-Gym rollout results. reward_penalties: @@ -468,7 +468,7 @@ policy: prefix_cache_ttl_seconds: 300.0 # MCore default coordinator cache-entry lifetime. prefix_caching_routing_alpha: 1.0 # MCore default load penalty for prefix-affinity routing. max_tokens: 16384 # Maximum number of tokens to use in a single step. Analogous to vllm's max_num_batched_tokens - kv_cache_management_mode: "persist" # KV cache lifecycle across suspend/resume. Options: "persist", "offload", "recompute". + kv_cache_management_mode: "recompute" # Must agree with recompute_kv_cache_after_weight_updates. Options: "persist", "offload", "recompute". async_sched_mode: "async" # Dynamic-batching scheduler mode. Options: "legacy", "async". "async" hides scheduler CPU time behind GPU compute. materialize_only_last_token_logits: true num_speculative_tokens: 0 diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index 0b8db9d2ba9..f1d8a8cc030 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -27,7 +27,7 @@ async_rl: # How far generation may run ahead of the trainer. max_lookahead_versions: 1 # Recompute generation KV caches after each weight update. - recompute_kv_cache_after_weight_updates: false + recompute_kv_cache_after_weight_updates: true # Min ready groups the streaming trainer waits for before dispatching a batch. min_groups_for_streaming_train: ${grpo.num_prompts_per_step} # Cap on in-flight generate_and_push calls in the rollout pump. diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index e344bc26b60..e397c05d66e 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -35,7 +35,7 @@ ppo: # Requires policy_training_start_step > 0. null uses max_trajectory_age_steps. warmup_generation_lead_steps: null in_flight_weight_updates: false - recompute_kv_cache_after_weight_updates: false + recompute_kv_cache_after_weight_updates: true # true discards partial restored rows and fills from subsequent prompts. drop_incomplete_targets_on_restore: false diff --git a/examples/configs/ppo_math_1B_megatron_single_controller.yaml b/examples/configs/ppo_math_1B_megatron_single_controller.yaml index 855b211a62a..50a80511129 100644 --- a/examples/configs/ppo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/ppo_math_1B_megatron_single_controller.yaml @@ -39,7 +39,7 @@ async_rl: # throughout. Buffer capacity below must cover the wider of the two. warmup_lookahead_versions: null # Recompute generation KV caches after each weight update. - recompute_kv_cache_after_weight_updates: false + recompute_kv_cache_after_weight_updates: true # Full-batch default. Reduce this to stream forwards, GAE, and policy backward; # requires ppo_epochs=1. normalize_advantages=true then normalizes per chunk. # Value training stays full batch. diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.yaml index 2638621cf9d..283d78017bd 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.yaml @@ -23,6 +23,7 @@ policy: backend: megatron refit_transport: mcore mcore_generation_config: + kv_cache_management_mode: recompute refit_backend: nccl tensor_model_parallel_size: 4 expert_model_parallel_size: 4 diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml index ba5b113e684..37e3074e215 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml @@ -19,6 +19,7 @@ policy: backend: megatron refit_transport: mcore mcore_generation_config: + kv_cache_management_mode: recompute refit_backend: nccl expert_model_parallel_size: 4 colocated: diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml index 02392c7f296..ab82f1ef235 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml @@ -4,6 +4,10 @@ grpo: val_period: 0 val_at_start: false val_at_end: false +policy: + generation: + mcore_generation_config: + kv_cache_management_mode: recompute data_plane: enabled: true simple: @@ -12,7 +16,7 @@ async_rl: sampler: name: in_order max_lookahead_versions: 1 - recompute_kv_cache_after_weight_updates: false + recompute_kv_cache_after_weight_updates: true min_groups_for_streaming_train: ${grpo.num_prompts_per_step} max_inflight_prompts: ${mul:${grpo.num_prompts_per_step}, 2} max_buffered_rollouts: ${mul:${grpo.num_prompts_per_step}, 2} diff --git a/examples/nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml b/examples/nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml index 369d829deb2..80d2d84ee34 100644 --- a/examples/nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml +++ b/examples/nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml @@ -43,6 +43,7 @@ policy: backend: "megatron" refit_transport: "mcore" mcore_generation_config: + kv_cache_management_mode: recompute refit_backend: nccl # NeMo-Gym drives rollouts through the engine's OpenAI server. expose_http_server: true diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index f45448f4107..aafb03288ec 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -286,7 +286,7 @@ class AsyncGRPOConfig(BaseModel, extra="allow"): # without waiting for the pending generations to finish. in_flight_weight_updates: bool = False # Recomputes the KV cache after weight updates. - recompute_kv_cache_after_weight_updates: bool = False + recompute_kv_cache_after_weight_updates: bool = True class RewardPenaltyTokenIdsConfig(BaseModel, extra="allow"): diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index 7ffb9e6255f..2dee89299ee 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -144,7 +144,7 @@ class AsyncPPOConfig(BaseModel, extra="allow"): # Allows weight updates while rollout requests are still in flight. in_flight_weight_updates: bool = False # Recomputes the KV cache after weight updates. - recompute_kv_cache_after_weight_updates: bool = False + recompute_kv_cache_after_weight_updates: bool = True # Drops partial restored targets; replacement rollouts use subsequent prompts. drop_incomplete_targets_on_restore: bool = False diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 050dddc1a57..54840893556 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -494,7 +494,7 @@ class AsyncRLConfig(BaseModel, extra="allow"): default_factory=GenerationRouterConfig ) # Recompute generation KV caches after each weight update. - recompute_kv_cache_after_weight_updates: bool = False + recompute_kv_cache_after_weight_updates: bool = True # Min ready groups the streaming trainer waits for before dispatching a batch. min_groups_for_streaming_train: int = 32 # Cap on in-flight generate_and_push calls in the rollout pump. diff --git a/tests/functional/nemotron_omni_clevr_megatron_single_controller_1n2g.sh b/tests/functional/nemotron_omni_clevr_megatron_single_controller_1n2g.sh index 9de2f6ba706..8aa7d811edd 100755 --- a/tests/functional/nemotron_omni_clevr_megatron_single_controller_1n2g.sh +++ b/tests/functional/nemotron_omni_clevr_megatron_single_controller_1n2g.sh @@ -158,6 +158,7 @@ uv run --no-sync python examples/run_grpo_single_controller.py \ ++async_rl.sampler.name=in_order \ ++async_rl.sampler.max_lookahead_versions=1 \ ++async_rl.recompute_kv_cache_after_weight_updates=false \ + policy.generation.mcore_generation_config.kv_cache_management_mode=persist \ ++async_rl.min_groups_for_streaming_train=1 \ ++async_rl.max_inflight_prompts=2 \ ++async_rl.max_buffered_rollouts=2 \ diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 37ea41a6b3e..71b68a32561 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -3889,7 +3889,7 @@ def build_nemo_gym_actors(_env_configs, **kwargs): }, "mcore_generation_config": { "expose_http_server": True, - "kv_cache_management_mode": "persist", + "kv_cache_management_mode": "recompute", }, } master_config.env = {"should_use_nemo_gym": True} diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index 9db501cc408..1e2d6e1b668 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -73,7 +73,7 @@ grpo: # when transient generation errors are expected and acceptable to drop. max_generation_failures: 0 in_flight_weight_updates: false # Set to true to enable in-flight weight updates - recompute_kv_cache_after_weight_updates: false # Set to true to recompute kv cache after weight updates. + recompute_kv_cache_after_weight_updates: true # Set false to retain pre-update KV caches (avoids repeated prefill). # Reward-zeroing penalties applied to NeMo-Gym rollout results. reward_penalties: @@ -442,7 +442,7 @@ policy: prefix_cache_ttl_seconds: 300.0 # MCore default coordinator cache-entry lifetime. prefix_caching_routing_alpha: 1.0 # MCore default load penalty for prefix-affinity routing. max_tokens: 16384 # Maximum number of tokens to use in a single step. Analogous to vllm's max_num_batched_tokens - kv_cache_management_mode: "persist" # KV cache lifecycle across suspend/resume. Options: "persist", "offload", "recompute". + kv_cache_management_mode: "recompute" # Must agree with recompute_kv_cache_after_weight_updates. Options: "persist", "offload", "recompute". async_sched_mode: "async" # Dynamic-batching scheduler mode. Options: "legacy", "async". "async" hides scheduler CPU time behind GPU compute. materialize_only_last_token_logits: true # Materialize logits only for the last token of each sequence during decode. num_speculative_tokens: 0 diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml index 472bbf29799..b95f4359481 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml @@ -30,7 +30,7 @@ ppo: max_trajectory_age_steps: 1 warmup_generation_lead_steps: null in_flight_weight_updates: false - recompute_kv_cache_after_weight_updates: false + recompute_kv_cache_after_weight_updates: true drop_incomplete_targets_on_restore: false reward_shaping: diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index 2e78bfa3865..e1ec4b1710e 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -161,7 +161,7 @@ def _make_master_config( # The megatron build path reads these before any generation factory runs. generation_config["mcore_generation_config"] = { "expose_http_server": False, - "kv_cache_management_mode": "persist", + "kv_cache_management_mode": "recompute", } policy_config["model_name"] = "test-model" return MasterConfig.model_construct( @@ -1490,11 +1490,11 @@ def test_invalid_config_fails_before_setup_factories( colocated=False, backend="megatron", megatron_enabled=False ) elif invalid_case == "megatron_recompute_mismatch": - # Flag says recompute; the engine mode (fixture default "persist") disagrees. + # The opt-out disagrees with the fixture's default recompute mode. mc = _make_master_config( colocated=False, backend="megatron", megatron_enabled=True ) - mc.async_rl.recompute_kv_cache_after_weight_updates = True + mc.async_rl.recompute_kv_cache_after_weight_updates = False elif invalid_case == "megatron_fleet_health": mc = _make_master_config( colocated=False, backend="megatron", megatron_enabled=True diff --git a/tests/unit/test_config_v2.py b/tests/unit/test_config_v2.py index f03a7a034f1..d67f4cf614f 100644 --- a/tests/unit/test_config_v2.py +++ b/tests/unit/test_config_v2.py @@ -37,6 +37,7 @@ from nemo_rl.algorithms.ppo import MasterConfig as PPOMasterConfig from nemo_rl.algorithms.rm import MasterConfig as RMMasterConfig from nemo_rl.algorithms.sft import MasterConfig as SFTMasterConfig +from nemo_rl.algorithms.single_controller_utils.config import AsyncRLConfig from nemo_rl.evals.eval import MasterConfig as EvalMasterConfig from nemo_rl.utils.config import ( load_config, @@ -58,6 +59,63 @@ ) +@pytest.mark.parametrize( + "config_class", [AsyncRLConfig, AsyncGRPOConfig, AsyncPPOConfig] +) +def test_async_kv_recompute_defaults_on_with_explicit_opt_out(config_class): + assert config_class().recompute_kv_cache_after_weight_updates is True + assert ( + config_class( + recompute_kv_cache_after_weight_updates=False + ).recompute_kv_cache_after_weight_updates + is False + ) + + +@pytest.mark.parametrize( + ("filename", "expected"), + [ + ("grpo_math_1B.yaml", True), + ("ppo_math_1B.yaml", True), + ("grpo_math_1B_megatron_single_controller.yaml", True), + ("ppo_math_1B_megatron_single_controller.yaml", True), + ("grpo_math_1B_megatron_generation_colocated_single_controller.yaml", True), + ( + "recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml", + True, + ), + ( + "recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.yaml", + True, + ), + ( + "../nemo_gym/grpo_qwen3_0_6b_megatron_generation_single_controller.yaml", + True, + ), + ( + "recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml", + True, + ), + ], +) +def test_async_kv_recompute_exemplars_and_megatron_mode(filename, expected): + config = load_config(real_configs_dir / filename) + if "async_rl" in config: + async_config = config.async_rl + config_class = AsyncRLConfig + elif "grpo" in config: + async_config = config.grpo.async_grpo + config_class = AsyncGRPOConfig + else: + async_config = config.ppo.async_ppo + config_class = AsyncPPOConfig + resolved = OmegaConf.to_container(async_config, resolve=True) + assert config_class(**resolved).recompute_kv_cache_after_weight_updates is expected + if config.policy.generation.backend == "megatron": + mode = config.policy.generation.mcore_generation_config.kv_cache_management_mode + assert expected == (mode == "recompute") + + def test_ppo_recipe_critic_epochs_follow_actor_override(): config = load_config(real_configs_dir / "ppo_math_1B.yaml") config = parse_hydra_overrides(config, ["ppo.ppo_epochs=7"]) From a5052e4bdfe4e017db27a828eb74844348aad12f Mon Sep 17 00:00:00 2001 From: Youngeun Kwon Date: Fri, 16 Oct 2026 15:42:24 -0700 Subject: [PATCH 10/13] test: align colocated Megatron smoke test with KV recomputation Signed-off-by: Youngeun Kwon --- ...colocated_reshard_gym_single_controller.sh | 1 + tests/unit/test_recipe_settings.py | 40 ++++++++++++++++++- 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/tests/functional/grpo_megatron_generation_colocated_reshard_gym_single_controller.sh b/tests/functional/grpo_megatron_generation_colocated_reshard_gym_single_controller.sh index cd085e99df9..39b0e6b0d19 100755 --- a/tests/functional/grpo_megatron_generation_colocated_reshard_gym_single_controller.sh +++ b/tests/functional/grpo_megatron_generation_colocated_reshard_gym_single_controller.sh @@ -74,6 +74,7 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE ++policy.generation.mcore_generation_config.transformer_impl=inference_optimized \ ++policy.generation.mcore_generation_config.tensor_model_parallel_size=1 \ policy.generation.mcore_generation_config.refit_backend=nccl \ + policy.generation.mcore_generation_config.kv_cache_management_mode=recompute \ policy.generation.max_new_tokens=128 \ policy.max_total_sequence_length=512 \ policy.generation.colocated.enabled=true \ diff --git a/tests/unit/test_recipe_settings.py b/tests/unit/test_recipe_settings.py index e41ab306d28..ff1a8c49eef 100644 --- a/tests/unit/test_recipe_settings.py +++ b/tests/unit/test_recipe_settings.py @@ -12,15 +12,53 @@ # See the License for the specific language governing permissions and # limitations under the License. +import shlex from pathlib import Path import pytest -from nemo_rl.utils.config import load_config, register_omegaconf_resolvers +from nemo_rl.utils.config import ( + load_config, + parse_hydra_overrides, + register_omegaconf_resolvers, +) RECIPE_ROOT = Path(__file__).parents[2] / "examples/configs/recipes" +@pytest.mark.parametrize( + "script_name", + [ + "grpo_megatron_generation_gym_single_controller.sh", + "grpo_megatron_generation_colocated_reshard_gym_single_controller.sh", + ], +) +def test_megatron_single_controller_gym_scripts_recompute_kv(script_name: str) -> None: + register_omegaconf_resolvers() + project_root = Path(__file__).parents[2] + script = (project_root / "tests/functional" / script_name).read_text() + command = next( + line + for line in script.replace("\\\n", " ").splitlines() + if line.startswith("uv run coverage run ") + ) + arguments = shlex.split(command) + config_path = arguments[arguments.index("--config") + 1].replace( + "$PROJECT_ROOT", str(project_root) + ) + overrides = [ + argument + for argument in arguments + if argument.lstrip("+").startswith(("policy.generation.", "async_rl.")) + ] + config = parse_hydra_overrides(load_config(config_path), overrides) + assert config.policy.generation.backend == "megatron" + assert ( + config.policy.generation.mcore_generation_config.kv_cache_management_mode + == "recompute" + ) + + @pytest.mark.parametrize( "recipe_name", [ From d001be04e610f3a1b2d5c5900606cb8d7fde1114 Mon Sep 17 00:00:00 2001 From: Youngeun Kwon Date: Fri, 16 Oct 2026 20:41:02 -0700 Subject: [PATCH 11/13] fix: address async KV refit review regressions Signed-off-by: Youngeun Kwon --- docs/guides/async-grpo.md | 4 +- ...8g-megatron-tp4ep4-async-gym-video.v1.yaml | 1 + nemo_rl/algorithms/grpo.py | 16 +-- nemo_rl/models/generation/interfaces.py | 23 ++-- nemo_rl/models/generation/megatron/config.py | 4 +- .../generation/sglang/sglang_generation.py | 25 +++++ .../models/generation/vllm/vllm_generation.py | 38 +++---- tests/unit/algorithms/test_grpo.py | 60 +++++++++++ .../models/generation/test_vllm_generation.py | 9 +- .../generation/test_vllm_refit_resume.py | 46 ++++++++ .../single_controller/test_checkpointing.py | 6 +- .../test_single_controller_actor.py | 100 ++++++++++++++++++ .../single_controller/test_train_pump_e2e.py | 6 +- tests/unit/test_recipe_settings.py | 90 ++++++++++++---- 14 files changed, 358 insertions(+), 70 deletions(-) diff --git a/docs/guides/async-grpo.md b/docs/guides/async-grpo.md index 57f9c2ed53b..9a76ceae772 100644 --- a/docs/guides/async-grpo.md +++ b/docs/guides/async-grpo.md @@ -221,9 +221,9 @@ If no `replay_buffer.pt` file is found in the latest checkpoint directory, train 3. **Resource Allocation**: Ensure sufficient GPU memory for both the training and generation clusters -4. **In-Flight Weight Updates**: Enable `in_flight_weight_updates: true` to refit without waiting for the longest in-flight generation to finish. Except for managed Dynamo, the collector requests a generation pause and resume from every async backend around the weight transfer. Async vLLM implements this contract while preserving request state. A backend that does not implement the hook emits a warning once per backend type per process and refits without a collector-side pause or drain; SGLang is in this group today and instead relies on the pause its own weight synchronizer performs around the transfer. Managed Dynamo always drains active trajectories before refit. vLLM requires `async_engine: true`; the Megatron backend is always async-engine. +4. **In-Flight Weight Updates**: Enable `in_flight_weight_updates: true` to refit without waiting for the longest in-flight generation to finish. Except for managed Dynamo, the collector requests a generation pause and resume from every async backend around the weight transfer. Async vLLM implements this contract while preserving request state. SGLang acknowledges these hooks through its weight synchronizer, which pauses, flushes KV, transfers weights, and resumes generation; its configured pause mode determines whether requests are requeued or aborted. A backend that does not implement the hook emits a warning once per backend type per process and refits without a collector-side pause or drain. Managed Dynamo always drains active trajectories before refit. vLLM requires `async_engine: true`; the Megatron backend is always async-engine. -5. **Recompute KV Cache After Weight Updates**: `recompute_kv_cache_after_weight_updates` defaults to `true` to invalidate reusable KV/prefix caches when weights change. On the native async vLLM in-flight path, caches are cleared while generation is paused, so preserved requests recompute their KV after resuming. This repeats prefill for those requests and can reduce throughput. Other refit paths keep their existing post-update invalidation behavior. Set it to `false` to retain pre-update KV caches and avoid this recomputation cost. On the Megatron generation backend, the flag must agree with `policy.generation.mcore_generation_config.kv_cache_management_mode`: use `true` with `recompute`, or explicitly set `false` with `persist`/`offload`; setup errors on a mismatch. +5. **Recompute KV Cache After Weight Updates**: `recompute_kv_cache_after_weight_updates` defaults to `true` to invalidate reusable KV/prefix caches when weights change. On the native async vLLM in-flight path, caches are cleared while generation is paused, so preserved requests recompute their KV after resuming. This repeats prefill for those requests and can reduce throughput. SGLang always flushes KV inside its weight synchronizer before the transfer, regardless of this flag; on its in-flight path, the collector skips a second flush after generation resumes. Other refit paths keep their existing post-update invalidation behavior. Set the flag to `false` to retain pre-update KV caches and avoid this recomputation cost on backends that support doing so. On the Megatron generation backend, the flag must agree with `policy.generation.mcore_generation_config.kv_cache_management_mode`: use `true` with `recompute`, or explicitly set `false` with `persist`/`offload`; setup errors on a mismatch. ## Why Importance Sampling Correction Is Required for Async diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml index 523e1ef32c7..de6ba643d50 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml @@ -9,6 +9,7 @@ grpo: async_grpo: enabled: true in_flight_weight_updates: true + recompute_kv_cache_after_weight_updates: false policy: model_name: nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 is_vlm: true diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index cb3fb34e18e..87001ca2a30 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -1351,13 +1351,15 @@ def _spinup_nemo_gym(base_urls, model_name, policy_api_key=None): policy_config["megatron_cfg"]["train_iters"] = total_train_iters # Megatron generation expresses recompute-after-refit engine-side via - # `kv_cache_management_mode="recompute"`; the loop-level flag must agree. - if generation_config["backend"] == "megatron": - async_grpo_config = grpo_config.async_grpo - recompute_kv_cache = bool( - async_grpo_config is not None - and async_grpo_config.recompute_kv_cache_after_weight_updates - ) + # `kv_cache_management_mode="recompute"`; the async loop's flag must agree. + # Synchronous GRPO does not consume this flag. + async_grpo_config = grpo_config.async_grpo + if ( + generation_config["backend"] == "megatron" + and async_grpo_config is not None + and async_grpo_config.enabled + ): + recompute_kv_cache = async_grpo_config.recompute_kv_cache_after_weight_updates kv_cache_mode = generation_config["mcore_generation_config"][ "kv_cache_management_mode" ] diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index a34ef255865..f17704455f9 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -725,13 +725,16 @@ def invalidate_kv_cache(self) -> bool: def pause_generation_for_refit( self, *, clear_cache: bool, timeout_s: Optional[float] = None ) -> bool: - """Pause in-flight generation while preserving request state. + """Arrange a safe generation pause around the upcoming weight refit. Backends with native in-flight refit support override this hook. The default implementation warns once per backend type and lets the refit continue with - the backend's existing in-flight behavior. On supported backends, in-flight - requests are frozen rather than aborted and resume from - :meth:`resume_generation_after_refit`; new requests queue until then. + the backend's existing in-flight behavior. Native pause implementations + freeze requests until :meth:`resume_generation_after_refit`. A backend + may instead acknowledge a lifecycle owned by its weight synchronizer, + which must pause and handle the requested cache clearing before transfer + and resume before returning. Such backends document their request handling + and timeout semantics; for example, SGLang may retract or abort requests. Args: clear_cache: Also clear the engine's reusable caches at pause time so @@ -739,8 +742,9 @@ def pause_generation_for_refit( timeout_s: Maximum time to wait for engines to pause; None waits indefinitely. Returns: - True if every engine paused; False when the backend has no native pause - support. Backends with native support raise when pausing fails. + True if every engine paused or the weight synchronizer guarantees the + lifecycle above; False when neither is supported. Pause failures raise + from this hook or the synchronizer that owns the pause. """ _warn_unsupported_in_flight_refit_pause_once(type(self).__name__) return False @@ -758,9 +762,10 @@ def resume_generation_after_refit( timeout_s: Maximum time to wait for engines to resume; None waits indefinitely. Returns: - True if every surviving engine resumed; False when the backend has no - native resume support. With fleet health, confirmed dead engines may be - excluded if enough survivors remain. Other resume failures raise. + True if every surviving engine resumed, including when the weight + synchronizer already performed the resume; False when neither path is + supported. With fleet health, confirmed dead engines may be excluded + if enough survivors remain. Other resume failures raise. """ _warn_unsupported_in_flight_refit_pause_once(type(self).__name__) return False diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py index fa6146ba7fd..cbaf3ec6584 100644 --- a/nemo_rl/models/generation/megatron/config.py +++ b/nemo_rl/models/generation/megatron/config.py @@ -170,9 +170,9 @@ class MCoreGenerationSpecificArgs(TypedDict): multimodal_prompt_config: NotRequired[MultimodalPromptConfigOverrides] # KV cache lifecycle across suspend/resume: - # - "persist": cache stays allocated; CUDA graphs remain valid (default) + # - "persist": cache stays allocated; CUDA graphs remain valid # - "offload": cache is moved off-GPU between iterations - # - "recompute": cache is dropped and rebuilt on resume + # - "recompute": cache is dropped and rebuilt on resume (default) kv_cache_management_mode: Literal["persist", "offload", "recompute"] logging_step_interval: NotRequired[int] diff --git a/nemo_rl/models/generation/sglang/sglang_generation.py b/nemo_rl/models/generation/sglang/sglang_generation.py index 6b2d36536f3..b40b0212b4b 100644 --- a/nemo_rl/models/generation/sglang/sglang_generation.py +++ b/nemo_rl/models/generation/sglang/sglang_generation.py @@ -433,6 +433,31 @@ def continue_generation(self) -> None: return ray.get([e.continue_generation.remote() for e in engines]) + def pause_generation_for_refit( + self, *, clear_cache: bool, timeout_s: Optional[float] = None + ) -> bool: + """Acknowledge the pause owned by the SGLang weight synchronizer. + + Both SGLang transports pause and flush KV inside ``sync_weights`` before + transferring weights. This hook issues no RPC; returning True prevents + callers from flushing again after the synchronizer resumes generation. + The synchronizer always clears KV, including when ``clear_cache=False``. + ``timeout_s`` does not add a deadline to that existing lifecycle. + """ + del clear_cache, timeout_s + return True + + def resume_generation_after_refit( + self, *, timeout_s: Optional[float] = None + ) -> bool: + """Acknowledge that the SGLang synchronizer already resumed generation. + + This hook issues no RPC and adds no deadline; ``sync_weights`` owns the + pause, cache flush, weight transfer, and resume as one lifecycle. + """ + del timeout_s + return True + def begin_weight_update(self) -> None: """Open a weight-update session on every node-0 engine.""" engines = [e for e in self.engines if e is not None] diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index b73aed39d8c..6195aa56326 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -939,8 +939,8 @@ def set_refit_membership(self, membership: "RefitMembership") -> None: """ self._refit_membership = membership - def _refit_leader_workers(self) -> list[Any]: - """DP leaders that should receive refit calls, in rank order. + def _refit_leaders_by_shard(self) -> list[tuple[int, int, Any]]: + """Snapshot (shard index, leader index, worker) in refit rank order. Falls back to every leader when no membership has been recorded, which is the state for the entire life of a run that never loses a shard. @@ -951,18 +951,25 @@ def _refit_leader_workers(self) -> list[Any]: membership = self._refit_membership if membership is None: per_shard = len(workers) // self.dp_size - return [workers[idx * per_shard] for idx in range(self.dp_size)] + shard_indices = range(self.dp_size) + else: + per_shard = membership.workers_per_shard + shard_indices = membership.shard_prefixes leaders = [] - for shard_idx in membership.shard_prefixes: - leader_idx = shard_idx * membership.workers_per_shard + for shard_idx in shard_indices: + leader_idx = shard_idx * per_shard if leader_idx >= len(workers): raise RuntimeError( f"shard {shard_idx} maps to worker {leader_idx}, but the group has " f"{len(workers)} workers" ) - leaders.append(workers[leader_idx]) + leaders.append((shard_idx, leader_idx, workers[leader_idx])) return leaders + def _refit_leader_workers(self) -> list[Any]: + """DP leaders that should receive refit calls, in rank order.""" + return [worker for _, _, worker in self._refit_leaders_by_shard()] + def rebuild_collective( self, membership: "RefitMembership", ip: str, port: int ) -> list[ray.ObjectRef]: @@ -1767,19 +1774,8 @@ def resume_generation_after_refit( if not self.worker_group or not self.worker_group.workers: raise RuntimeError("Worker group is not initialized") - membership = self._refit_membership - shard_indices = ( - list(membership.shard_prefixes) - if membership is not None - else list(range(self.dp_size)) - ) - per_shard = ( - membership.workers_per_shard - if membership is not None - else len(self.worker_group.workers) // self.dp_size - ) - leaders = self._refit_leader_workers() - futures = [worker.resume_generation_async.remote() for worker in leaders] + leaders = self._refit_leaders_by_shard() + futures = [worker.resume_generation_async.remote() for _, _, worker in leaders] # Settle all calls within one budget before inspecting failures. A dead leader # must not prevent us from waiting for the surviving engines to resume. _, pending = ray.wait(futures, num_returns=len(futures), timeout=timeout_s) @@ -1787,14 +1783,14 @@ def resume_generation_after_refit( raise RefitAborted( f"vLLM generation resume did not return within {timeout_s}s" ) - for shard_idx, worker, future in zip(shard_indices, leaders, futures): + for (shard_idx, leader_idx, worker), future in zip(leaders, futures): try: resumed = ray.get(future) except ray.exceptions.ActorDiedError as exc: if self.fleet_monitor is None: raise # A late error from an old actor must not condemn its replacement. - if self.worker_group.workers[shard_idx * per_shard] is worker: + if self.worker_group.workers[leader_idx] is worker: self.fleet_monitor.record_actor_death(shard_idx, error=str(exc)) continue if not resumed: diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index c556c799b6d..0c207ef5500 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -3072,6 +3072,66 @@ def test_dapo_dynamic_sampling_filters_on_raw_metric_after_overlong_shaping( assert surviving_prompts == ["prompt_1", "prompt_1", "prompt_1"] +@pytest.mark.parametrize("kv_cache_mode", ["persist", "offload", "recompute"]) +@pytest.mark.parametrize( + ("async_enabled", "recompute"), + [(None, None), (False, False), (False, True), (True, False), (True, True)], +) +def test_setup_megatron_kv_consistency_only_applies_to_async_grpo( + mock_grpo_components: dict[str, Any], + monkeypatch: pytest.MonkeyPatch, + kv_cache_mode: str, + async_enabled: bool | None, + recompute: bool | None, +) -> None: + master_config = mock_grpo_components["master_config"] + master_config.grpo.val_period = 0 + master_config.grpo.batch_multiplier = 1 + master_config.grpo.async_grpo = ( + None + if async_enabled is None + else AsyncGRPOConfig( + enabled=async_enabled, + recompute_kv_cache_after_weight_updates=recompute, + ) + ) + master_config.data.update(shuffle=False, num_workers=0) + master_config.policy.update( + model_name="test-model", + automodel_cfg={"enabled": False}, + megatron_cfg={"enabled": False}, + ) + master_config.policy["generation"].update( + backend="megatron", + mcore_generation_config={"kv_cache_management_mode": kv_cache_mode}, + ) + checkpointer = MagicMock() + checkpointer.get_latest_checkpoint_path.return_value = None + checkpointer.load_training_info.return_value = None + checkpointer.get_resume_paths.return_value = (None, None) + monkeypatch.setattr(grpo_mod, "Logger", MagicMock()) + monkeypatch.setattr(grpo_mod, "CheckpointManager", lambda _config: checkpointer) + monkeypatch.setattr(grpo_mod, "StatefulDataLoader", MagicMock()) + monkeypatch.setattr(grpo_mod, "RayVirtualCluster", MagicMock()) + monkeypatch.setattr( + grpo_mod, "prepare_segment_topology", lambda *_args: (None, [], {}) + ) + + class PolicyInitReached(Exception): + pass + + policy_init = MagicMock(side_effect=PolicyInitReached) + monkeypatch.setattr(grpo_mod, "Policy", policy_init) + mismatch = async_enabled and recompute != (kv_cache_mode == "recompute") + with ( + pytest.raises(ValueError, match="recompute_kv_cache_after_weight_updates") + if mismatch + else pytest.raises(PolicyInitReached) + ): + grpo_mod.setup(master_config, MagicMock(), MagicMock(), None) + assert policy_init.call_count == int(not mismatch) + + @pytest.mark.parametrize( ("kv_cache_dtype", "precision", "error"), [ diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index 321da7f05ba..3a839194775 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -622,8 +622,8 @@ def test_vllm_generation_broadcasts_native_refit_pause_and_resume( ) -> None: generation = VllmGeneration.__new__(VllmGeneration) generation.cfg = {"vllm_cfg": {"async_engine": True}} - generation.worker_group = MagicMock(workers=[object(), object()]) leaders = [MagicMock(), MagicMock()] + generation.worker_group = MagicMock(workers=leaders) pause_futures = [object(), object()] resume_futures = [object(), object()] for worker, pause_future, resume_future in zip( @@ -631,7 +631,6 @@ def test_vllm_generation_broadcasts_native_refit_pause_and_resume( ): worker.pause_generation_async.remote.return_value = pause_future worker.resume_generation_async.remote.return_value = resume_future - generation._refit_leader_workers = MagicMock(return_value=leaders) generation._refit_membership = None generation.dp_size = 2 generation.fleet_monitor = None @@ -661,9 +660,8 @@ def test_vllm_generation_refit_pause_and_resume_raise_on_timeout( ) -> None: generation = VllmGeneration.__new__(VllmGeneration) generation.cfg = {"vllm_cfg": {"async_engine": True}} - generation.worker_group = MagicMock(workers=[object()]) worker = MagicMock() - generation._refit_leader_workers = MagicMock(return_value=[worker]) + generation.worker_group = MagicMock(workers=[worker]) generation._refit_membership = None generation.dp_size = 1 ray_get = MagicMock(side_effect=ray.exceptions.GetTimeoutError("timed out")) @@ -691,9 +689,8 @@ def test_vllm_generation_rejects_partial_refit_pause_and_resume( ) -> None: generation = VllmGeneration.__new__(VllmGeneration) generation.cfg = {"vllm_cfg": {"async_engine": True}} - generation.worker_group = MagicMock(workers=[object(), object()]) leaders = [MagicMock(), MagicMock()] - generation._refit_leader_workers = MagicMock(return_value=leaders) + generation.worker_group = MagicMock(workers=leaders) generation._refit_membership = None generation.dp_size = 2 monkeypatch.setattr(ray, "wait", MagicMock(return_value=([], []))) diff --git a/tests/unit/models/generation/test_vllm_refit_resume.py b/tests/unit/models/generation/test_vllm_refit_resume.py index b19b6f6cf71..9ebb156a655 100644 --- a/tests/unit/models/generation/test_vllm_refit_resume.py +++ b/tests/unit/models/generation/test_vllm_refit_resume.py @@ -99,6 +99,52 @@ def replace_actor( assert monitor.serving_shards() == [2] +@pytest.mark.parametrize("membership_shards", [None, (2, 0)]) +def test_resume_maps_leader_failures_to_the_original_shard( + generation: VllmGeneration, + monkeypatch: pytest.MonkeyPatch, + membership_shards: tuple[int, ...] | None, +) -> None: + if membership_shards is None: + generation._refit_membership = None + generation.fleet_monitor = GenerationFleetHealth( + shard_count=3, policy=FleetHealthPolicy(min_healthy_shards=1) + ) + expected_shards = (0, 1, 2) + else: + generation._refit_membership.shard_prefixes = { + shard_idx: rank * 2 for rank, shard_idx in enumerate(membership_shards) + } + expected_shards = membership_shards + + workers = generation.worker_group.workers + futures = [ + workers[shard_idx * 2].resume_generation_async.remote.return_value + for shard_idx in expected_shards + ] + dead_future = workers[4].resume_generation_async.remote.return_value + + def result(ref: object) -> bool: + if ref is dead_future: + raise ray.exceptions.ActorDiedError() + return True + + monkeypatch.setattr(ray, "get", MagicMock(side_effect=result)) + + assert generation.resume_generation_after_refit(timeout_s=0.25) + + ray.wait.assert_called_once_with(futures, num_returns=len(futures), timeout=0.25) + for worker_idx, worker in enumerate(workers): + if worker_idx in {shard_idx * 2 for shard_idx in expected_shards}: + worker.resume_generation_async.remote.assert_called_once_with() + else: + worker.resume_generation_async.remote.assert_not_called() + assert generation.fleet_monitor.state_of(2) is ShardState.DEAD + assert generation.fleet_monitor.serving_shards() == sorted( + shard_idx for shard_idx in expected_shards if shard_idx != 2 + ) + + def test_resume_fails_if_the_remaining_fleet_is_exhausted( generation: VllmGeneration, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index fb29a410941..cf8b7dc35ea 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -463,11 +463,13 @@ def get_step_metrics(self) -> dict[str, float]: def pause_generation_for_refit( self, *, clear_cache: bool, timeout_s: float | None = None ) -> bool: + # Model native pause/clear support; unsupported backends are covered in + # test_single_controller_actor.py rather than these checkpoint tests. del clear_cache, timeout_s - return False + return True def resume_generation_after_refit(self, *, timeout_s: float | None = None) -> bool: - return False + return True def invalidate_kv_cache(self) -> bool: return False diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index 33accc7b170..c1fa7273a3a 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -75,12 +75,17 @@ from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.experience.rollout_reassembler_actor import RolloutReassemblerActor from nemo_rl.experience.rollout_recovery import RolloutRecoveryLedger +from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration from nemo_rl.models.generation.vllm.vllm_worker_async import ( VllmAsyncGenerationWorkerImpl, ) from nemo_rl.utils.logger import LoggerConfig from nemo_rl.utils.timer import TimeoutChecker, Timer from nemo_rl.utils.train_data_dump import TrainDataDump +from nemo_rl.weight_sync.sglang_weight_synchronizer import ( + SGLangColocatedWeightSynchronizer, + SGLangDisaggregatedWeightSynchronizer, +) class FakeWeightSynchronizer: @@ -881,6 +886,101 @@ async def test_sync_weights_honors_recompute_kv_cache_config( assert ctrl._rollout_permitted.is_set() +@pytest.mark.parametrize( + "synchronizer_cls", + [SGLangColocatedWeightSynchronizer, SGLangDisaggregatedWeightSynchronizer], +) +@pytest.mark.parametrize("recompute_kv_cache", [False, True]) +def test_sync_weights_sglang_flushes_only_while_paused( + monkeypatch: pytest.MonkeyPatch, + synchronizer_cls: type[ + SGLangColocatedWeightSynchronizer | SGLangDisaggregatedWeightSynchronizer + ], + recompute_kv_cache: bool, +) -> None: + """Use the real hooks and synchronizer; a flush after resume must fail.""" + events: list[str] = [] + paused = False + + def pause(*, mode: str) -> None: + nonlocal paused + assert mode == "retract" + paused = True + events.append("pause") + + def flush() -> None: + events.append("flush") + if not paused: + raise RuntimeError("SGLang refuses flush_cache with running requests") + + def resume() -> None: + nonlocal paused + paused = False + events.append("resume") + + def transfer(**_: Any) -> list: + assert paused + assert events[-1] == "begin" + events.append("transfer") + return [] + + engine = MagicMock() + engine.pause_generation.remote.side_effect = pause + engine.invalidate_kv_cache.remote.side_effect = flush + engine.continue_generation.remote.side_effect = resume + engine.begin_weight_update.remote.side_effect = lambda: events.append("begin") + engine.end_weight_update.remote.side_effect = lambda: events.append("end") + monkeypatch.setattr("ray.get", lambda refs: refs) + + gen = object.__new__(SGLangGeneration) + # No server/router was started; avoid resource cleanup in __del__. + gen.shutdown = MagicMock() + gen.all_engines = [engine] + gen.num_new_engines = 0 + gen.num_gpus_per_engine = 1 + gen.num_gpus_per_node = 1 + gen.gpu_offset = 0 + gen.pause_generation_mode = "retract" + gen.needs_offload = False + gen._health_monitor = None + gen.rollout_engine_lock = None + gen.sglang_cfg = { + "sglang_cfg": { + "quantization": {"scheme": "bf16"}, + "sglang_fault_tolerance_config": {"use_fault_tolerance": False}, + } + } + trainer = MagicMock() + trainer.update_weights_to_sglang_colocated.side_effect = transfer + trainer.update_weights_to_sglang_distributed.side_effect = transfer + + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + ctrl = object.__new__(controller_cls) + ctrl._async_cfg = AsyncRLConfig( + recompute_kv_cache_after_weight_updates=recompute_kv_cache + ) + ctrl._rollout_permitted = asyncio.Event() + ctrl._rollout_permitted.set() + ctrl._gen_fleet = None + ctrl._gen = gen + ctrl._trainer = trainer + ctrl._weight_synchronizer = synchronizer_cls(trainer, gen, refit_buffer_size_gb=1) + ctrl._timer = Timer() + ctrl._rollout_manager = SimpleNamespace(resume_request_deadlines=MagicMock()) + ctrl._inflight_by_group_id = {} + ctrl._rollout_recovery_enabled = False + ctrl._master_config = SimpleNamespace( + env={}, token_capture=SimpleNamespace(enabled=False) + ) + + assert asyncio.run(ctrl._sync_weights()) == 0 + + assert events == ["pause", "flush", "begin", "transfer", "end", "resume"] + engine.invalidate_kv_cache.remote.assert_called_once_with() + assert ctrl._rollout_permitted.is_set() + assert not paused + + def test_sync_weights_rejects_unconfirmed_inflight_kv_recomputation() -> None: controller_cls = SingleControllerActor.__ray_metadata__.modified_class ctrl = object.__new__(controller_cls) diff --git a/tests/unit/single_controller/test_train_pump_e2e.py b/tests/unit/single_controller/test_train_pump_e2e.py index d9cf7efb071..26e3a8385bd 100644 --- a/tests/unit/single_controller/test_train_pump_e2e.py +++ b/tests/unit/single_controller/test_train_pump_e2e.py @@ -188,7 +188,7 @@ def compute_advantage( class _FakeGeneration: - """Unsupported generation backend for this training-pump-only test.""" + """Generation stand-in with native pause/clear support for the training pump.""" def blocks_training(self) -> bool: return False @@ -203,10 +203,10 @@ def pause_generation_for_refit( self, *, clear_cache: bool, timeout_s: float | None = None ) -> bool: del clear_cache, timeout_s - return False + return True def resume_generation_after_refit(self, *, timeout_s: float | None = None) -> bool: - return False + return True @ray.remote(num_cpus=0) # pragma: no cover diff --git a/tests/unit/test_recipe_settings.py b/tests/unit/test_recipe_settings.py index ff1a8c49eef..7010235f5eb 100644 --- a/tests/unit/test_recipe_settings.py +++ b/tests/unit/test_recipe_settings.py @@ -16,6 +16,7 @@ from pathlib import Path import pytest +from omegaconf import DictConfig, OmegaConf from nemo_rl.utils.config import ( load_config, @@ -24,36 +25,89 @@ ) RECIPE_ROOT = Path(__file__).parents[2] / "examples/configs/recipes" +FUNCTIONAL_ROOT = Path(__file__).parents[2] / "tests/functional" +MEGATRON_FUNCTIONAL_SCRIPTS = sorted( + { + path.name + for path in FUNCTIONAL_ROOT.glob("*.sh") + if "policy.generation.backend=megatron" in path.read_text() + } + | {"grpo_megatron_generation_gym_single_controller.sh"} +) -@pytest.mark.parametrize( - "script_name", - [ - "grpo_megatron_generation_gym_single_controller.sh", - "grpo_megatron_generation_colocated_reshard_gym_single_controller.sh", - ], -) -def test_megatron_single_controller_gym_scripts_recompute_kv(script_name: str) -> None: +def _load_megatron_functional_config(script_name: str) -> tuple[DictConfig, bool]: register_omegaconf_resolvers() project_root = Path(__file__).parents[2] - script = (project_root / "tests/functional" / script_name).read_text() - command = next( - line + script = (FUNCTIONAL_ROOT / script_name).read_text() + commands = [ + shlex.split(line) for line in script.replace("\\\n", " ").splitlines() - if line.startswith("uv run coverage run ") - ) - arguments = shlex.split(command) - config_path = arguments[arguments.index("--config") + 1].replace( - "$PROJECT_ROOT", str(project_root) + if line.lstrip().startswith("uv run ") + and any(name in line for name in ("run_grpo", "run_vlm_grpo", "$SC_ENTRYPOINT")) + ] + arguments = next( + args + for args in commands + if any( + "run_grpo" in arg or "run_vlm_grpo" in arg or arg == "$SC_ENTRYPOINT" + for arg in args + ) ) + if "--config" in arguments: + config_path = arguments[arguments.index("--config") + 1].replace( + "$PROJECT_ROOT", str(project_root) + ) + else: + assert any(arg.endswith("run_grpo_sliding_puzzle.py") for arg in arguments) + config_path = str(project_root / "examples/configs/grpo_sliding_puzzle.yaml") overrides = [ argument for argument in arguments - if argument.lstrip("+").startswith(("policy.generation.", "async_rl.")) + if argument.lstrip("+").startswith( + ( + "policy.generation.backend=", + "policy.generation.mcore_generation_config.kv_cache_management_mode=", + "async_rl.", + "grpo.async_grpo.", + ) + ) ] config = parse_hydra_overrides(load_config(config_path), overrides) + single_controller = any( + "run_grpo_single_controller" in arg or arg == "$SC_ENTRYPOINT" + for arg in arguments + ) + return config, single_controller + + +@pytest.mark.parametrize("script_name", MEGATRON_FUNCTIONAL_SCRIPTS) +def test_megatron_functional_scripts_keep_kv_settings_consistent( + script_name: str, +) -> None: + # Load algorithm schemas only for these tests; other recipe tests need no Ray/Torch. + from nemo_rl.algorithms.grpo import AsyncGRPOConfig + from nemo_rl.algorithms.single_controller_utils.config import AsyncRLConfig + + config, single_controller = _load_megatron_functional_config(script_name) assert config.policy.generation.backend == "megatron" - assert ( + async_path = "async_rl" if single_controller else "grpo.async_grpo" + raw_async_config = OmegaConf.select(config, async_path) + async_values = ( + OmegaConf.to_container(raw_async_config, resolve=True) + if raw_async_config is not None + else {} + ) + if single_controller: + async_config = AsyncRLConfig.model_validate(async_values) + else: + if raw_async_config is None: + return + async_config = AsyncGRPOConfig.model_validate(async_values) + if not async_config.enabled: + return + # Preserve intentional persist/offload opt-outs as well as the recompute default. + assert async_config.recompute_kv_cache_after_weight_updates == ( config.policy.generation.mcore_generation_config.kv_cache_management_mode == "recompute" ) From 9c15e56804cb3054991b4a97477ab6c6eab2d354 Mon Sep 17 00:00:00 2001 From: Youngeun Kwon Date: Sat, 17 Oct 2026 11:49:46 -0700 Subject: [PATCH 12/13] test: resolve functional recipe paths from repository root Signed-off-by: Youngeun Kwon --- tests/unit/test_recipe_settings.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/unit/test_recipe_settings.py b/tests/unit/test_recipe_settings.py index 7010235f5eb..edba239166a 100644 --- a/tests/unit/test_recipe_settings.py +++ b/tests/unit/test_recipe_settings.py @@ -61,6 +61,7 @@ def _load_megatron_functional_config(script_name: str) -> tuple[DictConfig, bool else: assert any(arg.endswith("run_grpo_sliding_puzzle.py") for arg in arguments) config_path = str(project_root / "examples/configs/grpo_sliding_puzzle.yaml") + config_path = project_root / config_path overrides = [ argument for argument in arguments @@ -84,11 +85,15 @@ def _load_megatron_functional_config(script_name: str) -> tuple[DictConfig, bool @pytest.mark.parametrize("script_name", MEGATRON_FUNCTIONAL_SCRIPTS) def test_megatron_functional_scripts_keep_kv_settings_consistent( script_name: str, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, ) -> None: # Load algorithm schemas only for these tests; other recipe tests need no Ray/Torch. from nemo_rl.algorithms.grpo import AsyncGRPOConfig from nemo_rl.algorithms.single_controller_utils.config import AsyncRLConfig + # CI runs from tests/, while the launchers resolve paths from the repo root. + monkeypatch.chdir(tmp_path) config, single_controller = _load_megatron_functional_config(script_name) assert config.policy.generation.backend == "megatron" async_path = "async_rl" if single_controller else "grpo.async_grpo" From 064a49edbb8bfab325442828e61019557a0a0f20 Mon Sep 17 00:00:00 2001 From: Youngeun Kwon Date: Sat, 17 Oct 2026 19:54:00 -0700 Subject: [PATCH 13/13] test: bind refit shard helper in reshard rebuild fixture Signed-off-by: Youngeun Kwon --- tests/unit/weight_sync/test_reshard_rebuild.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/weight_sync/test_reshard_rebuild.py b/tests/unit/weight_sync/test_reshard_rebuild.py index 1ddda729bd5..d8e961772c1 100644 --- a/tests/unit/weight_sync/test_reshard_rebuild.py +++ b/tests/unit/weight_sync/test_reshard_rebuild.py @@ -93,6 +93,7 @@ def _reshard(dp_size=4, workers_per_shard=1, dead_shards=(), train_world_size=8) "prepare_nccl_reshard_refit_info", "set_refit_membership", "nccl_reshard_refit", + "_refit_leaders_by_shard", "_refit_leader_workers", ): setattr(