fix(megatron): synchronize parameters before inference refit

nemo_rl/algorithms/grpo.py

@@ -2528,22 +2528,28 @@ 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(
             "SGLang refits require policy_generation.weight_synchronizer to be "
             "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"])

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:

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(

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

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

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

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,21 +152,13 @@ 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"):
                 self._generation.prepare_for_generation(tags=["colocated_refit"])
             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"):

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

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