diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 60f09cef4ac..db6d7d0c4d5 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -2528,15 +2528,8 @@ def refit_policy_generation( Returns: Scalar metrics reported by the selected weight synchronizer. """ - # Every SGLang deployment reaches its refit through this hook: `setup` - # attaches an SGLang synchronizer that owns the whole lifecycle (phase - # transitions, engine recovery, pause/flush, transport), so SGLang never - # touches the branches below. synchronizer = getattr(policy_generation, "weight_synchronizer", None) - if synchronizer is not None: - return synchronizer.sync_weights(timer=timer, kv_scales=kv_scales) or {} - - if isinstance(policy_generation, SGLangGeneration): + if isinstance(policy_generation, SGLangGeneration) and synchronizer is None: # Fail loudly rather than falling through to the vLLM branches, which # would call methods the SGLang path does not implement. raise RuntimeError( @@ -2544,6 +2537,19 @@ def refit_policy_generation( "set. Attach one with create_weight_synchronizer(...) during setup." ) + # Materialize deferred Megatron parameter all-gathers before any transport + # reads policy weights, including synchronizers that return early below. + sync_context = ( + timer.time("prepare_for_generation/sync_policy_params") + if timer is not None + else nullcontext() + ) + with sync_context: + policy.sync_params_before_refit() + + if synchronizer is not None: + return synchronizer.sync_weights(timer=timer, kv_scales=kv_scales) or {} + if colocated_inference: policy.offload_before_refit() policy_generation.prepare_for_generation(tags=["weights"]) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 36dd5ffd103..0feef02f9c1 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -4008,8 +4008,9 @@ async def _sync_weights( 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 + 3. Materialize deferred policy parameter all-gathers. + 4. weight_synchronizer.sync_weights(kv_scales=...) + 5. _rollout_permitted.set() — resume Args: calibration_data: Optional data used to calibrate FP8 KV-cache @@ -4063,6 +4064,11 @@ async def _sync_weights( # once a shard was gone, because absent_shards() never empties again. await self._reconcile_refit_membership() + # Recovery may repeat the transport, but an optimizer update only needs + # one parameter all-gather, so keep this outside the retry block. + with self._timer.time("prepare_for_generation/sync_policy_params"): + await asyncio.to_thread(self._trainer.sync_params_before_refit) + try: await self._sync_weights_within(kv_scales, "first") except (RefitAborted, RayActorError) as failure: diff --git a/nemo_rl/models/policy/workers/base_policy_worker.py b/nemo_rl/models/policy/workers/base_policy_worker.py index d2fa634c09f..19d17d01d73 100644 --- a/nemo_rl/models/policy/workers/base_policy_worker.py +++ b/nemo_rl/models/policy/workers/base_policy_worker.py @@ -36,7 +36,8 @@ class AbstractPolicyWorker: def sync_params_before_refit(self) -> None: """Materialize optimizer updates before refit when the backend requires it.""" # DTensor policy parameters already contain the latest optimizer update. - # Megatron overrides this for overlapped MXFP8 parameter all-gather. + # Megatron overrides this whenever the distributed optimizer overlaps the + # parameter all-gather (MXFP8 shared-buffer and plain BF16 alike). pass def init_collective( diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index f436b7424ae..409c5344cff 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -4109,35 +4109,33 @@ def _clear_fp8_caches(self): @torch.no_grad() @wrap_with_nvtx_name("megatron_policy_worker/sync_params_before_refit") def sync_params_before_refit(self) -> None: - """Materialize optimizer updates before a refit reads model parameters.""" - # With MXFP8 overlap, the optimizer updates FP32 master shards and the - # next parameter all-gather requantizes them into the model weights. A - # refit happens between optimizer steps, before that next training - # forward, so force the gather now. This both gives generation the - # latest weights and leaves hooks disabled while the shared param/grad - # buffer is held across refit. The normal train-step transition - # re-enables them. - # Deliberately conditional on the hooks being enabled. Every state in - # which they are already off is one where the weights are current - # anyway: before the first train step the buffer holds the checkpoint; - # eval entry already forced a sync via disable_forward_pre_hook( - # param_sync=True) and runs no optimizer step; a skipped step leaves the - # masters unchanged; and a successful step re-enables the hooks before - # returning. If a stale case is ever found, note that staging alone does - # NOT fix it - with reuse_grad_buf_for_mxfp8_param_ag param_data aliases - # grad_data, so zero_grad_buffer() wipes the parameters and - # _copy_main_params_to_param_buffer restores only this rank's shard, - # leaving every other DP rank at zero. Upstream pairs that staging with a - # following start_param_sync (DistributedOptimizer. - # prepare_model_params_for_param_sync); any fix needs the sync too. + """Materialize optimizer updates before a refit reads model parameters. + + ``overlap_param_gather`` defers this work to the next training forward, + but refit reads the parameters first. Megatron-FSDP handles this in its + own module hooks and is intentionally not handled here. + """ if ( - self._uses_mxfp8_overlap_shared_param_buffer() - and self._forward_pre_hook_enabled() + not isinstance(self.model, DistributedDataParallel) + or not self.model.ddp_config.overlap_param_gather + or not self._forward_pre_hook_enabled() ): - # An in-flight async checkpoint may still read these tensors, so settle it - # before the explicit gather mutates the shared parameter buffer. + # Disabled hooks mean no optimizer update is waiting to be gathered. + return + + if self._uses_mxfp8_overlap_shared_param_buffer(): + # This path requantizes updated master shards into the shared buffer. + # Hold that buffer materialized until the next training step. self.finalize_async_save() self._disable_forward_pre_hook_until_next_train_step(param_sync=True) + return + + # BF16 master shards are already in the DDP parameter buffer; only the + # all-gather remains. Settle checkpoint reads before rewriting it. + self.finalize_async_save() + self.model.start_param_sync(force_sync=True) + # Ensure exporters cannot observe a partially gathered buffer. + torch.cuda.synchronize() @wrap_with_nvtx_name("megatron_policy_worker/offload_before_refit") def offload_before_refit(self): diff --git a/nemo_rl/weight_sync/collective_weight_synchronizer.py b/nemo_rl/weight_sync/collective_weight_synchronizer.py index 6dd75bae755..160c6b9cc16 100644 --- a/nemo_rl/weight_sync/collective_weight_synchronizer.py +++ b/nemo_rl/weight_sync/collective_weight_synchronizer.py @@ -20,10 +20,9 @@ established NCCL process group. Lifecycle per sync: - 1. policy.sync_params_before_refit() -- materialize optimizer updates - 2. policy.broadcast_weights_for_collective() -- send via NCCL + 1. policy.broadcast_weights_for_collective() -- send via NCCL generation.update_weights_from_collective() -- receive via NCCL - 3. Verify transfer success + 2. Verify transfer success No offload/restore steps are needed since policy and generation run on separate GPUs with dedicated memory. @@ -92,9 +91,6 @@ class CollectiveWeightSynchronizer(WeightSynchronizer): arms a watchdog and aborts its own communicator when it expires, which is what lets the controller rebuild over the survivors instead of blocking in NCCL forever. ``None`` disarms it entirely, so the hang protection is lost. - sync_policy_params: Whether this synchronizer owns the pre-transfer policy - parameter sync. A lifecycle wrapper may perform it earlier and disable it - here to avoid a duplicate worker round trip. """ def __init__( @@ -104,8 +100,6 @@ def __init__( train_cluster: Any, inference_cluster: Any, refit_timeout_s: Optional[float] = None, - *, - sync_policy_params: bool = True, ): # None disarms the abort watchdog in every worker, which is the default and # reproduces the pre-existing behaviour exactly. @@ -114,7 +108,6 @@ def __init__( self._generation = generation self._train_cluster = train_cluster self._inference_cluster = inference_cluster - self._sync_policy_params = sync_policy_params self._stale = True # The absent set this synchronizer's current communicator was built with, so a # membership that has not changed can skip the rebuild. None means "never rebuilt", @@ -135,8 +128,6 @@ def sync_weights( timer: Optional[Timer] = None, kv_scales: Optional[dict[str, float]] = None, ) -> None: - if self._sync_policy_params: - self._policy.sync_params_before_refit() timer_context = ( timer.time("prepare_for_generation/transfer_and_update_weights") if timer is not None diff --git a/nemo_rl/weight_sync/ipc_weight_synchronizer.py b/nemo_rl/weight_sync/ipc_weight_synchronizer.py index 8311012f8ce..a1f4dcb9521 100644 --- a/nemo_rl/weight_sync/ipc_weight_synchronizer.py +++ b/nemo_rl/weight_sync/ipc_weight_synchronizer.py @@ -19,13 +19,12 @@ transport for colocated vLLM deployments. Lifecycle per sync: - 1. policy.sync_params_before_refit() -- materialize optimizer updates - 2. policy.offload_before_refit() -- free GPU for weight staging - 3. generation.prepare_for_generation(tags=["weights"]) -- allocate buffers - 4. policy.stream_weights_via_ipc_zmq() -- send weights via ZMQ + 1. policy.offload_before_refit() -- free GPU for weight staging + 2. generation.prepare_for_generation(tags=["weights"]) -- allocate buffers + 3. policy.stream_weights_via_ipc_zmq() -- send weights via ZMQ generation.update_weights_via_ipc_zmq() -- receive weights - 5. policy.offload_after_refit() -- restore optimizer state - 6. generation.prepare_for_generation(tags=["kv_cache"]) -- rebuild KV cache + 4. policy.offload_after_refit() -- restore optimizer state + 5. generation.prepare_for_generation(tags=["kv_cache"]) -- rebuild KV cache """ import os @@ -70,7 +69,6 @@ def sync_weights( timer: Optional[Timer] = None, kv_scales: Optional[dict[str, float]] = None, ) -> None: - self._policy.sync_params_before_refit() self._policy.offload_before_refit() self._generation.prepare_for_generation(tags=["weights"]) diff --git a/nemo_rl/weight_sync/megatron_weight_synchronizer.py b/nemo_rl/weight_sync/megatron_weight_synchronizer.py index f37000d9f1f..9900f2b4a87 100644 --- a/nemo_rl/weight_sync/megatron_weight_synchronizer.py +++ b/nemo_rl/weight_sync/megatron_weight_synchronizer.py @@ -89,7 +89,6 @@ def __init__( train_cluster=train_cluster, inference_cluster=inference_cluster, refit_timeout_s=refit_timeout_s, - sync_policy_params=False, ) else: self._transport = CollectiveWeightSynchronizer( @@ -98,7 +97,6 @@ def __init__( train_cluster=train_cluster, inference_cluster=inference_cluster, refit_timeout_s=refit_timeout_s, - sync_policy_params=False, ) self._stale = True @@ -154,8 +152,6 @@ def timed_phase(name: str) -> AbstractContextManager[None]: # Tagging the call bypasses the worker's engine-awake early-return, so the reshard # copy riding this wake cannot be skipped. Any tag except "weights" works: the worker # treats "weights" as the wake-suppressing mid-refit call. - with timed_phase("prepare_for_generation/sync_policy_params"): - self._policy.sync_params_before_refit() with timed_phase("prepare_for_generation/offload_policy"): self._policy.offload_before_refit() with timed_phase("prepare_for_generation/prepare_weights"): @@ -163,12 +159,6 @@ def timed_phase(name: str) -> AbstractContextManager[None]: self._stale = False return {} - # Materialize optimizer updates before optional policy offload. Delegated - # transports skip their own copy of this prerequisite because this wrapper - # owns the Megatron generation lifecycle. - with timed_phase("prepare_for_generation/sync_policy_params"): - self._policy.sync_params_before_refit() - # The engine serves continuously in non-colocated mode; pause it exactly # around the swap. with timed_phase("prepare_for_generation/suspend_for_refit"): diff --git a/nemo_rl/weight_sync/nccl_reshard_weight_synchronizer.py b/nemo_rl/weight_sync/nccl_reshard_weight_synchronizer.py index de158302286..6ead67f30d0 100644 --- a/nemo_rl/weight_sync/nccl_reshard_weight_synchronizer.py +++ b/nemo_rl/weight_sync/nccl_reshard_weight_synchronizer.py @@ -112,9 +112,6 @@ class NcclReshardWeightSynchronizer(WeightSynchronizer): arms a watchdog and aborts its own communicator when it expires, which is what lets the controller rebuild over the survivors instead of blocking in NCCL forever. ``None`` disarms it entirely, so the hang protection is lost. - sync_policy_params: Whether this synchronizer owns the pre-transfer policy - parameter sync. A lifecycle wrapper may perform it earlier and disable it - here to avoid a duplicate worker round trip. """ def __init__( @@ -124,15 +121,12 @@ def __init__( train_cluster: Any, inference_cluster: Any, refit_timeout_s: Optional[float] = None, - *, - sync_policy_params: bool = True, ): self._policy = policy self._generation = generation self._train_cluster = train_cluster self._inference_cluster = inference_cluster self._refit_timeout_s = refit_timeout_s - self._sync_policy_params = sync_policy_params self._stale = True # The absent set this synchronizer's current communicator was built with, so a # membership that has not changed can skip the rebuild. None means "never rebuilt", @@ -197,8 +191,6 @@ def sync_weights( timer: Optional[Timer] = None, kv_scales: Optional[dict[str, float]] = None, ) -> None: - if self._sync_policy_params: - self._policy.sync_params_before_refit() timer_context = ( timer.time("prepare_for_generation/transfer_and_update_weights") if timer is not None diff --git a/nemo_rl/weight_sync/sglang_weight_synchronizer.py b/nemo_rl/weight_sync/sglang_weight_synchronizer.py index e83366da670..42c9afc8fe8 100644 --- a/nemo_rl/weight_sync/sglang_weight_synchronizer.py +++ b/nemo_rl/weight_sync/sglang_weight_synchronizer.py @@ -24,12 +24,11 @@ and own the GPU phase transitions around them. Colocated: - 1. policy.sync_params_before_refit() -- materialize updates - 2. policy.offload_before_refit() -- free GPU for staging - 3. generation.prepare_for_generation(tags=["weights"]) -- allocate buffers - 4. _refit() -- Ray CUDA-IPC transfer - 5. policy.offload_after_refit() -- restore optimizer state - 6. generation.prepare_for_generation(tags=["kv_cache"]) -- rebuild KV cache + 1. policy.offload_before_refit() -- free GPU for staging + 2. generation.prepare_for_generation(tags=["weights"]) -- allocate buffers + 3. _refit() -- Ray CUDA-IPC transfer + 4. policy.offload_after_refit() -- restore optimizer state + 5. generation.prepare_for_generation(tags=["kv_cache"]) -- rebuild KV cache Disaggregated: 1. generation.prepare_for_generation(tags=["weights"]) @@ -245,7 +244,6 @@ def sync_weights( kv_scales: Optional[dict[str, float]] = None, ) -> Optional[dict[str, float]]: self._reject_kv_scales(kv_scales) - self._policy.sync_params_before_refit() self._policy.offload_before_refit() sync_succeeded = False diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index b74cd266fcd..fc35f618d3e 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -219,21 +219,73 @@ def test_refit_policy_generation_forwards_kv_scales_on_colocated_ipc( ) -def test_megatron_m2n_refit_delegates_entirely_to_the_synchronizer() -> None: - """MegatronWeightSynchronizer owns the engine lifecycle; the caller must not duplicate it. +@patch("nemo_rl.algorithms.grpo.ray") +def test_legacy_noncolocated_refit_syncs_policy_params_first( + mock_ray: MagicMock, +) -> None: + mock_ray.get.return_value = [True] + events = [] + policy = MagicMock() + policy.sync_params_before_refit.side_effect = lambda: events.append("sync") + policy.broadcast_weights_for_collective.side_effect = lambda **_: ( + events.append("broadcast") or [MagicMock()] + ) + policy_generation = MagicMock() + policy_generation.weight_synchronizer = None + policy_generation.update_weights_from_collective.return_value = [MagicMock()] - ``refit_policy_generation`` returns as soon as a weight synchronizer is - present, so suspend/offload/prepare/resume must NOT be driven here — they - live inside ``MegatronWeightSynchronizer.sync_weights`` and are asserted in - ``tests/unit/weight_sync/test_weight_synchronizer.py``. + refit_policy_generation( + policy, + policy_generation, + colocated_inference=False, + ) + + assert events == ["sync", "broadcast"] + + +@patch("nemo_rl.algorithms.grpo.ray") +def test_legacy_colocated_refit_syncs_policy_params_before_offload( + mock_ray: MagicMock, +) -> None: + mock_ray.get.return_value = [True] + events = [] + policy = MagicMock() + policy.sync_params_before_refit.side_effect = lambda: events.append("sync") + policy.offload_before_refit.side_effect = lambda: events.append("offload") + policy.get_free_memory_bytes.return_value = 1 << 30 + policy.stream_weights_via_ipc_zmq.side_effect = lambda **_: ( + events.append("stream") or [MagicMock()] + ) + policy_generation = MagicMock() + policy_generation.weight_synchronizer = None + policy_generation.update_weights_via_ipc_zmq.return_value = [MagicMock()] + + refit_policy_generation( + policy, + policy_generation, + colocated_inference=True, + ) + + assert events == ["sync", "offload", "stream"] + + +def test_megatron_m2n_refit_syncs_params_then_delegates_to_synchronizer() -> None: + """The caller syncs params; MegatronWeightSynchronizer owns engine lifecycle. + + Suspend/offload/prepare/resume live inside the synchronizer and are asserted + in ``tests/unit/weight_sync/test_weight_synchronizer.py``. """ + events = [] policy = MagicMock() + policy.sync_params_before_refit.side_effect = lambda: events.append("sync") generation = object.__new__(MegatronGeneration) generation.suspend_for_refit = MagicMock() generation.prepare_for_generation = MagicMock() generation.resume_after_refit = MagicMock() generation.weight_synchronizer = MagicMock() - generation.weight_synchronizer.sync_weights.return_value = {"bytes": 16.0} + generation.weight_synchronizer.sync_weights.side_effect = lambda **_: ( + events.append("transfer") or {"bytes": 16.0} + ) metrics = refit_policy_generation( policy, @@ -243,6 +295,8 @@ def test_megatron_m2n_refit_delegates_entirely_to_the_synchronizer() -> None: ) assert metrics == {"bytes": 16.0} + assert events == ["sync", "transfer"] + policy.sync_params_before_refit.assert_called_once_with() generation.weight_synchronizer.sync_weights.assert_called_once_with( timer=None, kv_scales={"layer.0": 0.5} ) @@ -257,11 +311,10 @@ def test_refit_returns_empty_metrics_when_synchronizer_returns_none() -> None: generation = object.__new__(MegatronGeneration) generation.weight_synchronizer = MagicMock() generation.weight_synchronizer.sync_weights.return_value = None + policy = MagicMock() - assert ( - refit_policy_generation(MagicMock(), generation, colocated_inference=False) - == {} - ) + assert refit_policy_generation(policy, generation, colocated_inference=False) == {} + policy.sync_params_before_refit.assert_called_once_with() class TestMaskSampleFilter: diff --git a/tests/unit/algorithms/test_grpo_checkpoint_engine.py b/tests/unit/algorithms/test_grpo_checkpoint_engine.py index 0f9e1b00c99..d110a5caf72 100644 --- a/tests/unit/algorithms/test_grpo_checkpoint_engine.py +++ b/tests/unit/algorithms/test_grpo_checkpoint_engine.py @@ -79,7 +79,7 @@ def test_refit_policy_generation_uses_attached_checkpoint_engine_synchronizer(): from nemo_rl.algorithms import grpo as grpo_mod from nemo_rl.models.generation.vllm import VllmGeneration - policy = object() + policy = MagicMock() kv_scales = {"layer_0": 1.0} generation = MagicMock(spec=VllmGeneration) @@ -98,6 +98,7 @@ def test_refit_policy_generation_uses_attached_checkpoint_engine_synchronizer(): generation.weight_synchronizer.sync_weights.assert_called_once_with( timer=None, kv_scales=kv_scales ) + policy.sync_params_before_refit.assert_called_once_with() assert result == {"transfer_s": 1.0} diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index aace0f4cdb3..1cd0d86cd0f 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -1258,6 +1258,53 @@ def get_extra_state(self): raise AssertionError("moving a module must not serialize its extra state") +@pytest.mark.parametrize("hooks_enabled", [True, False]) +def test_sync_params_before_refit_gathers_pending_bf16_params( + monkeypatch, hooks_enabled +): + """Refit must see updated optimizer shards before it reads model parameters. + + The BF16 branch only needs the all-gather: the optimizer step already wrote + the updated shards into the DDP param buffer, and the MXFP8-only staging + helper must not be involved. + """ + from nemo_rl.models.policy.workers import megatron_policy_worker + + events = [] + + class FakeDDP: + ddp_config = SimpleNamespace(overlap_param_gather=True) + + def start_param_sync(self, *, force_sync): + events.append(("start_param_sync", force_sync)) + + monkeypatch.setattr(megatron_policy_worker, "DistributedDataParallel", FakeDDP) + monkeypatch.setattr( + torch.cuda, "synchronize", lambda: events.append(("cuda_synchronize", None)) + ) + + worker = object.__new__(megatron_policy_worker.MegatronPolicyWorkerImpl) + worker.model = FakeDDP() + worker._uses_mxfp8_overlap_shared_param_buffer = lambda: False + worker._forward_pre_hook_enabled = lambda: hooks_enabled + worker.finalize_async_save = lambda: events.append(("finalize_async_save", None)) + worker._copy_main_params_to_param_buffer = MagicMock() + + worker.sync_params_before_refit() + + expected = ( + [ + ("finalize_async_save", None), + ("start_param_sync", True), + ("cuda_synchronize", None), + ] + if hooks_enabled + else [] + ) + assert events == expected + worker._copy_main_params_to_param_buffer.assert_not_called() + + def test_megatron_offload_before_refit_finalizes_async_save_first(monkeypatch): """Async checkpoint tensor references must be released before GPU offload.""" from nemo_rl.models.policy.workers.megatron_policy_worker import ( @@ -1301,34 +1348,51 @@ def cuda(self): assert events.index("finalize_async_save") < events.index(("move_model", True)) -def test_megatron_sync_params_before_refit_materializes_latest_mxfp8_weights(): +@pytest.mark.parametrize("hooks_enabled", [True, False]) +def test_megatron_sync_params_before_refit_materializes_latest_mxfp8_weights( + monkeypatch, hooks_enabled +): """Refit must see optimizer updates before the next overlapped train forward.""" from nemo_rl.models.policy.workers.megatron_policy_worker import ( MegatronPolicyWorkerImpl, ) events = [] + + class FakeDDP: + ddp_config = SimpleNamespace(overlap_param_gather=True) + + from nemo_rl.models.policy.workers import megatron_policy_worker + + monkeypatch.setattr(megatron_policy_worker, "DistributedDataParallel", FakeDDP) worker = object.__new__(MegatronPolicyWorkerImpl) + worker.model = FakeDDP() worker.finalize_async_save = lambda: events.append("finalize_async_save") worker._uses_mxfp8_overlap_shared_param_buffer = lambda: True - worker._forward_pre_hook_enabled = lambda: True + worker._forward_pre_hook_enabled = lambda: hooks_enabled worker._disable_forward_pre_hook_until_next_train_step = ( lambda *, param_sync=False: events.append(("disable_hook", param_sync)) ) MegatronPolicyWorkerImpl.sync_params_before_refit(worker) - assert events == [ - "finalize_async_save", - ("disable_hook", True), - ] + expected = ["finalize_async_save", ("disable_hook", True)] if hooks_enabled else [] + assert events == expected -def test_megatron_sync_params_before_refit_is_noop_without_pending_mxfp8_gather(): +@pytest.mark.parametrize("ddp", [False, True]) +def test_megatron_sync_params_before_refit_is_noop_without_overlap(monkeypatch, ddp): from nemo_rl.models.policy.workers.megatron_policy_worker import ( MegatronPolicyWorkerImpl, ) worker = object.__new__(MegatronPolicyWorkerImpl) + from nemo_rl.models.policy.workers import megatron_policy_worker + + class FakeDDP: + ddp_config = SimpleNamespace(overlap_param_gather=False) + + monkeypatch.setattr(megatron_policy_worker, "DistributedDataParallel", FakeDDP) + worker.model = FakeDDP() if ddp else object() worker._uses_mxfp8_overlap_shared_param_buffer = lambda: False worker.finalize_async_save = MagicMock() worker._disable_forward_pre_hook_until_next_train_step = MagicMock() diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 668ec6d52de..e5999f1d462 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -157,6 +157,9 @@ def get_reference_policy_logprobs_from_meta(self, meta: KVBatchMeta) -> None: def prepare_for_training(self) -> None: pass + def sync_params_before_refit(self) -> None: + pass + def begin_train_step(self, loss_fn: Any) -> None: pass diff --git a/tests/unit/single_controller/test_refit_recovery.py b/tests/unit/single_controller/test_refit_recovery.py index b910e4a1f1b..e1133933f24 100644 --- a/tests/unit/single_controller/test_refit_recovery.py +++ b/tests/unit/single_controller/test_refit_recovery.py @@ -48,6 +48,7 @@ GenerationFleetHealth, ShardState, ) +from nemo_rl.utils.timer import Timer async def _completed(value=None): @@ -163,6 +164,8 @@ def _make_controller( ) ctrl._inflight_by_group_id = {} ctrl._rollout_recovery_enabled = False + ctrl._trainer = SimpleNamespace(sync_params_before_refit=MagicMock()) + ctrl._timer = Timer() return ctrl, monitor, sync @@ -198,6 +201,7 @@ def test_the_retry_runs_against_the_smaller_fleet(self): asyncio.run(ctrl._sync_weights()) assert sync.sync_calls == 2 assert sync.absent_at_retry == [0] + ctrl._trainer.sync_params_before_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 470ae99a3fb..046062bd448 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -436,6 +436,8 @@ def test_sync_weights_honors_recompute_kv_cache_config( # monitor there is nothing to reconcile. ctrl._gen_fleet = None ctrl._weight_synchronizer = SimpleNamespace(sync_weights=MagicMock()) + ctrl._trainer = SimpleNamespace(sync_params_before_refit=MagicMock()) + ctrl._timer = Timer() ctrl._rollout_manager = SimpleNamespace(resume_request_deadlines=MagicMock()) ctrl._gen = SimpleNamespace( invalidate_kv_cache=MagicMock(), @@ -452,6 +454,7 @@ def test_sync_weights_honors_recompute_kv_cache_config( asyncio.run(ctrl._sync_weights()) ctrl._weight_synchronizer.sync_weights.assert_called_once_with(kv_scales=None) + ctrl._trainer.sync_params_before_refit.assert_called_once_with() assert ctrl._gen.invalidate_kv_cache.call_count == expected_invalidation_calls assert ctrl._rollout_permitted.is_set() @@ -472,8 +475,10 @@ def test_sync_weights_calibrates_and_forwards_fp8_kv_scales() -> None: requires_kv_scale_sync=True, ) ctrl._trainer = SimpleNamespace( - calibrate_qkv_fp8_scales=MagicMock(return_value={"layers": {"layer.0": 0.5}}) + calibrate_qkv_fp8_scales=MagicMock(return_value={"layers": {"layer.0": 0.5}}), + sync_params_before_refit=MagicMock(), ) + ctrl._timer = Timer() ctrl._inflight_by_group_id = {} ctrl._rollout_recovery_enabled = False # env={} -> should_use_nemo_gym is False, so _sync_weights takes the native @@ -497,6 +502,7 @@ def test_sync_weights_calibrates_and_forwards_fp8_kv_scales() -> None: ctrl._weight_synchronizer.sync_weights.assert_called_once_with( kv_scales={"layer.0": 0.5} ) + ctrl._trainer.sync_params_before_refit.assert_called_once_with() class _AdvantageDataPlane: diff --git a/tests/unit/weight_sync/test_weight_synchronizer.py b/tests/unit/weight_sync/test_weight_synchronizer.py index f80820bdc78..5da08d266c9 100644 --- a/tests/unit/weight_sync/test_weight_synchronizer.py +++ b/tests/unit/weight_sync/test_weight_synchronizer.py @@ -14,7 +14,7 @@ """Unit tests for the WeightSynchronizer abstraction and its implementations.""" -from unittest.mock import MagicMock, call, patch +from unittest.mock import MagicMock, patch import pytest @@ -131,7 +131,7 @@ def test_sync_weights_calls_full_lifecycle(self, mock_ray): sync.sync_weights() assert not sync.is_stale - policy.sync_params_before_refit.assert_called_once_with() + policy.sync_params_before_refit.assert_not_called() policy.offload_before_refit.assert_called_once() gen.prepare_for_generation.assert_any_call(tags=["weights"]) policy.stream_weights_via_ipc_zmq.assert_called_once() @@ -281,7 +281,7 @@ def test_sync_weights_calls_full_lifecycle(self, mock_ray): sync.sync_weights() assert not sync.is_stale - policy.sync_params_before_refit.assert_called_once_with() + policy.sync_params_before_refit.assert_not_called() policy.offload_before_refit.assert_called_once() gen.prepare_for_generation.assert_any_call(tags=["weights"]) gen.pause_generation.assert_called_once_with(mode="retract") @@ -536,7 +536,7 @@ def test_sync_weights_calls_broadcast_and_receive(self, mock_ray): sync.sync_weights() assert not sync.is_stale - policy.sync_params_before_refit.assert_called_once_with() + policy.sync_params_before_refit.assert_not_called() policy.broadcast_weights_for_collective.assert_called_once_with( kv_scales=None, refit_timeout_s=None, @@ -544,16 +544,6 @@ def test_sync_weights_calls_broadcast_and_receive(self, mock_ray): num_buffers=None, ) gen.update_weights_from_collective.assert_called_once() - assert policy.mock_calls.index(call.sync_params_before_refit()) < ( - policy.mock_calls.index( - call.broadcast_weights_for_collective( - kv_scales=None, - refit_timeout_s=None, - buffer_size_bytes=None, - num_buffers=None, - ) - ) - ) @patch("nemo_rl.weight_sync.collective_weight_synchronizer.ray") def test_sync_weights_passes_kv_scales(self, mock_ray): @@ -650,7 +640,7 @@ def test_backend_sender_contract_controls_geometry_and_world_size(self, mock_ray class TestNcclReshardWeightSynchronizer: @patch("nemo_rl.weight_sync.nccl_reshard_weight_synchronizer.ray") - def test_sync_weights_materializes_policy_params_before_transfer(self, mock_ray): + def test_sync_weights_leaves_policy_param_sync_to_caller(self, mock_ray): mock_ray.get.return_value = [True] policy = _mock_policy() policy.nccl_reshard_refit.return_value = [MagicMock()] @@ -662,11 +652,9 @@ def test_sync_weights_materializes_policy_params_before_transfer(self, mock_ray) sync.sync_weights() - policy.sync_params_before_refit.assert_called_once_with() - assert policy.mock_calls.index(call.sync_params_before_refit()) < ( - policy.mock_calls.index( - call.nccl_reshard_refit(kv_scales=None, refit_timeout_s=None) - ) + policy.sync_params_before_refit.assert_not_called() + policy.nccl_reshard_refit.assert_called_once_with( + kv_scales=None, refit_timeout_s=None ) @patch("nemo_rl.weight_sync.nccl_reshard_weight_synchronizer.ray") @@ -812,7 +800,6 @@ def test_m2n_refit_delegates_transfer_and_keeps_megatron_lifecycle( train_cluster=sync._train_cluster, inference_cluster=sync._inference_cluster, refit_timeout_s=17.0, - sync_policy_params=False, ) sync.init_communicator() @@ -821,7 +808,7 @@ def test_m2n_refit_delegates_transfer_and_keeps_megatron_lifecycle( transport.init_communicator.assert_called_once() transport.sync_weights.assert_called_once_with(kv_scales={"scale": 1.0}) gen.suspend_for_refit.assert_called_once() - policy.sync_params_before_refit.assert_called_once_with() + policy.sync_params_before_refit.assert_not_called() policy.offload_before_refit.assert_not_called() assert [ call.kwargs.get("tags") @@ -860,7 +847,7 @@ def test_non_colocated_sync_sequence(self, mock_ray, refit_backend): assert sync.sync_weights() == {} gen.suspend_for_refit.assert_called_once() - policy.sync_params_before_refit.assert_called_once_with() + policy.sync_params_before_refit.assert_not_called() policy.offload_before_refit.assert_not_called() policy.swap_weights_via_reshard.assert_called_once_with(is_source=True) gen.update_weights_from_collective.assert_called_once_with(refit_timeout_s=None) @@ -899,7 +886,7 @@ def test_colocated_sync_is_offload_and_wake(self): assert sync.is_stale assert sync.sync_weights() == {} - policy.sync_params_before_refit.assert_called_once_with() + policy.sync_params_before_refit.assert_not_called() policy.offload_before_refit.assert_called_once() # The refit-protocol tag makes the wake bypass the worker's # engine-awake early-return (the reshard copy rides this wake). @@ -929,12 +916,9 @@ def test_non_colocated_policy_offload_is_configurable( sync.init_communicator() sync.sync_weights() - policy.sync_params_before_refit.assert_called_once_with() + policy.sync_params_before_refit.assert_not_called() if offload_policy_before_refit: policy.offload_before_refit.assert_called_once_with() - assert policy.mock_calls.index(call.sync_params_before_refit()) < ( - policy.mock_calls.index(call.offload_before_refit()) - ) else: policy.offload_before_refit.assert_not_called() @@ -996,7 +980,6 @@ def test_non_colocated_null_transport_uses_collective_transport(self): inference_cluster=_mock_cluster(), ) assert isinstance(sync._transport, CollectiveWeightSynchronizer) - assert sync._transport._sync_policy_params is False class TestFactory: @@ -1034,7 +1017,6 @@ def test_non_colocated_vllm_returns_collective(self): inference_cluster=_mock_cluster(), ) assert isinstance(sync, CollectiveWeightSynchronizer) - assert sync._sync_policy_params is True def test_colocated_megatron_returns_megatron_synchronizer(self): sync = create_weight_synchronizer(