From caafe62fa884a0da9d342ac60b8af7ee64348747 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 4 Aug 2026 12:13:26 -0500
Subject: [PATCH 01/25] feat: colocated Megatron reshard

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 nemo_rl/models/generation/megatron/config.py  |  13 +-
 .../megatron/megatron_generation.py           |   4 +-
 .../generation/megatron/megatron_worker.py    | 147 ++++++++++++++----
 nemo_rl/models/megatron/memory_saver.py       |  57 +++++++
 nemo_rl/models/megatron/setup.py              |  80 ++++++++++
 .../policy/workers/megatron_policy_worker.py  |  61 ++++++++
 pyrefly.toml                                  |   1 +
 7 files changed, 333 insertions(+), 30 deletions(-)
 create mode 100644 nemo_rl/models/megatron/memory_saver.py

diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py
index af431e602c2..c13b4058132 100644
--- a/nemo_rl/models/generation/megatron/config.py
+++ b/nemo_rl/models/generation/megatron/config.py
@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Literal, NotRequired, TypedDict
+from typing import Any, Literal, NotRequired, TypedDict, cast
 
 from nemo_rl.models.generation.interfaces import GenerationConfig
+from nemo_rl.models.policy import PolicyConfig
 
 
 class MCoreGenerationSpecificArgs(TypedDict):
@@ -69,3 +70,13 @@ class MCoreGenerationConfig(GenerationConfig):
     """Generation config for Megatron Inference."""
 
     mcore_generation_config: MCoreGenerationSpecificArgs
+
+
+def apply_megatron_inference_overrides(policy_config: PolicyConfig) -> None:
+    """Apply inference-only Megatron configs on top of the normal `megatron_cfg`."""
+    generation_config = cast(MCoreGenerationConfig, policy_config["generation"])
+    # The overlay is intentionally dynamic: any inference-only key overwrites
+    # its training-side counterpart, so treat the config as a plain dict.
+    megatron_cfg = cast(dict[str, Any], policy_config["megatron_cfg"])
+    megatron_cfg.update(generation_config["mcore_generation_config"])
+    megatron_cfg["activation_checkpointing"] = False
diff --git a/nemo_rl/models/generation/megatron/megatron_generation.py b/nemo_rl/models/generation/megatron/megatron_generation.py
index 00bb10464ac..c41ae8e00da 100644
--- a/nemo_rl/models/generation/megatron/megatron_generation.py
+++ b/nemo_rl/models/generation/megatron/megatron_generation.py
@@ -25,7 +25,9 @@
     GenerationInterface,
     GenerationOutputSpec,
 )
-from nemo_rl.models.generation.megatron.config import MCoreGenerationConfig
+from nemo_rl.models.generation.megatron.config import (
+    MCoreGenerationConfig,
+)
 from nemo_rl.models.policy import PolicyConfig
 
 if TYPE_CHECKING:
diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py
index 5b6857a0704..13cb65a88f7 100644
--- a/nemo_rl/models/generation/megatron/megatron_worker.py
+++ b/nemo_rl/models/generation/megatron/megatron_worker.py
@@ -29,6 +29,12 @@
 )
 from megatron.core.inference.engines.dynamic_engine import EngineState
 from megatron.core.inference.sampling_params import SamplingParams
+from megatron.core.resharding.copy_services.gloo_copy_service import GlooCopyService
+from megatron.core.resharding.copy_services.nccl_copy_service import NCCLCopyService
+from megatron.core.resharding.refit import (
+    prepare_swap_model_weights,
+    swap_model_weights,
+)
 from megatron.core.transformer.enums import InferenceCudaGraphScope
 from megatron.core.transformer.utils import toggle_cuda_graphs
 from megatron.core.utils import unwrap_model
@@ -43,6 +49,11 @@
     log_gpu_memory,
     resolve_torch_dtype,
 )
+from nemo_rl.models.megatron.memory_saver import (
+    HAVE_TORCH_MEMORY_SAVER,
+    pause_inference_weights,
+    resume_inference_weights,
+)
 from nemo_rl.utils.nsys import wrap_with_nvtx_name
 
 
@@ -59,6 +70,18 @@ class MegatronGenerationMixin:
      - is_generation_colocated: Whether colocated or distributed.
     """
 
+    # Colocated-reshard hosts assign the dedicated inference-layout model here
+    # (see MegatronPolicyWorkerImpl._maybe_build_colocated_inference_model).
+    inference_model = None
+
+    def _gen_model(self):
+        """The model the inference engine wraps.
+
+        Returns the dedicated inference-layout model when one exists (colocated
+        reshard), otherwise the shared training model.
+        """
+        return self.inference_model if self.inference_model is not None else self.model
+
     def _init_inference_engine_state(self) -> None:
         """Reset all inference-engine attributes to their uninitialized state."""
         self.dynamic_inference_engine = None
@@ -94,7 +117,8 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None:
         )
         from megatron.core.utils import get_attr_wrapped_model
 
-        pg_collection = get_attr_wrapped_model(self.model, "pg_collection")
+        gen_model = self._gen_model()
+        pg_collection = get_attr_wrapped_model(gen_model, "pg_collection")
 
         buffer_size_gb = mcore_generation_config["buffer_size_gb"]
         num_cuda_graphs = mcore_generation_config["num_cuda_graphs"]
@@ -115,7 +139,7 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None:
         num_speculative_tokens = mcore_generation_config["num_speculative_tokens"]
         max_requests = mcore_generation_config.get("max_requests")
 
-        mamba_inference_state_config = MambaInferenceStateConfig.from_model(self.model)
+        mamba_inference_state_config = MambaInferenceStateConfig.from_model(gen_model)
         is_hybrid_model = mamba_inference_state_config is not None
         if is_hybrid_model:
             if (
@@ -140,7 +164,7 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None:
             logging_step_interval = 0
 
         # flashinfer's fused-RoPE kernel only dispatches fp16/bf16 q/k.
-        use_flashinfer_fused_rope = self.model.config.params_dtype in (
+        use_flashinfer_fused_rope = gen_model.config.params_dtype in (
             torch.float16,
             torch.bfloat16,
         )
@@ -176,15 +200,15 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None:
         )
 
         if "inference_cuda_graph_scope" in mcore_generation_config:
-            self.model.config.inference_cuda_graph_scope = InferenceCudaGraphScope[
+            gen_model.config.inference_cuda_graph_scope = InferenceCudaGraphScope[
                 mcore_generation_config["inference_cuda_graph_scope"]
             ]
 
         self.inference_context = DynamicInferenceContext(
-            self.model.config, inference_config
+            gen_model.config, inference_config
         )
         self.inference_wrapped_model = GPTInferenceWrapper(
-            self.model, self.inference_context
+            gen_model, self.inference_context
         )
         text_generation_controller = TextGenerationController(
             inference_wrapped_model=self.inference_wrapped_model,
@@ -347,7 +371,7 @@ def finish_generation(self) -> None:
         print(f"[Rank {self.rank}] finishing generation", flush=True)
         log_gpu_memory("finish_generation START")
 
-        lang_module = unwrap_model(self.model)
+        lang_module = unwrap_model(self._gen_model())
 
         if self.is_generation_colocated:
             if self._inference_engine_initialized and not self._inference_engine_asleep:
@@ -365,6 +389,9 @@ def finish_generation(self) -> None:
             rotary_module.forward.cache_clear()
 
         if self.is_generation_colocated:
+            # Offload the inference weights to CPU.
+            if self.inference_model is not None:
+                self._offload_inference_model()
             gc.collect()
             torch.cuda.empty_cache()
 
@@ -380,20 +407,33 @@ def prepare_for_generation(self, tags=None, **kwargs) -> None:
         log_gpu_memory("prepare_for_generation START")
         mcore_generation_config = self.cfg["generation"]["mcore_generation_config"]
 
-        self.model.config.flash_decode = False
-        if self.is_generation_colocated and self.should_disable_forward_pre_hook:
-            # Bring offloaded params back to CUDA before colocated generation.
+        # Colocated reshard: build the dedicated inference-layout model on the first cycle.
+        if (
+            getattr(self, "_colocated_reshard_eligible", False)
+            and not self._colocated_inference_model_checked
+        ):
+            self._maybe_build_colocated_inference_model(self.cfg)
+            self._colocated_inference_model_checked = True
+
+        gen_model = self._gen_model()
+        gen_model.config.flash_decode = False
+        if self.is_generation_colocated and self.inference_model is None:
             self.model = self.move_model(
                 self.model, "cuda", move_params=True, move_grads=False
             )
-            # DP inference schedules requests independently, so a forward pre-hook
-            # cannot safely launch a parameter all-gather from only the rank that
-            # received work. Gather once across every worker, then keep the hooks
-            # disabled until the next training step completes.
-            if self._forward_pre_hook_enabled():
+            # Forward pre-hooks are not safe with cross-DP inference.
+            if (
+                self.should_disable_forward_pre_hook
+                and self._forward_pre_hook_enabled()
+            ):
                 self._disable_forward_pre_hook_until_next_train_step(param_sync=True)
+            gen_model = self.model
 
-        lang_module = unwrap_model(self.model)
+        # Colocated reshard (hosts without a dedicated inference model skip it).
+        if self.inference_model is not None:
+            self._reshard_into_inference_model()
+
+        lang_module = unwrap_model(gen_model)
         lang_module.eval()
 
         rotary_module = getattr(lang_module, "rotary_pos_emb", None)
@@ -764,26 +804,17 @@ def init_collective_mcore_generation(
         _world.pg_names[pg] = group_name
 
         if refit_backend == "nvshmem":
+            # Deferred: importing NVSHMEMCopyService loads the optional nvshmem bindings.
             from megatron.core.resharding.copy_services.nvshmem_copy_service import (
                 NVSHMEMCopyService,
             )
 
             self.refit_copy_service = NVSHMEMCopyService(group=self.refit_pg)
         elif refit_backend == "nccl":
-            from megatron.core.resharding.copy_services.nccl_copy_service import (
-                NCCLCopyService,
-            )
-
             self.refit_copy_service = NCCLCopyService(group=self.refit_pg)
         else:
-            from megatron.core.resharding.copy_services.gloo_copy_service import (
-                GlooCopyService,
-            )
-
             self.refit_copy_service = GlooCopyService(group=self.refit_pg)
 
-        from megatron.core.resharding.refit import prepare_swap_model_weights
-
         is_source = rank_offset == 0
         # Cache for later refit calls (swap_weights_via_reshard).
         self.refit_dst_rank_offset = (
@@ -822,8 +853,6 @@ def swap_weights_via_reshard(self, is_source: bool) -> bool:
         Returns:
             True on success.
         """
-        from megatron.core.resharding.refit import swap_model_weights
-
         src_model = self.model if is_source else None
         dst_model = None if is_source else self.model
 
@@ -838,6 +867,68 @@ def swap_weights_via_reshard(self, is_source: bool) -> bool:
 
         return True
 
+    def _onload_inference_model(self) -> None:
+        """Restore the colocated inference weights to GPU before resharding / generation."""
+        if not getattr(self, "_inference_model_offloaded", False):
+            return
+        resume_inference_weights()
+        self._inference_model_offloaded = False
+
+    def _offload_inference_model(self) -> None:
+        """Offload the colocated inference weights to CPU while training runs."""
+        if (
+            getattr(self, "inference_model", None) is None
+            or getattr(self, "_inference_model_offloaded", False)
+            or not HAVE_TORCH_MEMORY_SAVER
+        ):
+            return
+        pause_inference_weights()
+        self._inference_model_offloaded = True
+
+    def _reshard_into_inference_model(self) -> None:
+        """Reshard current training weights into the colocated inference-layout model."""
+        inference_model = getattr(self, "inference_model", None)
+        if inference_model is None:
+            return
+
+        # Bring the inference weights back to GPU.
+        self._onload_inference_model()
+        self.model = self.move_model(
+            self.model, "cuda", move_params=True, move_grads=False
+        )
+
+        # The swap reads the training params as its source;
+        # under overlap_param_gather they stay stale after the optimizer step until gathered,
+        if self.should_disable_forward_pre_hook and self._forward_pre_hook_enabled():
+            self._disable_forward_pre_hook_until_next_train_step(param_sync=True)
+
+        # Build + cache the same-rank reshard plan once, before the first CUDA-graph capture.
+        if not self._colocated_reshard_plan_ready:
+            prepare_swap_model_weights(
+                src_model=self.model,
+                target_model=inference_model,
+                group=None,
+                src_rank_offset=0,
+                dst_rank_offset=0,
+            )
+            self._colocated_reshard_plan_ready = True
+
+        swap_model_weights(
+            self.model,
+            inference_model,
+            refit_method=self.cfg["generation"]["mcore_generation_config"][
+                "refit_backend"
+            ],
+            group=None,
+            src_rank_offset=0,
+            dst_rank_offset=0,
+        )
+
+        # Offload training model.
+        self.model = self.move_model(
+            self.model, "cpu", move_params=True, move_grads=False
+        )
+
     def suspend_for_refit(self) -> None:
         """Pause+suspend the inference engine before a weight refit."""
         if not self._inference_engine_initialized:
diff --git a/nemo_rl/models/megatron/memory_saver.py b/nemo_rl/models/megatron/memory_saver.py
new file mode 100644
index 00000000000..9eb07e9110b
--- /dev/null
+++ b/nemo_rl/models/megatron/memory_saver.py
@@ -0,0 +1,57 @@
+# 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 warnings
+from contextlib import nullcontext
+from typing import ContextManager
+
+try:
+    import torch_memory_saver  # pyrefly: ignore[import-error]
+
+    HAVE_TORCH_MEMORY_SAVER = True
+except ImportError:
+    HAVE_TORCH_MEMORY_SAVER = False
+
+# torch_memory_saver region tag for the colocated inference model's weights.
+_INFERENCE_MODEL_OFFLOAD_TAG = "nemo_rl_megatron_inference_model"
+
+
+def inference_model_alloc_region() -> ContextManager[None]:
+    """Allocation region to build the colocated inference model under.
+
+    Returns a CPU-backup-enabled torch_memory_saver region, or a null context.
+    """
+    if HAVE_TORCH_MEMORY_SAVER:
+        return torch_memory_saver.region(
+            tag=_INFERENCE_MODEL_OFFLOAD_TAG, enable_cpu_backup=True
+        )
+    warnings.warn(
+        "torch_memory_saver is unavailable; the colocated inference model will stay "
+        "GPU-resident alongside the training model (higher peak memory). Install "
+        "torch_memory_saver to enable inference-weight offload.",
+        stacklevel=2,
+    )
+    return nullcontext()
+
+
+def pause_inference_weights() -> None:
+    """Back the colocated inference weights to CPU (no-op without torch_memory_saver)."""
+    if HAVE_TORCH_MEMORY_SAVER:
+        torch_memory_saver.pause(_INFERENCE_MODEL_OFFLOAD_TAG)
+
+
+def resume_inference_weights() -> None:
+    """Restore the inference weights to their GPU addresses (no-op without torch_memory_saver)."""
+    if HAVE_TORCH_MEMORY_SAVER:
+        torch_memory_saver.resume(_INFERENCE_MODEL_OFFLOAD_TAG)
diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py
index bdf4d7849ac..0161eed07d5 100644
--- a/nemo_rl/models/megatron/setup.py
+++ b/nemo_rl/models/megatron/setup.py
@@ -61,6 +61,7 @@
 from megatron.bridge.utils.cuda_graph import set_cuda_graph_modules
 from megatron.bridge.utils.vocab_utils import calculate_padded_vocab_size
 from megatron.core import parallel_state
+from megatron.core.inference.shards import build_inference_pg_collection
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.transformer import MegatronModule
 from megatron.core.transformer.enums import AttnBackend, InferenceCudaGraphScope
@@ -234,6 +235,7 @@ def _sync_distrib_opt(distrib_opt):
     find_draft_owner_chunk,
     get_attached_draft_model,
 )
+from nemo_rl.models.megatron.memory_saver import inference_model_alloc_region
 from nemo_rl.models.megatron.router_replay import (
     clear_global_router_replay_instances,
     router_replay_enabled,
@@ -1477,6 +1479,84 @@ def main_thread_only_enter(self):
     _BRIDGE_SIGNAL_HANDLER_PATCHED = True
 
 
+def build_inference_model(
+    policy_cfg: PolicyConfig,
+    megatron_cfg: ConfigContainer,
+) -> MegatronModule:
+    """Build a second, inference-layout model for colocated Megatron refit.
+
+    The returned model is resident on GPU; its weights are uninitialized until the first reshard.
+
+    Args:
+        policy_cfg: The inference config
+        megatron_cfg: The training config
+
+    Returns:
+        The inference model module (single element; not DDP-wrapped, no optimizer).
+    """
+    # Deep-copy the training config and apply the inference overrides on top.
+    inference_provider = copy.deepcopy(megatron_cfg.model)
+    train_pipeline_model_parallel_size = inference_provider.pipeline_model_parallel_size
+    _apply_parallelism_config(inference_provider, policy_cfg)
+    _apply_moe_config(inference_provider, policy_cfg)
+    if "transformer_impl" in policy_cfg["megatron_cfg"]:
+        inference_provider.transformer_impl = policy_cfg["megatron_cfg"][
+            "transformer_impl"
+        ]
+    # A custom (uneven) pipeline split is tuned for the training PP; reset to an even split
+    # when inference uses a different PP (the reshard maps params across stages by name).
+    if (
+        inference_provider.pipeline_model_parallel_size
+        != train_pipeline_model_parallel_size
+    ):
+        inference_provider.num_layers_in_first_pipeline_stage = None
+        inference_provider.num_layers_in_last_pipeline_stage = None
+    # Sequence parallelism requires TP > 1; force it off otherwise (Megatron asserts this).
+    inference_provider.sequence_parallel = (
+        inference_provider.sequence_parallel
+        and inference_provider.tensor_model_parallel_size > 1
+    )
+    # Inference never trains: disable recompute.
+    inference_provider.recompute_granularity = None
+    inference_provider.recompute_method = None
+    inference_provider.recompute_num_layers = None
+
+    world_size = torch.distributed.get_world_size()
+    inference_pg_collection = build_inference_pg_collection(
+        world_size,
+        tp_size=inference_provider.tensor_model_parallel_size,
+        pp_size=inference_provider.pipeline_model_parallel_size,
+        cp_size=inference_provider.context_parallel_size,
+        ep_size=inference_provider.expert_model_parallel_size,
+        expt_tp_size=inference_provider.expert_tensor_parallel_size,
+        use_tp_pp_dp_mapping=megatron_cfg.dist.use_tp_pp_dp_mapping,
+        rank_offset=0,  # colocated: the same ranks hold both the training and inference models
+    )
+    setattr(inference_provider, "_pg_collection", inference_pg_collection)
+
+    # Match the training mixed-precision wrapper.
+    mixed_precision_wrapper = (
+        MoEFloat16Module
+        if policy_cfg["megatron_cfg"]["freeze_moe_router"]
+        else Float16Module
+    )
+
+    # Only one model's weights stay resident at a time; swap weights in and out at the same address.
+    with inference_model_alloc_region():
+        inference_model = get_model(
+            inference_provider,
+            megatron_cfg.ddp,
+            use_torch_fsdp2=megatron_cfg.dist.use_torch_fsdp2,
+            data_parallel_random_init=megatron_cfg.rng.data_parallel_random_init,
+            mixed_precision_wrapper=mixed_precision_wrapper,
+            pg_collection=inference_pg_collection,
+            wrap_with_ddp=False,  # never trained: no DDP, no grad buffers, no optimizer
+        )
+    inference_model = inference_model[0]
+    inference_model.eval()
+    return inference_model
+
+
 def setup_model_and_optimizer(
     policy_cfg: PolicyConfig,
     megatron_cfg: ConfigContainer,
diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py
index 56afb8aefd3..a6f06e0c4b6 100644
--- a/nemo_rl/models/policy/workers/megatron_policy_worker.py
+++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py
@@ -54,6 +54,9 @@
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.distributed.named_sharding import NamedSharding
 from nemo_rl.models.generation.interfaces import GenerationDatumSpec
+from nemo_rl.models.generation.megatron.config import (
+    apply_megatron_inference_overrides,
+)
 from nemo_rl.models.generation.megatron.megatron_worker import (
     MegatronGenerationMixin,
     MegatronGenerationRefitMixin,
@@ -71,6 +74,7 @@
 )
 from nemo_rl.models.megatron.router_replay import router_replay_enabled
 from nemo_rl.models.megatron.setup import (
+    build_inference_model,
     finalize_megatron_setup,
     handle_model_import,
     setup_distributed,
@@ -617,6 +621,20 @@ def __init__(
                     "virtual pipeline parallelism."
                 )
 
+        # Colocated reshard: build a dedicated inference-layout model container.
+        self.inference_model = None
+        self._colocated_reshard_plan_ready = False
+        self._inference_model_offloaded = False
+        self._colocated_inference_model_checked = False
+        gen_cfg = config.get("generation")
+        # The build itself is deferred to the first prepare_for_generation.
+        self._colocated_reshard_eligible = (
+            init_optimizer
+            and self.is_generation_colocated
+            and gen_cfg is not None
+            and gen_cfg.get("backend") == "megatron"
+        )
+
         # vars used for refit
         ## will be initialized in prepare_refit_info
         # refit_param_info_mcore combines the conversion tasks with the param memory
@@ -2737,6 +2755,49 @@ def prepare_for_lp_inference(self):
         gc.collect()
         torch.cuda.empty_cache()
 
+    def _maybe_build_colocated_inference_model(self, config) -> None:
+        """Build a separate inference-layout model when the colocated layout differs."""
+        # Resolve the inference layout the same way the non-colocated generation policy does:
+        # overlay the sparse mcore_generation_config onto a copy of megatron_cfg.
+        inference_config = copy.deepcopy(config)
+        apply_megatron_inference_overrides(inference_config)
+        # Inference never uses CP: pin CP=1, so CP>1 training builds a separate inference model.
+        inference_config["megatron_cfg"]["context_parallel_size"] = 1
+
+        train_mcfg = config["megatron_cfg"]
+        inf_mcfg = inference_config["megatron_cfg"]
+        layout_keys = (
+            "tensor_model_parallel_size",
+            "pipeline_model_parallel_size",
+            "expert_model_parallel_size",
+            "expert_tensor_parallel_size",
+            "context_parallel_size",
+        )
+        layout_differs = any(inf_mcfg[k] != train_mcfg[k] for k in layout_keys)
+        impl_differs = inf_mcfg.get("transformer_impl") != train_mcfg.get(
+            "transformer_impl"
+        )
+        if not (layout_differs or impl_differs):
+            return
+
+        peft_cfg = train_mcfg.get("peft")
+        if peft_cfg is not None and peft_cfg.get("enabled"):
+            raise NotImplementedError(
+                "Colocated generation with a differing inference parallel layout is not "
+                "supported with PEFT. Use a matched layout or non-colocated generation."
+            )
+        draft_cfg = config.get("draft")
+        if draft_cfg is not None and draft_cfg.get("enabled"):
+            raise NotImplementedError(
+                "Colocated generation with a differing inference parallel layout is not "
+                "supported with a speculative draft model."
+            )
+        # Built inside the first prepare_for_generation, immediately before the
+        # reshard needs it resident; finish_generation offloads it afterwards.
+        self.inference_model = build_inference_model(
+            inference_config, self.megatron_cfg
+        )
+
     def prepare_for_training(self, *args, **kwargs):
         # onload models and optimizer state to cuda
         self.model = self.move_model(
diff --git a/pyrefly.toml b/pyrefly.toml
index bd92f2c37df..004ee094136 100644
--- a/pyrefly.toml
+++ b/pyrefly.toml
@@ -197,6 +197,7 @@ project-includes = [
   "nemo_rl/models/huggingface/__init__.py",
   "nemo_rl/models/megatron/__init__.py",
   "nemo_rl/models/megatron/draft/__init__.py",
+  "nemo_rl/models/megatron/memory_saver.py",
   "nemo_rl/models/policy/__init__.py",
   "nemo_rl/models/policy/interfaces.py",
   "nemo_rl/models/policy/utils.py",

From e78254ca07cbc16a48bfc62d591a9735887d38c5 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 4 Aug 2026 12:13:27 -0500
Subject: [PATCH 02/25] colocated-reshard nightly (nanov3)

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 ...0BA3B-2n8g-megatron_colocated_reshard.yaml | 47 +++++++++++++++++++
 ...-30BA3B-2n8g-megatron_colocated_reshard.sh | 38 +++++++++++++++
 tests/test_suites/nightly.txt                 |  1 +
 3 files changed, 86 insertions(+)
 create mode 100644 examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.yaml
 create mode 100755 tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.sh

diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.yaml
new file mode 100644
index 00000000000..12f9d5cbacd
--- /dev/null
+++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.yaml
@@ -0,0 +1,47 @@
+defaults: ../../grpo_math_1B.yaml
+# Sync (non-async) colocated GRPO with Megatron generation (Nemotron-3-Nano-30B-A3B), with reshard allowed.
+grpo:
+  num_prompts_per_step: 2
+  num_generations_per_prompt: 8
+  max_num_steps: 500
+checkpointing:
+  enabled: false
+  checkpoint_dir: results/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard
+  save_period: 100
+policy:
+  model_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16
+  tokenizer:
+    name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
+  train_global_batch_size: 16
+  train_micro_batch_size: 1
+  logprob_batch_size: 1
+  max_total_sequence_length: 2048
+  megatron_cfg:
+    enabled: true
+    bias_activation_fusion: false
+    tensor_model_parallel_size: 2
+    expert_model_parallel_size: 8
+    sequence_parallel: true
+  dtensor_cfg:
+    enabled: false
+  sequence_packing:
+    enabled: false
+  generation:
+    backend: megatron
+    mcore_generation_config:
+      transformer_impl: inference_optimized
+      activation_checkpointing: false
+      tensor_model_parallel_size: 4
+      expert_model_parallel_size: 4
+      sequence_parallel: true
+      refit_backend: nccl
+logger:
+  log_dir: logs/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard
+  wandb_enabled: true
+  tensorboard_enabled: true
+  wandb:
+    project: nemo-rl
+    name: grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard
+cluster:
+  gpus_per_node: 8
+  num_nodes: 2
diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.sh
new file mode 100755
index 00000000000..43b19a166e4
--- /dev/null
+++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.sh
@@ -0,0 +1,38 @@
+#!/bin/bash
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+source $SCRIPT_DIR/common.env
+
+# ===== BEGIN CONFIG =====
+NUM_NODES=2
+GPUS_PER_NODE=8
+STEPS_PER_RUN=10
+MAX_STEPS=10
+NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN ))  # Round up
+NUM_MINUTES=60
+# ===== END CONFIG =====
+
+exit_if_max_steps_reached
+
+# Run the experiment
+cd $PROJECT_ROOT
+uv run examples/run_grpo.py \
+    --config $CONFIG_PATH \
+    grpo.max_num_steps=$MAX_STEPS \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=True \
+    logger.wandb.project=nemo-rl \
+    logger.wandb.name=$EXP_NAME \
+    logger.monitor_gpus=True \
+    logger.tensorboard_enabled=True \
+    checkpointing.enabled=False \
+    $@ \
+    2>&1 | tee $RUN_LOG
+
+# Convert tensorboard logs to json
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+# Only run metrics if the target step is reached
+if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then
+    uv run tests/check_metrics.py $JSON_METRICS \
+        'max(data["train/reward"]) > 0.0'
+fi
diff --git a/tests/test_suites/nightly.txt b/tests/test_suites/nightly.txt
index 64fb93790e4..3ce95e53c82 100644
--- a/tests/test_suites/nightly.txt
+++ b/tests/test_suites/nightly.txt
@@ -117,6 +117,7 @@ tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-lora.sh
 tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp.sh
 tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation.sh
 tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym.sh
+tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.sh
 
 # Nano-v3.5
 tests/test_suites/llm/dapo-nanov3.5-30BA3B-4n8g-automodel.sh

From f7d1ef469fb932cbac08edaacc7a74f55911c0a8 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 4 Aug 2026 12:13:28 -0500
Subject: [PATCH 03/25] colocated reshard functional test

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 .../L1_Functional_Tests_Megatron_4.sh         |  1 +
 ...o_megatron_generation_colocated_reshard.sh | 52 +++++++++++++++++++
 2 files changed, 53 insertions(+)
 create mode 100755 tests/functional/grpo_megatron_generation_colocated_reshard.sh

diff --git a/tests/functional/L1_Functional_Tests_Megatron_4.sh b/tests/functional/L1_Functional_Tests_Megatron_4.sh
index 3ad59ff4699..2462db4d6e2 100644
--- a/tests/functional/L1_Functional_Tests_Megatron_4.sh
+++ b/tests/functional/L1_Functional_Tests_Megatron_4.sh
@@ -52,6 +52,7 @@ if megatron_generation_supported; then
     run_test fast uv run --no-sync bash ./tests/functional/grpo_megatron_generation.sh
     run_test      uv run --no-sync bash ./tests/functional/grpo_megatron_generation_topology.sh
     run_test      uv run --no-sync bash ./tests/functional/grpo_megatron_generation_non_colocated.sh
+    run_test      uv run --no-sync bash ./tests/functional/grpo_megatron_generation_colocated_reshard.sh
     run_test      uv run --no-sync bash ./tests/functional/grpo_megatron_generation_async_gym.sh
     run_test fast uv run --no-sync bash ./tests/functional/grpo_megatron_generation_topp_topk.sh
     # Disabled: token_mult_prob_error ~2.0 > 1.1 under top_p/top_k after the
diff --git a/tests/functional/grpo_megatron_generation_colocated_reshard.sh b/tests/functional/grpo_megatron_generation_colocated_reshard.sh
new file mode 100755
index 00000000000..8acf12bf47f
--- /dev/null
+++ b/tests/functional/grpo_megatron_generation_colocated_reshard.sh
@@ -0,0 +1,52 @@
+#!/bin/bash
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..)
+# Mark the current repo as safe, since wandb fetches metadata about the repo
+git config --global --add safe.directory $PROJECT_ROOT
+
+set -eou pipefail
+
+EXP_NAME=$(basename $0 .sh)
+EXP_DIR=$SCRIPT_DIR/$EXP_NAME
+LOG_DIR=$EXP_DIR/logs
+JSON_METRICS=$EXP_DIR/metrics.json
+RUN_LOG=$EXP_DIR/run.log
+export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-}
+
+rm -rf $EXP_DIR $LOG_DIR
+mkdir -p $EXP_DIR $LOG_DIR
+
+# Colocated reshard: training runs TP2 while inference runs TP1 x DP2, so
+# prepare_for_generation must build a dedicated inference model and swap
+# weights across differing layouts. A wrong-weights reshard shows up as
+# generation/training logprob disagreement, hence the mult_prob_error gate.
+cd $PROJECT_ROOT
+uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \
+    $PROJECT_ROOT/examples/run_grpo.py \
+    --config $PROJECT_ROOT/examples/configs/grpo_math_1B_megatron.yaml \
+    policy.model_name=Qwen/Qwen2.5-0.5B \
+    grpo.num_prompts_per_step=2 \
+    grpo.num_generations_per_prompt=4 \
+    policy.train_global_batch_size=4 \
+    policy.logprob_batch_size=4 \
+    policy.train_micro_batch_size=1 \
+    policy.megatron_cfg.tensor_model_parallel_size=2 \
+    policy.generation.backend=megatron \
+    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 \
+    cluster.gpus_per_node=2 \
+    grpo.max_num_steps=2 \
+    logger.tensorboard_enabled=true \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=false \
+    logger.monitor_gpus=true \
+    checkpointing.enabled=false \
+    $@ \
+    2>&1 | tee $RUN_LOG
+
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+uv run tests/check_metrics.py $JSON_METRICS \
+    'max(data["train/token_mult_prob_error"]) < 1.05'

From 84412d334d0055652275107235dd8ea3f4f20384 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Wed, 5 Aug 2026 02:09:57 -0500
Subject: [PATCH 04/25] Fix issues uncovered by tests

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 .../models/generation/megatron/megatron_worker.py  |  5 ++++-
 nemo_rl/models/megatron/setup.py                   | 14 ++++++++++++--
 .../grpo_megatron_generation_colocated_reshard.sh  |  4 ++--
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py
index 13cb65a88f7..b392f6fa80e 100644
--- a/nemo_rl/models/generation/megatron/megatron_worker.py
+++ b/nemo_rl/models/generation/megatron/megatron_worker.py
@@ -896,6 +896,8 @@ def _reshard_into_inference_model(self) -> None:
         self.model = self.move_model(
             self.model, "cuda", move_params=True, move_grads=False
         )
+        # TODO: Optimize away the full synchronization.
+        torch.cuda.synchronize()
 
         # The swap reads the training params as its source;
         # under overlap_param_gather they stay stale after the optimizer step until gathered,
@@ -923,11 +925,12 @@ def _reshard_into_inference_model(self) -> None:
             src_rank_offset=0,
             dst_rank_offset=0,
         )
-
         # Offload training model.
         self.model = self.move_model(
             self.model, "cpu", move_params=True, move_grads=False
         )
+        # TODO: Optimize away the full synchronization.
+        torch.cuda.synchronize()
 
     def suspend_for_refit(self) -> None:
         """Pause+suspend the inference engine before a weight refit."""
diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py
index 0161eed07d5..b158921402a 100644
--- a/nemo_rl/models/megatron/setup.py
+++ b/nemo_rl/models/megatron/setup.py
@@ -1494,8 +1494,9 @@ def build_inference_model(
     Returns:
         The inference model module (single element; not DDP-wrapped, no optimizer).
     """
-    # Deep-copy the training config and apply the inference overrides on top.
-    inference_provider = copy.deepcopy(megatron_cfg.model)
+    # Derive the inference provider from the initial snapshot taken by setup_model_and_optimizer.
+    inference_provider = megatron_cfg._initial_model_provider
+    del megatron_cfg._initial_model_provider
     train_pipeline_model_parallel_size = inference_provider.pipeline_model_parallel_size
     _apply_parallelism_config(inference_provider, policy_cfg)
     _apply_moe_config(inference_provider, policy_cfg)
@@ -1702,6 +1703,15 @@ def freeze_moe_router(megatron_model):
 
     megatron_cfg.peft = peft
 
+    # Snapshot the provider before any runtime state is added onto it.
+    generation_cfg = policy_cfg.get("generation")
+    if (
+        generation_cfg is not None
+        and generation_cfg.get("backend") == "megatron"
+        and generation_cfg.get("colocated", {}).get("enabled", False)
+    ):
+        megatron_cfg._initial_model_provider = copy.deepcopy(megatron_cfg.model)
+
     if megatron_cfg.peft is not None:
         pre_peft_hook = _create_peft_pre_wrap_hook(megatron_cfg, state)
         megatron_cfg.model.register_pre_wrap_hook(pre_peft_hook)
diff --git a/tests/functional/grpo_megatron_generation_colocated_reshard.sh b/tests/functional/grpo_megatron_generation_colocated_reshard.sh
index 8acf12bf47f..de87c759379 100755
--- a/tests/functional/grpo_megatron_generation_colocated_reshard.sh
+++ b/tests/functional/grpo_megatron_generation_colocated_reshard.sh
@@ -33,8 +33,8 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE
     policy.train_micro_batch_size=1 \
     policy.megatron_cfg.tensor_model_parallel_size=2 \
     policy.generation.backend=megatron \
-    policy.generation.mcore_generation_config.transformer_impl=inference_optimized \
-    policy.generation.mcore_generation_config.tensor_model_parallel_size=1 \
+    ++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 \
     cluster.gpus_per_node=2 \
     grpo.max_num_steps=2 \

From 6e650b538e9ad7c75471eb6c7854d820858a5f88 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Fri, 7 Aug 2026 03:36:29 -0500
Subject: [PATCH 05/25] Single source of truth for config merge

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 nemo_rl/models/generation/megatron/config.py   | 18 +++++++++++-------
 .../generation/megatron/megatron_generation.py | 11 +++--------
 .../policy/workers/megatron_policy_worker.py   |  4 ++--
 3 files changed, 16 insertions(+), 17 deletions(-)

diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py
index c13b4058132..e37bdb61fb2 100644
--- a/nemo_rl/models/generation/megatron/config.py
+++ b/nemo_rl/models/generation/megatron/config.py
@@ -72,11 +72,15 @@ class MCoreGenerationConfig(GenerationConfig):
     mcore_generation_config: MCoreGenerationSpecificArgs
 
 
-def apply_megatron_inference_overrides(policy_config: PolicyConfig) -> None:
-    """Apply inference-only Megatron configs on top of the normal `megatron_cfg`."""
+def merged_inference_megatron_cfg(policy_config: PolicyConfig) -> dict[str, Any]:
+    """The `megatron_cfg` a dedicated inference model runs with.
+
+    Overlays the sparse `mcore_generation_config` onto `megatron_cfg`,
+    intentionally overwriting any training-side config with inference-side config.
+    """
     generation_config = cast(MCoreGenerationConfig, policy_config["generation"])
-    # The overlay is intentionally dynamic: any inference-only key overwrites
-    # its training-side counterpart, so treat the config as a plain dict.
-    megatron_cfg = cast(dict[str, Any], policy_config["megatron_cfg"])
-    megatron_cfg.update(generation_config["mcore_generation_config"])
-    megatron_cfg["activation_checkpointing"] = False
+    return {
+        **cast(dict[str, Any], policy_config["megatron_cfg"]),
+        **(generation_config.get("mcore_generation_config") or {}),
+        "activation_checkpointing": False,
+    }
diff --git a/nemo_rl/models/generation/megatron/megatron_generation.py b/nemo_rl/models/generation/megatron/megatron_generation.py
index c41ae8e00da..69c22d09813 100644
--- a/nemo_rl/models/generation/megatron/megatron_generation.py
+++ b/nemo_rl/models/generation/megatron/megatron_generation.py
@@ -27,6 +27,7 @@
 )
 from nemo_rl.models.generation.megatron.config import (
     MCoreGenerationConfig,
+    merged_inference_megatron_cfg,
 )
 from nemo_rl.models.policy import PolicyConfig
 
@@ -45,13 +46,9 @@ def effective_megatron_cfg(config: PolicyConfig) -> dict[str, Any]:
         values apply; non-colocated builds a dedicated policy with
         mcore_generation_config merged on top. Always returns a fresh dict.
         """
-        megatron_cfg = config["megatron_cfg"]
         if config["generation"]["colocated"]["enabled"]:
-            return dict(megatron_cfg)
-        return {
-            **megatron_cfg,
-            **config["generation"].get("mcore_generation_config", {}),
-        }
+            return dict(config["megatron_cfg"])
+        return merged_inference_megatron_cfg(config)
 
     @classmethod
     def nvlink_domain_span(cls, config: PolicyConfig) -> int:
@@ -139,8 +136,6 @@ def __init__(
             **config,
             "megatron_cfg": self.effective_megatron_cfg(config),
         }
-        # Activation checkpointing is not compatible or useful in inference.
-        self._policy_config["megatron_cfg"]["activation_checkpointing"] = False
         # Reserve GPUs before Policy workers grab them, to prevent disjoint NVLS domains.
         self.init_cluster_placement_groups(cluster, self._policy_config)
         self._policy = Policy(
diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py
index a6f06e0c4b6..ecfc0e130a1 100644
--- a/nemo_rl/models/policy/workers/megatron_policy_worker.py
+++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py
@@ -55,7 +55,7 @@
 from nemo_rl.distributed.named_sharding import NamedSharding
 from nemo_rl.models.generation.interfaces import GenerationDatumSpec
 from nemo_rl.models.generation.megatron.config import (
-    apply_megatron_inference_overrides,
+    merged_inference_megatron_cfg,
 )
 from nemo_rl.models.generation.megatron.megatron_worker import (
     MegatronGenerationMixin,
@@ -2760,7 +2760,7 @@ def _maybe_build_colocated_inference_model(self, config) -> None:
         # Resolve the inference layout the same way the non-colocated generation policy does:
         # overlay the sparse mcore_generation_config onto a copy of megatron_cfg.
         inference_config = copy.deepcopy(config)
-        apply_megatron_inference_overrides(inference_config)
+        inference_config["megatron_cfg"] = merged_inference_megatron_cfg(inference_config)
         # Inference never uses CP: pin CP=1, so CP>1 training builds a separate inference model.
         inference_config["megatron_cfg"]["context_parallel_size"] = 1
 

From 4a2dfafca1f97f2dccfe8d60153369d7e037ac88 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Fri, 7 Aug 2026 12:16:20 -0500
Subject: [PATCH 06/25] lint

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 nemo_rl/models/policy/workers/megatron_policy_worker.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py
index ecfc0e130a1..271d90d9079 100644
--- a/nemo_rl/models/policy/workers/megatron_policy_worker.py
+++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py
@@ -2760,7 +2760,9 @@ def _maybe_build_colocated_inference_model(self, config) -> None:
         # Resolve the inference layout the same way the non-colocated generation policy does:
         # overlay the sparse mcore_generation_config onto a copy of megatron_cfg.
         inference_config = copy.deepcopy(config)
-        inference_config["megatron_cfg"] = merged_inference_megatron_cfg(inference_config)
+        inference_config["megatron_cfg"] = merged_inference_megatron_cfg(
+            inference_config
+        )
         # Inference never uses CP: pin CP=1, so CP>1 training builds a separate inference model.
         inference_config["megatron_cfg"]["context_parallel_size"] = 1
 

From fecd2e3843e7857cd1920124f9a14c5b48ac5762 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Fri, 7 Aug 2026 12:46:10 -0500
Subject: [PATCH 07/25] Unify Megatron refit behind a WeightSynchronizer

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 nemo_rl/algorithms/grpo.py                    |  74 +++------
 .../megatron/megatron_generation.py           |   3 +
 .../checkpoint_engine_weight_synchronizer.py  |   3 -
 .../collective_weight_synchronizer.py         |   3 -
 nemo_rl/weight_sync/factory.py                |  15 ++
 .../weight_sync/http_weight_synchronizer.py   |   3 -
 nemo_rl/weight_sync/interfaces.py             |  23 +--
 .../weight_sync/ipc_weight_synchronizer.py    |   3 -
 .../megatron_weight_synchronizer.py           | 148 ++++++++++++++++++
 .../nccl_reshard_weight_synchronizer.py       |   3 -
 .../vllm_remote_sparse_weight_synchronizer.py |   3 -
 ...t_checkpoint_engine_weight_synchronizer.py |   2 -
 ..._vllm_remote_sparse_weight_synchronizer.py |   1 -
 .../weight_sync/test_weight_synchronizer.py   | 146 ++++++++++++++---
 14 files changed, 321 insertions(+), 109 deletions(-)
 create mode 100644 nemo_rl/weight_sync/megatron_weight_synchronizer.py

diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py
index 7321169e58b..aeeaf49e09a 100644
--- a/nemo_rl/algorithms/grpo.py
+++ b/nemo_rl/algorithms/grpo.py
@@ -1452,8 +1452,27 @@ def init_trtllm():
             "https://github.com/NVIDIA-NeMo/RL/issues/3288."
         )
 
+    if backend == "megatron":
+        t0 = time.perf_counter()
+        policy_generation.weight_synchronizer = create_weight_synchronizer(
+            policy=policy,
+            generation=policy_generation,
+            generation_backend=backend,
+            colocated=colocated_inference,
+            train_cluster=train_cluster,
+            inference_cluster=None if colocated_inference else inference_cluster,
+        )
+        policy_generation.weight_synchronizer.init_communicator()
+        worker_init_timing_metrics["collective_init_time_s"] = time.perf_counter() - t0
+        if not colocated_inference:
+            # Load the model weights now.
+            t0 = time.perf_counter()
+            policy_generation.weight_synchronizer.sync_weights()
+            worker_init_timing_metrics["initial_weight_sync_time_s"] = (
+                time.perf_counter() - t0
+            )
     # if it is not colocated inference, initialize collective communication for update weights
-    if (
+    elif (
         not colocated_inference
         and remote_transport is None
         and checkpoint_engine_config is None
@@ -1467,26 +1486,7 @@ def init_trtllm():
         world_size = train_world_size + inference_world_size
 
         # init collective
-        if backend == "megatron":
-            refit_backend = policy_config["generation"]["mcore_generation_config"][
-                "refit_backend"
-            ]
-            futures_train = policy.init_collective_mcore_generation(
-                ip,
-                port,
-                world_size,
-                rank_offset=0,
-                refit_backend=refit_backend,
-            )
-            futures_inference = policy_generation.init_collective(
-                ip,
-                port,
-                world_size,
-                train_world_size=train_world_size,
-                refit_backend=refit_backend,
-            )
-            ray.get(futures_train + futures_inference)
-        elif nccl_reshard_refit_enabled:
+        if nccl_reshard_refit_enabled:
             policy_generation.weight_synchronizer = create_weight_synchronizer(
                 policy=policy,
                 generation=policy_generation,
@@ -2331,27 +2331,10 @@ def refit_policy_generation(
     if synchronizer is not None:
         return synchronizer.sync_weights(timer=timer, kv_scales=kv_scales) or {}
 
-    # Megatron generation backend needs explicit suspend/resume around refits.
-    if isinstance(policy_generation, MegatronGeneration):
-        policy_generation.suspend_for_refit()
-
-    if colocated_inference or isinstance(policy_generation, MegatronGeneration):
+    if colocated_inference:
         policy.offload_before_refit()
-    # Colocated inference needs to prepare for generation.
-    # Megatron non-colocated inference needs to enter inference mode after refit.
-    if colocated_inference or isinstance(policy_generation, MegatronGeneration):
         policy_generation.prepare_for_generation(tags=["weights"])
 
-    if (
-        not colocated_inference
-        and isinstance(policy_generation, MegatronGeneration)
-        and policy_generation.cfg["mcore_generation_config"]["refit_backend"]
-        == "nvshmem"
-    ):
-        futures_train = policy.preinit_nvshmem()
-        futures_inference = policy_generation.preinit_nvshmem_collective()
-        ray.get(futures_train + futures_inference)
-
     # Create a context manager that does nothing when timer is None
     timer_context = (
         timer.time("prepare_for_generation/transfer_and_update_weights")
@@ -2403,12 +2386,7 @@ def refit_policy_generation(
                 raise NotImplementedError(
                     "SGLang haven't implemented non-colocated inference mode. "
                 )
-            if isinstance(policy_generation, MegatronGeneration):
-                futures_train = policy.swap_weights_via_reshard(is_source=True)
-            else:
-                futures_train = policy.broadcast_weights_for_collective(
-                    kv_scales=kv_scales
-                )
+            futures_train = policy.broadcast_weights_for_collective(kv_scales=kv_scales)
             futures_inference = policy_generation.update_weights_from_collective()
             # wait for all futures to complete
             ray.get(futures_train)
@@ -2427,14 +2405,8 @@ def refit_policy_generation(
 
     if colocated_inference:
         policy.offload_after_refit()
-    # Colocated inference needs to prepare for generation.
-    # Megatron non-colocated inference needs to enter inference mode after refit.
-    if colocated_inference or isinstance(policy_generation, MegatronGeneration):
         policy_generation.prepare_for_generation(tags=["kv_cache"])
 
-    if isinstance(policy_generation, MegatronGeneration):
-        policy_generation.resume_after_refit()
-
     return {}
 
 
diff --git a/nemo_rl/models/generation/megatron/megatron_generation.py b/nemo_rl/models/generation/megatron/megatron_generation.py
index 69c22d09813..4f59b1ee1da 100644
--- a/nemo_rl/models/generation/megatron/megatron_generation.py
+++ b/nemo_rl/models/generation/megatron/megatron_generation.py
@@ -30,6 +30,7 @@
     merged_inference_megatron_cfg,
 )
 from nemo_rl.models.policy import PolicyConfig
+from nemo_rl.weight_sync.interfaces import WeightSynchronizer
 
 if TYPE_CHECKING:
     from nemo_rl.models.policy.lm_policy import Policy
@@ -120,6 +121,8 @@ def __init__(
         self.cfg: MCoreGenerationConfig = config["generation"]
         # Populated after the first prepare_for_generation (which starts the HTTP server).
         self.dp_openai_server_base_urls: list[Optional[str]] = []
+        # Installed by setup via create_weight_synchronizer.
+        self.weight_synchronizer: Optional["WeightSynchronizer"] = None
 
         if policy is not None:
             # Reuse the existing training policy.
diff --git a/nemo_rl/weight_sync/checkpoint_engine_weight_synchronizer.py b/nemo_rl/weight_sync/checkpoint_engine_weight_synchronizer.py
index 9e035fb4ff3..946547adbd6 100644
--- a/nemo_rl/weight_sync/checkpoint_engine_weight_synchronizer.py
+++ b/nemo_rl/weight_sync/checkpoint_engine_weight_synchronizer.py
@@ -74,9 +74,6 @@ def init_communicator(self) -> None:
     def is_stale(self) -> bool:
         return self._stale
 
-    def mark_stale(self) -> None:
-        self._stale = True
-
     def _release_after_refit(self) -> bool:
         cfg = self._checkpoint_engine_config
         return bool(cfg["engine_kwargs"][cfg["backend"]]["release_after_refit"])
diff --git a/nemo_rl/weight_sync/collective_weight_synchronizer.py b/nemo_rl/weight_sync/collective_weight_synchronizer.py
index a047915c81c..84e842dc018 100644
--- a/nemo_rl/weight_sync/collective_weight_synchronizer.py
+++ b/nemo_rl/weight_sync/collective_weight_synchronizer.py
@@ -98,9 +98,6 @@ def sync_weights(
     def is_stale(self) -> bool:
         return self._stale
 
-    def mark_stale(self) -> None:
-        self._stale = True
-
     def init_communicator(self) -> None:
         # prepare_refit_info is called before init_collective. This matches
         # distillation.py ordering. Neither call depends on the other today,
diff --git a/nemo_rl/weight_sync/factory.py b/nemo_rl/weight_sync/factory.py
index b76abf9933a..a6f80b4cc91 100644
--- a/nemo_rl/weight_sync/factory.py
+++ b/nemo_rl/weight_sync/factory.py
@@ -93,6 +93,21 @@ def create_weight_synchronizer(
     if refit_buffer_size_gb is not None and refit_buffer_size_gb <= 0:
         raise ValueError("refit_buffer_size_gb must be > 0")
 
+    if generation_backend == MEGATRON_BACKEND:
+        from nemo_rl.weight_sync.megatron_weight_synchronizer import (
+            MegatronWeightSynchronizer,
+        )
+
+        # One synchronizer serves both colocation modes; colocated is the
+        # degenerate path (no cross-group collective to wire).
+        return MegatronWeightSynchronizer(
+            policy=policy,
+            generation=generation,
+            colocated=colocated,
+            train_cluster=train_cluster,
+            inference_cluster=inference_cluster,
+        )
+
     if not colocated:
         if generation_backend == SGLANG_BACKEND:
             raise NotImplementedError(
diff --git a/nemo_rl/weight_sync/http_weight_synchronizer.py b/nemo_rl/weight_sync/http_weight_synchronizer.py
index 13cba3849b7..30274f01be6 100644
--- a/nemo_rl/weight_sync/http_weight_synchronizer.py
+++ b/nemo_rl/weight_sync/http_weight_synchronizer.py
@@ -100,9 +100,6 @@ def sync_weights(
     def is_stale(self) -> bool:
         return self._stale
 
-    def mark_stale(self) -> None:
-        self._stale = True
-
     def init_communicator(self) -> None:
         state_dict_info = self._policy.prepare_refit_info()
         self._generation.prepare_refit_info(state_dict_info)
diff --git a/nemo_rl/weight_sync/interfaces.py b/nemo_rl/weight_sync/interfaces.py
index 5b7623e7c68..175adb642ff 100644
--- a/nemo_rl/weight_sync/interfaces.py
+++ b/nemo_rl/weight_sync/interfaces.py
@@ -48,8 +48,9 @@ class WeightSynchronizer(ABC):
 
     Implementations handle the weight transfer for a specific transport
     mechanism (ZMQ IPC, HTTP, NCCL collectives). The orchestrator calls
-    sync_weights() and mark_stale() without knowing which transport is
-    being used or whether components are colocated.
+    sync_weights() without knowing which transport is being used or
+    whether components are colocated; per-step staleness bookkeeping is
+    owned by the training loop.
 
     Colocated transports (IPC, HTTP) own phase transitions internally
     (offload_before_refit, prepare_for_generation, offload_after_refit).
@@ -101,19 +102,11 @@ def sync_weights(
     def is_stale(self) -> bool:
         """Whether the generation backend's weights are out of date.
 
-        Returns True after mark_stale() is called and before the next
-        successful sync_weights() completes.
-        """
-        pass
-
-    @abstractmethod
-    def mark_stale(self) -> None:
-        """Mark weights as stale after a training step.
-
-        Should be called after every training step so the orchestrator
-        knows a sync is needed before the next generation phase. This
-        applies globally — all generation workers are considered stale
-        and will be updated atomically on the next ``sync_weights()`` call.
+        Returns True until the first successful sync_weights()
+        completes, so a fresh run always performs its initial sync (a
+        synchronizer that seeds current weights at construction may start
+        False to skip it). Per-step staleness is tracked by the training
+        loop, not here.
         """
         pass
 
diff --git a/nemo_rl/weight_sync/ipc_weight_synchronizer.py b/nemo_rl/weight_sync/ipc_weight_synchronizer.py
index ab62a9d6782..b23f4906268 100644
--- a/nemo_rl/weight_sync/ipc_weight_synchronizer.py
+++ b/nemo_rl/weight_sync/ipc_weight_synchronizer.py
@@ -108,9 +108,6 @@ def sync_weights(
     def is_stale(self) -> bool:
         return self._stale
 
-    def mark_stale(self) -> None:
-        self._stale = True
-
     def init_communicator(self) -> None:
         state_dict_info = self._policy.prepare_refit_info()
         self._generation.prepare_refit_info(state_dict_info)
diff --git a/nemo_rl/weight_sync/megatron_weight_synchronizer.py b/nemo_rl/weight_sync/megatron_weight_synchronizer.py
new file mode 100644
index 00000000000..4688ba26fb0
--- /dev/null
+++ b/nemo_rl/weight_sync/megatron_weight_synchronizer.py
@@ -0,0 +1,148 @@
+# 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.
+
+from contextlib import nullcontext
+from typing import Any, Optional
+
+import ray
+
+from nemo_rl.utils.timer import Timer
+from nemo_rl.weight_sync.interfaces import WeightSynchronizer
+
+
+class MegatronWeightSynchronizer(WeightSynchronizer):
+    """Weight synchronization for the Megatron generation backend, both colocation modes.
+
+    Colocated is the degenerate path: generation either aliases the training
+    weights outright (dual-mode) or re-partitions them into the worker's
+    dedicated inference model inside ``prepare_for_generation`` (when the
+    configured inference layout/impl differs) — a genuine parallelism-changing
+    transfer, but one the worker performs internally on wake. Sync therefore
+    reduces to dropping training-only buffers and re-entering inference mode.
+
+    Non-colocated adds the cross-group collective: the training and inference
+    workers are disjoint actor groups that rendezvous in mcore's
+    reshard-capable weight swap, with the engine suspended around the
+    transfer. That wiring (a joint refit process group over the configured
+    copy-service backend) is established once in ``init_communicator``.
+    """
+
+    def __init__(
+        self,
+        policy: Any,
+        generation: Any,
+        *,
+        colocated: bool,
+        train_cluster: Optional[Any] = None,
+        inference_cluster: Optional[Any] = None,
+    ):
+        if not colocated and (train_cluster is None or inference_cluster is None):
+            raise ValueError(
+                "train_cluster and inference_cluster are required for "
+                "non-colocated Megatron weight synchronization."
+            )
+        self._policy = policy
+        self._generation = generation
+        self._colocated = colocated
+        self._train_cluster = train_cluster
+        self._inference_cluster = inference_cluster
+        self._refit_backend: Optional[str] = None
+        self._stale = True
+
+    def init_communicator(self) -> None:
+        """Wire the cross-group refit collective (non-colocated only).
+
+        Colocated generation shares the training worker group, so there is
+        nothing to wire.
+        """
+        if self._colocated:
+            return
+        ip, port = self._train_cluster.get_master_address_and_port()
+        print(f"Using ip: {ip}, port: {port} for collective communication", flush=True)
+        train_world_size = self._train_cluster.world_size()
+        world_size = train_world_size + self._inference_cluster.world_size()
+        self._refit_backend = self._generation.cfg["mcore_generation_config"][
+            "refit_backend"
+        ]
+        futures_train = self._policy.init_collective_mcore_generation(
+            ip,
+            port,
+            world_size,
+            rank_offset=0,
+            refit_backend=self._refit_backend,
+        )
+        futures_inference = self._generation.init_collective(
+            ip,
+            port,
+            world_size,
+            train_world_size=train_world_size,
+            refit_backend=self._refit_backend,
+        )
+        ray.get(futures_train + futures_inference)
+
+    def sync_weights(
+        self,
+        *,
+        timer: Optional[Timer] = None,
+        kv_scales: Optional[dict[str, float]] = None,
+    ) -> Optional[dict[str, float]]:
+        if self._colocated:
+            # The wake below carries any configured reshard; the loop already
+            # slept the engine before training (or it has not started yet),
+            # so no suspend is needed.
+            self._policy.offload_before_refit()
+            self._generation.prepare_for_generation()
+            self._stale = False
+            return {}
+
+        # The engine serves continuously in non-colocated mode; pause it
+        # exactly around the swap.
+        self._generation.suspend_for_refit()
+        self._policy.offload_before_refit()
+        self._generation.prepare_for_generation(tags=["weights"])
+
+        if self._refit_backend == "nvshmem":
+            futures_train = self._policy.preinit_nvshmem()
+            futures_inference = self._generation.preinit_nvshmem_collective()
+            ray.get(futures_train + futures_inference)
+
+        timer_context = (
+            timer.time("prepare_for_generation/transfer_and_update_weights")
+            if timer is not None
+            else nullcontext()
+        )
+        with timer_context:
+            futures_train = self._policy.swap_weights_via_reshard(is_source=True)
+            futures_inference = self._generation.update_weights_from_collective()
+            ray.get(futures_train)
+            results = ray.get(futures_inference)
+            if not all(result for result in results if result is not None):
+                raise RuntimeError(
+                    "❌ Error: Updating weights for the generation policy failed "
+                    "during refit.\nThis often indicates an issue with the "
+                    "refit copy service or a problem within the generation "
+                    "backend.\n"
+                )
+
+        self._generation.prepare_for_generation(tags=["kv_cache"])
+        self._generation.resume_after_refit()
+        self._stale = False
+        return {}
+
+    @property
+    def is_stale(self) -> bool:
+        return self._stale
+
+    def shutdown(self) -> None:
+        """Nothing to tear down; the collective lives in the worker groups."""
diff --git a/nemo_rl/weight_sync/nccl_reshard_weight_synchronizer.py b/nemo_rl/weight_sync/nccl_reshard_weight_synchronizer.py
index 1d89341d36a..ab39bafde86 100644
--- a/nemo_rl/weight_sync/nccl_reshard_weight_synchronizer.py
+++ b/nemo_rl/weight_sync/nccl_reshard_weight_synchronizer.py
@@ -135,9 +135,6 @@ def sync_weights(
     def is_stale(self) -> bool:
         return self._stale
 
-    def mark_stale(self) -> None:
-        self._stale = True
-
     def init_communicator(self) -> None:
         train_parallelism = self._train_parallelism()
         gen_parallelism = self._gen_parallelism()
diff --git a/nemo_rl/weight_sync/vllm_remote_sparse_weight_synchronizer.py b/nemo_rl/weight_sync/vllm_remote_sparse_weight_synchronizer.py
index 021afd0429a..b124d0b78f1 100644
--- a/nemo_rl/weight_sync/vllm_remote_sparse_weight_synchronizer.py
+++ b/nemo_rl/weight_sync/vllm_remote_sparse_weight_synchronizer.py
@@ -286,9 +286,6 @@ def sync_weights(
     def is_stale(self) -> bool:
         return self._stale
 
-    def mark_stale(self) -> None:
-        self._stale = True
-
     def _run_policy_workers(self, method_name: str, **kwargs: Any) -> list[Any]:
         workers = self._policy.worker_group
         count = len(workers.workers)
diff --git a/tests/unit/weight_sync/test_checkpoint_engine_weight_synchronizer.py b/tests/unit/weight_sync/test_checkpoint_engine_weight_synchronizer.py
index 8a02257e8fa..1897a8924ec 100644
--- a/tests/unit/weight_sync/test_checkpoint_engine_weight_synchronizer.py
+++ b/tests/unit/weight_sync/test_checkpoint_engine_weight_synchronizer.py
@@ -244,8 +244,6 @@ def test_sync_weights_runs_checkpoint_engine_lifecycle(self, mock_ray):
             (0, 2, 2, ["policy-0", "policy-1", "generation-0", "generation-1"]),
             (2, 2, 2, ["policy-0", "policy-1", "generation-0", "generation-1"]),
         ]
-        sync.mark_stale()
-        assert sync.is_stale
         sync.shutdown()
         assert sync._generation.worker_group.calls[-1] == (
             "checkpoint_engine_rpc",
diff --git a/tests/unit/weight_sync/test_vllm_remote_sparse_weight_synchronizer.py b/tests/unit/weight_sync/test_vllm_remote_sparse_weight_synchronizer.py
index d55acb5a055..3724f93835f 100644
--- a/tests/unit/weight_sync/test_vllm_remote_sparse_weight_synchronizer.py
+++ b/tests/unit/weight_sync/test_vllm_remote_sparse_weight_synchronizer.py
@@ -187,7 +187,6 @@ def test_shutdown_cancels_pending_work_and_stops_zmq(self, mock_ray):
         sync._targets = ["tcp://relay"]
         sync._stale = False
 
-        sync.mark_stale()
         sync.shutdown()
 
         assert mock_ray.cancel.call_count == 2
diff --git a/tests/unit/weight_sync/test_weight_synchronizer.py b/tests/unit/weight_sync/test_weight_synchronizer.py
index 094511135b4..17f553b6518 100644
--- a/tests/unit/weight_sync/test_weight_synchronizer.py
+++ b/tests/unit/weight_sync/test_weight_synchronizer.py
@@ -31,6 +31,9 @@
     HTTPWeightSynchronizer,
 )
 from nemo_rl.weight_sync.interfaces import WeightSynchronizer
+from nemo_rl.weight_sync.megatron_weight_synchronizer import (
+    MegatronWeightSynchronizer,
+)
 from nemo_rl.weight_sync.ipc_weight_synchronizer import (
     IPCWeightSynchronizer,
 )
@@ -174,16 +177,6 @@ def test_dynamic_buffer_size(self, mock_ray, monkeypatch):
         expected = int(10 * (1024**3) * 0.3)
         assert call_kwargs.kwargs["buffer_size_bytes"] == expected
 
-    def test_mark_stale(self):
-        policy = _mock_policy()
-        gen = _mock_generation()
-        sync = IPCWeightSynchronizer(policy, gen)
-
-        sync._stale = False
-        assert not sync.is_stale
-        sync.mark_stale()
-        assert sync.is_stale
-
     def test_init_communicator(self):
         policy = _mock_policy()
         gen = _mock_generation()
@@ -284,16 +277,6 @@ def test_fixed_buffer_size(self, mock_ray):
         assert call_kwargs.kwargs["rollout_engine_urls"] == ["http://localhost:30000"]
         assert call_kwargs.kwargs["buffer_size_bytes"] == 2 * (1024**3)
 
-    def test_mark_stale(self):
-        policy = _mock_policy()
-        gen = _mock_generation()
-        sync = HTTPWeightSynchronizer(policy, gen)
-
-        sync._stale = False
-        assert not sync.is_stale
-        sync.mark_stale()
-        assert sync.is_stale
-
     def test_init_communicator(self):
         policy = _mock_policy()
         gen = _mock_generation()
@@ -496,6 +479,112 @@ def test_shutdown_drops_the_generation_handle(self):
 # ---------------------------------------------------------------------------
 
 
+def _mock_megatron_generation(refit_backend="nccl", **overrides):
+    gen = _mock_generation(**overrides)
+    gen.cfg = {"mcore_generation_config": {"refit_backend": refit_backend}}
+    gen.suspend_for_refit.return_value = None
+    gen.resume_after_refit.return_value = None
+    gen.preinit_nvshmem_collective.return_value = [MagicMock()]
+    return gen
+
+
+def _mock_megatron_policy(**overrides):
+    policy = _mock_policy(**overrides)
+    policy.swap_weights_via_reshard.return_value = [MagicMock()]
+    policy.init_collective_mcore_generation.return_value = [MagicMock()]
+    policy.preinit_nvshmem.return_value = [MagicMock()]
+    return policy
+
+
+class TestMegatronWeightSynchronizer:
+    def test_non_colocated_requires_clusters(self):
+        with pytest.raises(ValueError):
+            MegatronWeightSynchronizer(
+                _mock_megatron_policy(), _mock_megatron_generation(), colocated=False
+            )
+
+    def test_colocated_sync_is_offload_and_wake(self):
+        policy = _mock_megatron_policy()
+        gen = _mock_megatron_generation()
+        sync = MegatronWeightSynchronizer(policy, gen, colocated=True)
+
+        sync.init_communicator()  # no collective to wire
+        policy.init_collective_mcore_generation.assert_not_called()
+
+        assert sync.is_stale
+        assert sync.sync_weights() == {}
+        policy.offload_before_refit.assert_called_once()
+        gen.prepare_for_generation.assert_called_once_with()
+        gen.suspend_for_refit.assert_not_called()
+        policy.swap_weights_via_reshard.assert_not_called()
+        assert not sync.is_stale
+
+    @patch("nemo_rl.weight_sync.megatron_weight_synchronizer.ray")
+    def test_non_colocated_sync_sequence(self, mock_ray):
+        mock_ray.get.side_effect = lambda futures: [True for _ in futures]
+        policy = _mock_megatron_policy()
+        gen = _mock_megatron_generation()
+        sync = MegatronWeightSynchronizer(
+            policy,
+            gen,
+            colocated=False,
+            train_cluster=_mock_cluster(),
+            inference_cluster=_mock_cluster(),
+        )
+
+        sync.init_communicator()
+        policy.init_collective_mcore_generation.assert_called_once()
+        gen.init_collective.assert_called_once()
+
+        assert sync.sync_weights() == {}
+        gen.suspend_for_refit.assert_called_once()
+        policy.offload_before_refit.assert_called_once()
+        policy.swap_weights_via_reshard.assert_called_once_with(is_source=True)
+        gen.update_weights_from_collective.assert_called_once()
+        gen.resume_after_refit.assert_called_once()
+        # prepare called for the weights phase and then the kv_cache phase
+        tags = [c.kwargs.get("tags") for c in gen.prepare_for_generation.call_args_list]
+        assert tags == [["weights"], ["kv_cache"]]
+        # no nvshmem preinit on the nccl backend
+        policy.preinit_nvshmem.assert_not_called()
+        assert not sync.is_stale
+
+    @patch("nemo_rl.weight_sync.megatron_weight_synchronizer.ray")
+    def test_non_colocated_nvshmem_preinits(self, mock_ray):
+        mock_ray.get.side_effect = lambda futures: [True for _ in futures]
+        policy = _mock_megatron_policy()
+        gen = _mock_megatron_generation(refit_backend="nvshmem")
+        sync = MegatronWeightSynchronizer(
+            policy,
+            gen,
+            colocated=False,
+            train_cluster=_mock_cluster(),
+            inference_cluster=_mock_cluster(),
+        )
+        sync.init_communicator()
+        sync.sync_weights()
+        policy.preinit_nvshmem.assert_called_once()
+        gen.preinit_nvshmem_collective.assert_called_once()
+
+    @patch("nemo_rl.weight_sync.megatron_weight_synchronizer.ray")
+    def test_non_colocated_failed_update_raises(self, mock_ray):
+        # swap futures resolve fine; the inference-side results report failure
+        mock_ray.get.side_effect = lambda futures: [False for _ in futures]
+        policy = _mock_megatron_policy()
+        gen = _mock_megatron_generation()
+        sync = MegatronWeightSynchronizer(
+            policy,
+            gen,
+            colocated=False,
+            train_cluster=_mock_cluster(),
+            inference_cluster=_mock_cluster(),
+        )
+        sync.init_communicator()
+        with pytest.raises(RuntimeError):
+            sync.sync_weights()
+        assert sync.is_stale
+
+
 class TestFactory:
     def test_colocated_vllm_returns_ipc(self):
         policy = _mock_policy()
@@ -519,7 +608,7 @@ def test_colocated_sglang_returns_http(self):
         )
         assert isinstance(sync, HTTPWeightSynchronizer)
 
-    def test_colocated_megatron_returns_ipc(self):
+    def test_colocated_megatron_returns_megatron_synchronizer(self):
         policy = _mock_policy()
         gen = _mock_generation()
         sync = create_weight_synchronizer(
@@ -528,7 +617,20 @@ def test_colocated_megatron_returns_ipc(self):
             generation_backend=MEGATRON_BACKEND,
             colocated=True,
         )
-        assert isinstance(sync, IPCWeightSynchronizer)
+        assert isinstance(sync, MegatronWeightSynchronizer)
+
+    def test_non_colocated_megatron_returns_megatron_synchronizer(self):
+        policy = _mock_policy()
+        gen = _mock_generation()
+        sync = create_weight_synchronizer(
+            policy=policy,
+            generation=gen,
+            generation_backend=MEGATRON_BACKEND,
+            colocated=False,
+            train_cluster=_mock_cluster(),
+            inference_cluster=_mock_cluster(),
+        )
+        assert isinstance(sync, MegatronWeightSynchronizer)
 
     def test_non_colocated_vllm_returns_collective(self):
         policy = _mock_policy()

From 75b1c81eb7389605bed03e3ce179e72aed038e71 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Mon, 10 Aug 2026 05:59:52 -0500
Subject: [PATCH 08/25] Add offload_optimizer_for_refit config

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 examples/configs/distillation_math.yaml       |  1 +
 examples/configs/grpo_math_1B.yaml            |  3 +-
 examples/configs/ppo_math_1B.yaml             |  1 +
 nemo_rl/algorithms/grpo.py                    | 25 +++++--------
 nemo_rl/models/megatron/config.py             |  1 +
 nemo_rl/models/megatron/setup.py              |  2 ++
 .../policy/workers/megatron_policy_worker.py  |  2 ++
 .../models/megatron/test_megatron_setup.py    |  3 ++
 .../models/policy/test_megatron_worker.py     | 35 +++++++++++++++++++
 .../reference_configs/distillation_math.yaml  |  1 +
 .../unit/reference_configs/grpo_math_1B.yaml  |  3 +-
 .../ppo_math_1B_megatron.yaml                 |  1 +
 12 files changed, 60 insertions(+), 18 deletions(-)

diff --git a/examples/configs/distillation_math.yaml b/examples/configs/distillation_math.yaml
index 59f7616e0e1..eff54539827 100644
--- a/examples/configs/distillation_math.yaml
+++ b/examples/configs/distillation_math.yaml
@@ -44,6 +44,7 @@ policy: &POLICY_BASE
     logprob_chunk_size: null
 
     offload_optimizer_for_logprob: false
+    offload_optimizer_for_refit: true # Offload the optimizer state to CPU around each colocated-generation refit; configs that do not inherit this default keep the optimizer on GPU (needs the extra headroom)
 
     dtensor_cfg: &DTENSOR_BASE
         enabled: true
diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml
index 62e94d7b0b4..05825750e37 100644
--- a/examples/configs/grpo_math_1B.yaml
+++ b/examples/configs/grpo_math_1B.yaml
@@ -136,7 +136,8 @@ policy:
   max_total_sequence_length: 512
   precision: "bfloat16"
   logprob_chunk_size: null
-  offload_optimizer_for_logprob: false # Only useful for non-colocated generation since colocated generation will always offload optimizer to cuda before refit
+  offload_optimizer_for_logprob: false # Only useful for non-colocated generation since colocated generation offloads the optimizer before refit (see offload_optimizer_for_refit)
+  offload_optimizer_for_refit: true # Offload the optimizer state to CPU around each colocated-generation refit; configs that do not inherit this default keep the optimizer on GPU (needs the extra headroom)
 
   dtensor_cfg:
     _v2: true
diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml
index 97c5f9688de..e99ee8301f4 100644
--- a/examples/configs/ppo_math_1B.yaml
+++ b/examples/configs/ppo_math_1B.yaml
@@ -102,6 +102,7 @@ policy:
   precision: "bfloat16"
   logprob_chunk_size: null
   offload_optimizer_for_logprob: false
+  offload_optimizer_for_refit: true # Offload the optimizer state to CPU around each colocated-generation refit; configs that do not inherit this default keep the optimizer on GPU (needs the extra headroom)
 
   dtensor_cfg:
     _v2: true
diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py
index aeeaf49e09a..b75d67c1025 100644
--- a/nemo_rl/algorithms/grpo.py
+++ b/nemo_rl/algorithms/grpo.py
@@ -1149,22 +1149,15 @@ def init_sglang():
     def init_megatron_generation(policy=None):
         """Initialize Megatron generation."""
         t0 = time.perf_counter()
-        if colocated_inference:
-            mg = MegatronGeneration(
-                policy=policy,
-                config=policy_config,
-                tokenizer=tokenizer,
-                processor=processor,
-            )
-        else:
-            mg = MegatronGeneration(
-                cluster=inference_cluster,
-                config=policy_config,
-                tokenizer=tokenizer,
-                processor=processor,
-                weights_path=weights_path,
-                skip_weight_load=True,
-            )
+        mg = MegatronGeneration(
+            config=policy_config,
+            tokenizer=tokenizer,
+            cluster=None if colocated_inference else inference_cluster,
+            policy=policy if colocated_inference else None,
+            processor=processor,
+            weights_path=weights_path,
+            skip_weight_load=not colocated_inference,
+        )
         return mg, time.perf_counter() - t0
 
     def initialize_generation_with_policy(
diff --git a/nemo_rl/models/megatron/config.py b/nemo_rl/models/megatron/config.py
index ccda91f3325..06a8a3b3df1 100644
--- a/nemo_rl/models/megatron/config.py
+++ b/nemo_rl/models/megatron/config.py
@@ -37,6 +37,7 @@ class RuntimeConfig(NamedTuple):
     dtype: torch.dtype
     optimizer_cpu_offload: bool
     offload_optimizer_for_logprob: bool
+    offload_optimizer_for_refit: bool
     is_generation_colocated: Optional[bool]
     sampling_params: Optional[TrainingSamplingParams]
     final_padded_vocab_size: int
diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py
index b158921402a..b63e1aa9412 100644
--- a/nemo_rl/models/megatron/setup.py
+++ b/nemo_rl/models/megatron/setup.py
@@ -352,6 +352,7 @@ def validate_and_set_config(
     # Optimizer configuration
     optimizer_cpu_offload = config["megatron_cfg"]["optimizer"]["optimizer_cpu_offload"]
     offload_optimizer_for_logprob = config["offload_optimizer_for_logprob"]
+    offload_optimizer_for_refit = bool(config.get("offload_optimizer_for_refit"))
 
     # Reward models are not yet supported with Megatron.
     if "reward_model_cfg" in config and config["reward_model_cfg"]["enabled"]:
@@ -402,6 +403,7 @@ def validate_and_set_config(
         dtype,
         optimizer_cpu_offload,
         offload_optimizer_for_logprob,
+        offload_optimizer_for_refit,
         is_generation_colocated,
         sampling_params,
         final_padded_vocab_size,
diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py
index 271d90d9079..0a3528f41a6 100644
--- a/nemo_rl/models/policy/workers/megatron_policy_worker.py
+++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py
@@ -499,6 +499,7 @@ def __init__(
         self.offload_optimizer_for_logprob = (
             runtime_config.offload_optimizer_for_logprob
         )
+        self.offload_optimizer_for_refit = runtime_config.offload_optimizer_for_refit
         self.is_generation_colocated = runtime_config.is_generation_colocated
         self.final_padded_vocab_size = runtime_config.final_padded_vocab_size
         self.sampling_params = runtime_config.sampling_params
@@ -2936,6 +2937,7 @@ def offload_before_refit(self):
             hasattr(self, "optimizer")
             and self.optimizer is not None
             and not self.optimizer_cpu_offload
+            and self.offload_optimizer_for_refit
         ):
             self.move_optimizer("cpu")
 
diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py
index e3f1c4c5c50..8b7d2a9375b 100644
--- a/tests/unit/models/megatron/test_megatron_setup.py
+++ b/tests/unit/models/megatron/test_megatron_setup.py
@@ -2150,6 +2150,7 @@ def test_generation_colocation_detection(self):
                 )
 
                 assert runtime_config.is_generation_colocated is True
+                assert runtime_config.offload_optimizer_for_refit is False
 
 
 @pytest.mark.mcore
@@ -2166,6 +2167,7 @@ def test_runtime_config_fields(self):
             dtype=torch.bfloat16,
             optimizer_cpu_offload=False,
             offload_optimizer_for_logprob=True,
+            offload_optimizer_for_refit=False,
             is_generation_colocated=True,
             sampling_params=None,
             final_padded_vocab_size=32000,
@@ -2175,6 +2177,7 @@ def test_runtime_config_fields(self):
         assert runtime_config.optimizer_cpu_offload is False
         assert runtime_config.offload_optimizer_for_logprob is True
         assert runtime_config.is_generation_colocated is True
+        assert runtime_config.offload_optimizer_for_refit is False
         assert runtime_config.sampling_params is None
         assert runtime_config.final_padded_vocab_size == 32000
 
diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py
index f6bc54dd0bb..825aedc6203 100644
--- a/tests/unit/models/policy/test_megatron_worker.py
+++ b/tests/unit/models/policy/test_megatron_worker.py
@@ -266,6 +266,41 @@ def cuda(self):
     assert events.index("finalize_async_save") < events.index("move_model")
 
 
+@pytest.mark.parametrize("offload_optimizer", [False, True])
+def test_megatron_offload_before_refit_honors_offload_optimizer_for_refit(
+    monkeypatch, offload_optimizer
+):
+    """offload_optimizer_for_refit=False must leave the optimizer untouched."""
+    from nemo_rl.models.policy.workers.megatron_policy_worker import (
+        MegatronPolicyWorkerImpl,
+    )
+
+    moved = []
+    worker = object.__new__(MegatronPolicyWorkerImpl)
+    worker.model = object()
+    worker.optimizer = object()
+    worker.optimizer_cpu_offload = False
+    worker.offload_optimizer_for_refit = offload_optimizer
+    worker.fp8_cfg = None
+    worker.cfg = {"megatron_cfg": {"clear_memory_caches_before_refit": False}}
+    worker.finalize_async_save = lambda: None
+    worker.move_model = lambda model, device, move_params, move_grads: model
+    worker.move_optimizer = lambda device: moved.append(device)
+
+    class _AllocatorWakeup:
+        def cuda(self):
+            pass
+
+    monkeypatch.setattr(torch.cuda, "memory_allocated", lambda *args, **kwargs: 0)
+    monkeypatch.setattr(torch.cuda, "memory_reserved", lambda *args, **kwargs: 0)
+    monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None)
+    monkeypatch.setattr(torch, "randn", lambda *args, **kwargs: _AllocatorWakeup())
+
+    MegatronPolicyWorkerImpl.offload_before_refit(worker)
+
+    assert moved == (["cpu"] if offload_optimizer else [])
+
+
 def test_megatron_offload_after_refit_finalizes_before_model_move(monkeypatch):
     """Checkpoint CUDA IPC handles must be dropped before model storage is replaced."""
     from nemo_rl.models.policy.workers.megatron_policy_worker import (
diff --git a/tests/unit/reference_configs/distillation_math.yaml b/tests/unit/reference_configs/distillation_math.yaml
index 5f5f2d8ddfb..4ac6b56b997 100644
--- a/tests/unit/reference_configs/distillation_math.yaml
+++ b/tests/unit/reference_configs/distillation_math.yaml
@@ -44,6 +44,7 @@ policy: &POLICY_BASE
     logprob_chunk_size: null
 
     offload_optimizer_for_logprob: false
+    offload_optimizer_for_refit: true # Offload the optimizer state to CPU around each colocated-generation refit; configs that do not inherit this default keep the optimizer on GPU (needs the extra headroom)
 
     dtensor_cfg: &DTENSOR_BASE
         enabled: true
diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml
index d5738e8c762..22c602f5dfe 100644
--- a/tests/unit/reference_configs/grpo_math_1B.yaml
+++ b/tests/unit/reference_configs/grpo_math_1B.yaml
@@ -142,7 +142,8 @@ policy:
   max_total_sequence_length: 512
   precision: "bfloat16"
   logprob_chunk_size: null
-  offload_optimizer_for_logprob: false # Only useful for non-colocated generation since colocated generation will always offload optimizer to cuda before refit
+  offload_optimizer_for_logprob: false # Only useful for non-colocated generation since colocated generation offloads the optimizer before refit (see offload_optimizer_for_refit)
+  offload_optimizer_for_refit: true # Offload the optimizer state to CPU around each colocated-generation refit; configs that do not inherit this default keep the optimizer on GPU (needs the extra headroom)
 
   dtensor_cfg:
     _v2: true
diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml
index 9858962d779..2d6620c4ef1 100644
--- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml
+++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml
@@ -101,6 +101,7 @@ policy:
   precision: "bfloat16"
   logprob_chunk_size: null
   offload_optimizer_for_logprob: false
+  offload_optimizer_for_refit: true # Offload the optimizer state to CPU around each colocated-generation refit; configs that do not inherit this default keep the optimizer on GPU (needs the extra headroom)
 
   dtensor_cfg:
     _v2: true

From 6d0a720b34e5213359e26b941295ced88ebccf55 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Mon, 10 Aug 2026 09:22:33 -0500
Subject: [PATCH 09/25] Adopt setup_timing_metrics in the synchronizer init
 block

The #3499 rebase replaced worker_init_timing_metrics with the
SetupTimingMetrics dataclass; the eager initial sync is the
generation load in that scheme.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 nemo_rl/algorithms/grpo.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py
index b75d67c1025..6676b81af05 100644
--- a/nemo_rl/algorithms/grpo.py
+++ b/nemo_rl/algorithms/grpo.py
@@ -1456,14 +1456,12 @@ def init_trtllm():
             inference_cluster=None if colocated_inference else inference_cluster,
         )
         policy_generation.weight_synchronizer.init_communicator()
-        worker_init_timing_metrics["collective_init_time_s"] = time.perf_counter() - t0
+        setup_timing_metrics.collective_init_time_s = time.perf_counter() - t0
         if not colocated_inference:
             # Load the model weights now.
             t0 = time.perf_counter()
             policy_generation.weight_synchronizer.sync_weights()
-            worker_init_timing_metrics["initial_weight_sync_time_s"] = (
-                time.perf_counter() - t0
-            )
+            setup_timing_metrics.generation_init_load_time_s = time.perf_counter() - t0
     # if it is not colocated inference, initialize collective communication for update weights
     elif (
         not colocated_inference

From 356d25f189c2296df23b006f0e7af8d68561e05a Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Mon, 10 Aug 2026 09:50:35 -0500
Subject: [PATCH 10/25] lint

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 pyrefly.toml                                       | 1 +
 tests/unit/weight_sync/test_weight_synchronizer.py | 6 +++---
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/pyrefly.toml b/pyrefly.toml
index 004ee094136..634d5ed0bdd 100644
--- a/pyrefly.toml
+++ b/pyrefly.toml
@@ -238,6 +238,7 @@ project-includes = [
   "nemo_rl/weight_sync/http_weight_synchronizer.py",
   "nemo_rl/weight_sync/interfaces.py",
   "nemo_rl/weight_sync/ipc_weight_synchronizer.py",
+  "nemo_rl/weight_sync/megatron_weight_synchronizer.py",
   "nemo_rl/weight_sync/vllm_remote_sparse_weight_synchronizer.py",
   "nemo_rl/weight_sync/nccl_reshard_utils.py",
   "nemo_rl/weight_sync/nccl_reshard_weight_synchronizer.py",
diff --git a/tests/unit/weight_sync/test_weight_synchronizer.py b/tests/unit/weight_sync/test_weight_synchronizer.py
index 17f553b6518..be870962f0b 100644
--- a/tests/unit/weight_sync/test_weight_synchronizer.py
+++ b/tests/unit/weight_sync/test_weight_synchronizer.py
@@ -31,12 +31,12 @@
     HTTPWeightSynchronizer,
 )
 from nemo_rl.weight_sync.interfaces import WeightSynchronizer
-from nemo_rl.weight_sync.megatron_weight_synchronizer import (
-    MegatronWeightSynchronizer,
-)
 from nemo_rl.weight_sync.ipc_weight_synchronizer import (
     IPCWeightSynchronizer,
 )
+from nemo_rl.weight_sync.megatron_weight_synchronizer import (
+    MegatronWeightSynchronizer,
+)
 from nemo_rl.weight_sync.nccl_reshard_utils import build_nccl_reshard_refit_info
 from nemo_rl.weight_sync.nccl_reshard_weight_synchronizer import (
     NcclReshardWeightSynchronizer,

From ea559a474c25dd2e434a43bee590b15346afcbe8 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Mon, 10 Aug 2026 11:06:26 -0500
Subject: [PATCH 11/25] Fix the non-colocated refit test to dispatch via the
 synchronizer

refit_policy_generation no longer hand-rolls the Megatron collective;
without a wired synchronizer it fell through to the vLLM broadcast
path (AttributeError: model_update_group).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 .../generation/test_megatron_generation.py    | 28 ++++++++-----------
 1 file changed, 12 insertions(+), 16 deletions(-)

diff --git a/tests/unit/models/generation/test_megatron_generation.py b/tests/unit/models/generation/test_megatron_generation.py
index 0ad052c1057..c41b6938267 100644
--- a/tests/unit/models/generation/test_megatron_generation.py
+++ b/tests/unit/models/generation/test_megatron_generation.py
@@ -16,7 +16,6 @@
 from copy import deepcopy
 
 import pytest
-import ray
 import torch
 
 from nemo_rl.algorithms.grpo import refit_policy_generation
@@ -26,6 +25,9 @@
 from nemo_rl.models.generation.megatron import MegatronGeneration
 from nemo_rl.models.policy import PolicyConfig
 from nemo_rl.models.policy.lm_policy import Policy
+from nemo_rl.weight_sync.megatron_weight_synchronizer import (
+    MegatronWeightSynchronizer,
+)
 
 model_name = "Qwen/Qwen3-0.6B"
 
@@ -476,22 +478,16 @@ def test_megatron_generation_non_colocated_refit(
             skip_weight_load=skip_weight_load,
         )
 
-        # init the refit collective on both sides.
-        ip, port = policy_cluster_separate.get_master_address_and_port()
-        train_world_size = policy_cluster_separate.world_size()
-        world_size = train_world_size + generation_cluster.world_size()
-        refit_backend = config["generation"]["mcore_generation_config"]["refit_backend"]
-        futures_train = policy.init_collective_mcore_generation(
-            ip, port, world_size, rank_offset=0, refit_backend=refit_backend
-        )
-        futures_inference = mg.init_collective(
-            ip,
-            port,
-            world_size,
-            train_world_size=train_world_size,
-            refit_backend=refit_backend,
+        # Wire the refit collective the way grpo.setup does: through the
+        # weight synchronizer, which refit_policy_generation dispatches to.
+        mg.weight_synchronizer = MegatronWeightSynchronizer(
+            policy,
+            mg,
+            colocated=False,
+            train_cluster=policy_cluster_separate,
+            inference_cluster=generation_cluster,
         )
-        ray.get(futures_train + futures_inference)
+        mg.weight_synchronizer.init_communicator()
 
         # refit the inference engine from the training weights, then generate
         refit_policy_generation(policy, mg, False)

From 9ffe4490aa38af5c312dcf620211ac72f1d510ae Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Mon, 10 Aug 2026 22:18:57 -0500
Subject: [PATCH 12/25] Fix inference-optimized incompatibility

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 nemo_rl/models/megatron/setup.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py
index b63e1aa9412..0f6ac62af17 100644
--- a/nemo_rl/models/megatron/setup.py
+++ b/nemo_rl/models/megatron/setup.py
@@ -913,6 +913,12 @@ def _apply_moe_config(model_cfg: Any, config: PolicyConfig) -> None:
         model_cfg.moe_router_group_topk = config["megatron_cfg"][
             "moe_router_group_topk"
         ]
+    if (
+        config["megatron_cfg"].get("transformer_impl") == "inference_optimized"
+        and getattr(model_cfg, "moe_router_num_groups", None) == 1
+    ):
+        model_cfg.moe_router_num_groups = None
+        model_cfg.moe_router_group_topk = None
     if "moe_pad_experts_for_cuda_graph_inference" in config["megatron_cfg"]:
         model_cfg.moe_pad_experts_for_cuda_graph_inference = config["megatron_cfg"][
             "moe_pad_experts_for_cuda_graph_inference"

From 844e96426041aa803ca2cde39a536e63a27afa4d Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 11 Aug 2026 11:46:00 -0500
Subject: [PATCH 13/25] Properly fix inference config mismatch

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 ...0BA3B-2n8g-megatron_colocated_reshard.yaml |  1 +
 nemo_rl/models/megatron/setup.py              | 26 +++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.yaml
index 12f9d5cbacd..e3afbd27f99 100644
--- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.yaml
+++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.yaml
@@ -30,6 +30,7 @@ policy:
     backend: megatron
     mcore_generation_config:
       transformer_impl: inference_optimized
+      moe_router_dtype: fp32
       activation_checkpointing: false
       tensor_model_parallel_size: 4
       expert_model_parallel_size: 4
diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py
index 0f6ac62af17..0a10faefd95 100644
--- a/nemo_rl/models/megatron/setup.py
+++ b/nemo_rl/models/megatron/setup.py
@@ -923,6 +923,28 @@ def _apply_moe_config(model_cfg: Any, config: PolicyConfig) -> None:
         model_cfg.moe_pad_experts_for_cuda_graph_inference = config["megatron_cfg"][
             "moe_pad_experts_for_cuda_graph_inference"
         ]
+    generation_cfg = config.get("generation")
+    mcore_gen_cfg = (
+        (generation_cfg.get("mcore_generation_config") or {})
+        if generation_cfg is not None and generation_cfg.get("backend") == "megatron"
+        else {}
+    )
+    if (
+        mcore_gen_cfg.get("cuda_graph_impl") == "local"
+        and mcore_gen_cfg.get(
+            "transformer_impl", config["megatron_cfg"].get("transformer_impl")
+        )
+        != "inference_optimized"
+        and model_cfg.expert_model_parallel_size > 1
+        and "moe_pad_experts_for_cuda_graph_inference" not in config["megatron_cfg"]
+        and "moe_pad_experts_for_cuda_graph_inference" not in mcore_gen_cfg
+    ):
+        print(
+            "[_apply_moe_config] Setting "
+            "moe_pad_experts_for_cuda_graph_inference=True: CUDA-graph "
+            "inference with expert parallelism requires padded experts."
+        )
+        model_cfg.moe_pad_experts_for_cuda_graph_inference = True
     model_cfg.moe_shared_expert_overlap = config["megatron_cfg"][
         "moe_shared_expert_overlap"
     ]
@@ -1529,6 +1551,10 @@ def build_inference_model(
     inference_provider.recompute_granularity = None
     inference_provider.recompute_method = None
     inference_provider.recompute_num_layers = None
+    if inference_provider.transformer_impl == "inference_optimized":
+        inference_provider.moe_pad_experts_for_cuda_graph_inference = False
+    # Need to run __post__init__ with the correct config.
+    TransformerConfig.__post_init__(inference_provider)
 
     world_size = torch.distributed.get_world_size()
     inference_pg_collection = build_inference_pg_collection(

From 0cd9e5e4740c613b5bb6df5ad39a35059fc803b6 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 11 Aug 2026 21:11:47 -0500
Subject: [PATCH 14/25] Use a QKV-bias-free model in the reshard test

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 tests/functional/grpo_megatron_generation_colocated_reshard.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/functional/grpo_megatron_generation_colocated_reshard.sh b/tests/functional/grpo_megatron_generation_colocated_reshard.sh
index de87c759379..9aea0df5c39 100755
--- a/tests/functional/grpo_megatron_generation_colocated_reshard.sh
+++ b/tests/functional/grpo_megatron_generation_colocated_reshard.sh
@@ -25,7 +25,7 @@ cd $PROJECT_ROOT
 uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \
     $PROJECT_ROOT/examples/run_grpo.py \
     --config $PROJECT_ROOT/examples/configs/grpo_math_1B_megatron.yaml \
-    policy.model_name=Qwen/Qwen2.5-0.5B \
+    policy.model_name=Qwen/Qwen3-0.6B \
     grpo.num_prompts_per_step=2 \
     grpo.num_generations_per_prompt=4 \
     policy.train_global_batch_size=4 \

From 7cf972ae38984a11165ac8fd39a1cb783571e61f Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 11 Aug 2026 22:56:38 -0500
Subject: [PATCH 15/25] Add refit test that doesn't need TMS

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 ...nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml} | 10 +++++-----
 ...o-nanov3-30BA3B-4n4g-megatron_colocated_reshard.sh} |  4 ++--
 tests/test_suites/nightly.txt                          |  1 -
 tests/test_suites/nightly_gb200.txt                    |  3 +++
 4 files changed, 10 insertions(+), 8 deletions(-)
 rename examples/configs/recipes/llm/{grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.yaml => grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml} (84%)
 rename tests/test_suites/llm/{grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.sh => grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.sh} (97%)

diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml
similarity index 84%
rename from examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.yaml
rename to examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml
index e3afbd27f99..b24bdb06f0a 100644
--- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.yaml
+++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml
@@ -6,7 +6,7 @@ grpo:
   max_num_steps: 500
 checkpointing:
   enabled: false
-  checkpoint_dir: results/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard
+  checkpoint_dir: results/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard
   save_period: 100
 policy:
   model_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16
@@ -37,12 +37,12 @@ policy:
       sequence_parallel: true
       refit_backend: nccl
 logger:
-  log_dir: logs/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard
+  log_dir: logs/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard
   wandb_enabled: true
   tensorboard_enabled: true
   wandb:
     project: nemo-rl
-    name: grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard
+    name: grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard
 cluster:
-  gpus_per_node: 8
-  num_nodes: 2
+  gpus_per_node: 4
+  num_nodes: 4
diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.sh
similarity index 97%
rename from tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.sh
rename to tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.sh
index 43b19a166e4..d0868cf7cda 100755
--- a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.sh
+++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.sh
@@ -3,8 +3,8 @@ SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
 source $SCRIPT_DIR/common.env
 
 # ===== BEGIN CONFIG =====
-NUM_NODES=2
-GPUS_PER_NODE=8
+NUM_NODES=4
+GPUS_PER_NODE=4
 STEPS_PER_RUN=10
 MAX_STEPS=10
 NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN ))  # Round up
diff --git a/tests/test_suites/nightly.txt b/tests/test_suites/nightly.txt
index 3ce95e53c82..64fb93790e4 100644
--- a/tests/test_suites/nightly.txt
+++ b/tests/test_suites/nightly.txt
@@ -117,7 +117,6 @@ tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-lora.sh
 tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp.sh
 tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation.sh
 tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym.sh
-tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_colocated_reshard.sh
 
 # Nano-v3.5
 tests/test_suites/llm/dapo-nanov3.5-30BA3B-4n8g-automodel.sh
diff --git a/tests/test_suites/nightly_gb200.txt b/tests/test_suites/nightly_gb200.txt
index 872f6c0bfd8..aa2dccea837 100644
--- a/tests/test_suites/nightly_gb200.txt
+++ b/tests/test_suites/nightly_gb200.txt
@@ -26,6 +26,9 @@ tests/test_suites/llm/grpo-moonlight-16ba3b-4n4g-megatron.sh
 tests/test_suites/llm/grpo-nanov3-30ba3b-4n4g-megatron-qa-nvfp4-w4a16-real.sh
 tests/test_suites/llm/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.sh
 
+# Nemotron 3 Nano 30B colocated reshard (training/inference layout swap)
+tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.sh
+
 # Functional VLM run
 tests/test_suites/vlm/vlm_grpo-qwen2.5-vl-3b-instruct-clevr-1n4g-dtensor2tp1.v1.sh
 tests/test_suites/vlm/vlm_grpo-qwen2.5-vl-3b-instruct-clevr-1n4g-megatrontp1.v1.sh

From 6fa85ca28d62e788e47e579ef8e471a4e5488d91 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Tue, 11 Aug 2026 23:24:58 -0500
Subject: [PATCH 16/25] Minimize the 4n4g reshard recipe config

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 .../llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml  | 1 -
 1 file changed, 1 deletion(-)

diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml
index b24bdb06f0a..2417a3803e6 100644
--- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml
+++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml
@@ -1,5 +1,4 @@
 defaults: ../../grpo_math_1B.yaml
-# Sync (non-async) colocated GRPO with Megatron generation (Nemotron-3-Nano-30B-A3B), with reshard allowed.
 grpo:
   num_prompts_per_step: 2
   num_generations_per_prompt: 8

From c7d9592747804d8e09bef979c6e655d005b8fdaf Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Wed, 12 Aug 2026 06:41:58 -0500
Subject: [PATCH 17/25] Fix segment-size on nightlies

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 .../llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml  | 1 +
 .../llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.sh    | 1 +
 2 files changed, 2 insertions(+)

diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml
index 2417a3803e6..cfc218cc9d1 100644
--- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml
+++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml
@@ -45,3 +45,4 @@ logger:
 cluster:
   gpus_per_node: 4
   num_nodes: 4
+  segment_size: 2
diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.sh
index d0868cf7cda..77974b3e33d 100755
--- a/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.sh
+++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.sh
@@ -5,6 +5,7 @@ source $SCRIPT_DIR/common.env
 # ===== BEGIN CONFIG =====
 NUM_NODES=4
 GPUS_PER_NODE=4
+SEGMENT_SIZE=2
 STEPS_PER_RUN=10
 MAX_STEPS=10
 NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN ))  # Round up

From 7695ba53fdb448dbeea508418e70767ff27cc8ae Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Wed, 12 Aug 2026 14:44:44 -0500
Subject: [PATCH 18/25] Address reviewer comments

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 nemo_rl/models/generation/megatron/megatron_worker.py | 4 +++-
 nemo_rl/models/megatron/setup.py                      | 8 +++++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py
index b392f6fa80e..756430d9537 100644
--- a/nemo_rl/models/generation/megatron/megatron_worker.py
+++ b/nemo_rl/models/generation/megatron/megatron_worker.py
@@ -421,7 +421,9 @@ def prepare_for_generation(self, tags=None, **kwargs) -> None:
             self.model = self.move_model(
                 self.model, "cuda", move_params=True, move_grads=False
             )
-            # Forward pre-hooks are not safe with cross-DP inference.
+            # Because DP inference requests are asynchronously scheduled per rank, pre-forward hooks that trigger DP collectives (such as an overlapped param gather after optimizer steps) will stall or hang.
+            # Instead, synchronously gather all model compute weights from the sharded model state here, and deactivate all pre-forward hooks.
+            # Incompatible with FSDP2 or Megatron-FSDP for inference.
             if (
                 self.should_disable_forward_pre_hook
                 and self._forward_pre_hook_enabled()
diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py
index 0a10faefd95..4e39cbebf8a 100644
--- a/nemo_rl/models/megatron/setup.py
+++ b/nemo_rl/models/megatron/setup.py
@@ -1524,6 +1524,12 @@ def build_inference_model(
     Returns:
         The inference model module (single element; not DDP-wrapped, no optimizer).
     """
+    if megatron_cfg.dist.use_torch_fsdp2:
+        raise ValueError(
+            "A dedicated inference model (reshard) is not supported with use_torch_fsdp2 training: "
+            "DP inference disables the training model's forward pre-hooks, "
+            "which requires Megatron-core DistributedDataParallel."
+        )
     # Derive the inference provider from the initial snapshot taken by setup_model_and_optimizer.
     inference_provider = megatron_cfg._initial_model_provider
     del megatron_cfg._initial_model_provider
@@ -1581,7 +1587,7 @@ def build_inference_model(
         inference_model = get_model(
             inference_provider,
             megatron_cfg.ddp,
-            use_torch_fsdp2=megatron_cfg.dist.use_torch_fsdp2,
+            use_torch_fsdp2=False,  # the inference model is never trained
             data_parallel_random_init=megatron_cfg.rng.data_parallel_random_init,
             mixed_precision_wrapper=mixed_precision_wrapper,
             pg_collection=inference_pg_collection,

From 5d7f2c4b3dcba3a159dbffdcc47578cbe4378d62 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Wed, 12 Aug 2026 22:29:22 -0500
Subject: [PATCH 19/25] Address review-pr-team comments

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 nemo_rl/models/generation/megatron/megatron_worker.py | 11 ++++++-----
 nemo_rl/models/megatron/setup.py                      |  4 ++--
 .../models/policy/workers/megatron_policy_worker.py   |  2 +-
 3 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py
index 756430d9537..b9b7bc80bb3 100644
--- a/nemo_rl/models/generation/megatron/megatron_worker.py
+++ b/nemo_rl/models/generation/megatron/megatron_worker.py
@@ -35,6 +35,7 @@
     prepare_swap_model_weights,
     swap_model_weights,
 )
+from megatron.core.transformer import MegatronModule
 from megatron.core.transformer.enums import InferenceCudaGraphScope
 from megatron.core.transformer.utils import toggle_cuda_graphs
 from megatron.core.utils import unwrap_model
@@ -74,7 +75,7 @@ class MegatronGenerationMixin:
     # (see MegatronPolicyWorkerImpl._maybe_build_colocated_inference_model).
     inference_model = None
 
-    def _gen_model(self):
+    def _gen_model(self) -> MegatronModule:
         """The model the inference engine wraps.
 
         Returns the dedicated inference-layout model when one exists (colocated
@@ -871,7 +872,7 @@ def swap_weights_via_reshard(self, is_source: bool) -> bool:
 
     def _onload_inference_model(self) -> None:
         """Restore the colocated inference weights to GPU before resharding / generation."""
-        if not getattr(self, "_inference_model_offloaded", False):
+        if not self._inference_model_offloaded:
             return
         resume_inference_weights()
         self._inference_model_offloaded = False
@@ -879,8 +880,8 @@ def _onload_inference_model(self) -> None:
     def _offload_inference_model(self) -> None:
         """Offload the colocated inference weights to CPU while training runs."""
         if (
-            getattr(self, "inference_model", None) is None
-            or getattr(self, "_inference_model_offloaded", False)
+            self.inference_model is None
+            or self._inference_model_offloaded
             or not HAVE_TORCH_MEMORY_SAVER
         ):
             return
@@ -889,7 +890,7 @@ def _offload_inference_model(self) -> None:
 
     def _reshard_into_inference_model(self) -> None:
         """Reshard current training weights into the colocated inference-layout model."""
-        inference_model = getattr(self, "inference_model", None)
+        inference_model = self.inference_model
         if inference_model is None:
             return
 
diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py
index 4e39cbebf8a..1cc2ef15004 100644
--- a/nemo_rl/models/megatron/setup.py
+++ b/nemo_rl/models/megatron/setup.py
@@ -1559,8 +1559,8 @@ def build_inference_model(
     inference_provider.recompute_num_layers = None
     if inference_provider.transformer_impl == "inference_optimized":
         inference_provider.moe_pad_experts_for_cuda_graph_inference = False
-    # Need to run __post__init__ with the correct config.
-    TransformerConfig.__post_init__(inference_provider)
+    # Re-run the deferred MCore post-init (virtual, idempotent).
+    inference_provider.finalize()
 
     world_size = torch.distributed.get_world_size()
     inference_pg_collection = build_inference_pg_collection(
diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py
index 0a3528f41a6..362a01ea8cc 100644
--- a/nemo_rl/models/policy/workers/megatron_policy_worker.py
+++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py
@@ -2756,7 +2756,7 @@ def prepare_for_lp_inference(self):
         gc.collect()
         torch.cuda.empty_cache()
 
-    def _maybe_build_colocated_inference_model(self, config) -> None:
+    def _maybe_build_colocated_inference_model(self, config: PolicyConfig) -> None:
         """Build a separate inference-layout model when the colocated layout differs."""
         # Resolve the inference layout the same way the non-colocated generation policy does:
         # overlay the sparse mcore_generation_config onto a copy of megatron_cfg.

From 4c3e03c608ae3ae160c22780ef2abcaa76c1fe20 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Fri, 14 Aug 2026 05:45:57 -0500
Subject: [PATCH 20/25] Import torch_memory_saver correctly

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 nemo_rl/models/megatron/memory_saver.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/nemo_rl/models/megatron/memory_saver.py b/nemo_rl/models/megatron/memory_saver.py
index 9eb07e9110b..e1d900d22b2 100644
--- a/nemo_rl/models/megatron/memory_saver.py
+++ b/nemo_rl/models/megatron/memory_saver.py
@@ -17,7 +17,11 @@
 from typing import ContextManager
 
 try:
-    import torch_memory_saver  # pyrefly: ignore[import-error]
+    from torch_memory_saver import (  # pyrefly: ignore[import-error]
+        torch_memory_saver,
+    )
+
+    torch_memory_saver.hook_mode = "torch"
 
     HAVE_TORCH_MEMORY_SAVER = True
 except ImportError:

From 3ad74f79f415aab3387867e7adf14f24875c0f2a Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Mon, 17 Aug 2026 12:22:20 -0500
Subject: [PATCH 21/25] lint rebase conflict

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 .../llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml  | 1 -
 1 file changed, 1 deletion(-)

diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml
index cfc218cc9d1..3c45dc1c8b3 100644
--- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml
+++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml
@@ -34,7 +34,6 @@ policy:
       tensor_model_parallel_size: 4
       expert_model_parallel_size: 4
       sequence_parallel: true
-      refit_backend: nccl
 logger:
   log_dir: logs/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard
   wandb_enabled: true

From 3ca38b768b5da1227195485f19da69d96d984ffb Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Mon, 17 Aug 2026 17:38:51 -0500
Subject: [PATCH 22/25] Address reviewer comments

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 nemo_rl/models/generation/megatron/config.py  | 41 +++++++++--
 .../generation/megatron/megatron_worker.py    | 13 ++--
 nemo_rl/models/megatron/config.py             | 17 +++++
 nemo_rl/models/megatron/setup.py              | 58 ++++++++++++----
 .../policy/workers/megatron_policy_worker.py  | 68 +++++--------------
 ...o_megatron_generation_colocated_reshard.sh |  8 +++
 6 files changed, 128 insertions(+), 77 deletions(-)

diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py
index e37bdb61fb2..4d2511523b1 100644
--- a/nemo_rl/models/generation/megatron/config.py
+++ b/nemo_rl/models/generation/megatron/config.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Literal, NotRequired, TypedDict, cast
+from typing import Any, Literal, NotRequired, Optional, TypedDict, cast
 
 from nemo_rl.models.generation.interfaces import GenerationConfig
 from nemo_rl.models.policy import PolicyConfig
@@ -73,14 +73,43 @@ class MCoreGenerationConfig(GenerationConfig):
 
 
 def merged_inference_megatron_cfg(policy_config: PolicyConfig) -> dict[str, Any]:
-    """The `megatron_cfg` a dedicated inference model runs with.
-
-    Overlays the sparse `mcore_generation_config` onto `megatron_cfg`,
-    intentionally overwriting any training-side config with inference-side config.
-    """
+    """The `megatron_cfg` a dedicated inference model runs with."""
     generation_config = cast(MCoreGenerationConfig, policy_config["generation"])
     return {
         **cast(dict[str, Any], policy_config["megatron_cfg"]),
         **(generation_config.get("mcore_generation_config") or {}),
         "activation_checkpointing": False,
     }
+
+
+def dedicated_inference_megatron_cfg(
+    policy_config: PolicyConfig,
+) -> Optional[dict[str, Any]]:
+    """The `megatron_cfg` for a dedicated colocated inference model, or None.
+
+    Colocated Megatron generation shares the training model unless the resolved
+    inference layout or `transformer_impl` differs from training; then the worker
+    builds a second model and reshards into it on every wake. Inference never
+    uses CP, so CP is pinned to 1 (CP>1 training therefore always differs).
+
+    Returns None when the resolved config matches training (dual-mode: generate
+    directly on the shared training model).
+    """
+    inference_mcfg = merged_inference_megatron_cfg(policy_config)
+    inference_mcfg["context_parallel_size"] = 1
+
+    train_mcfg = cast(dict[str, Any], policy_config["megatron_cfg"])
+    layout_keys = (
+        "tensor_model_parallel_size",
+        "pipeline_model_parallel_size",
+        "expert_model_parallel_size",
+        "expert_tensor_parallel_size",
+        "context_parallel_size",
+    )
+    layout_differs = any(inference_mcfg[k] != train_mcfg[k] for k in layout_keys)
+    impl_differs = inference_mcfg.get("transformer_impl") != train_mcfg.get(
+        "transformer_impl"
+    )
+    if not (layout_differs or impl_differs):
+        return None
+    return inference_mcfg
diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py
index b9b7bc80bb3..c2420ade784 100644
--- a/nemo_rl/models/generation/megatron/megatron_worker.py
+++ b/nemo_rl/models/generation/megatron/megatron_worker.py
@@ -72,8 +72,9 @@ class MegatronGenerationMixin:
     """
 
     # Colocated-reshard hosts assign the dedicated inference-layout model here
-    # (see MegatronPolicyWorkerImpl._maybe_build_colocated_inference_model).
+    # (see MegatronPolicyWorkerImpl._build_colocated_inference_model).
     inference_model = None
+    _colocated_reshard_plan = None
 
     def _gen_model(self) -> MegatronModule:
         """The model the inference engine wraps.
@@ -409,12 +410,8 @@ def prepare_for_generation(self, tags=None, **kwargs) -> None:
         mcore_generation_config = self.cfg["generation"]["mcore_generation_config"]
 
         # Colocated reshard: build the dedicated inference-layout model on the first cycle.
-        if (
-            getattr(self, "_colocated_reshard_eligible", False)
-            and not self._colocated_inference_model_checked
-        ):
-            self._maybe_build_colocated_inference_model(self.cfg)
-            self._colocated_inference_model_checked = True
+        if self._colocated_reshard_plan is not None:
+            self._build_colocated_inference_model(self.cfg)
 
         gen_model = self._gen_model()
         gen_model.config.flash_decode = False
@@ -903,7 +900,7 @@ def _reshard_into_inference_model(self) -> None:
         torch.cuda.synchronize()
 
         # The swap reads the training params as its source;
-        # under overlap_param_gather they stay stale after the optimizer step until gathered,
+        # under overlap_param_gather they stay stale after the optimizer step until gathered.
         if self.should_disable_forward_pre_hook and self._forward_pre_hook_enabled():
             self._disable_forward_pre_hook_until_next_train_step(param_sync=True)
 
diff --git a/nemo_rl/models/megatron/config.py b/nemo_rl/models/megatron/config.py
index 06a8a3b3df1..5cf32e9b09a 100644
--- a/nemo_rl/models/megatron/config.py
+++ b/nemo_rl/models/megatron/config.py
@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from dataclasses import dataclass
 from typing import Any, Callable, NamedTuple, Optional
 
 import torch
+from megatron.bridge.models.model_provider import ModelProviderMixin
 from megatron.bridge.training.config import ConfigContainer
 from megatron.bridge.training.state import GlobalState
 from megatron.core.optimizer import MegatronOptimizer
@@ -43,6 +45,20 @@ class RuntimeConfig(NamedTuple):
     final_padded_vocab_size: int
 
 
+@dataclass
+class ColocatedReshardPlan:
+    """Setup-time plan for building the dedicated colocated inference model.
+
+    Produced by `setup_model_and_optimizer` when the inference and training layouts differ.
+    Consumed by the first `prepare_for_generation` call.
+    """
+
+    # Pre-wrap provider snapshot; build_inference_model mutates it to the inference layout.
+    initial_model_provider: ModelProviderMixin
+    # The resolved megatron_cfg the inference model runs with.
+    inference_megatron_cfg: dict[str, Any]
+
+
 ## returned from setup_model_and_optimizer
 class ModelAndOptimizerState(NamedTuple):
     """Container for model and optimizer state.
@@ -58,3 +74,4 @@ class ModelAndOptimizerState(NamedTuple):
     checkpointing_context: dict[str, Any]
     param_sync_func: Optional[Callable]
     draft_model: Optional[MegatronModule] = None
+    colocated_reshard_plan: Optional[ColocatedReshardPlan] = None
diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py
index 1cc2ef15004..866242f9e1a 100644
--- a/nemo_rl/models/megatron/setup.py
+++ b/nemo_rl/models/megatron/setup.py
@@ -25,7 +25,7 @@
 
 import torch
 from megatron.bridge import AutoBridge
-from megatron.bridge.models.model_provider import get_model
+from megatron.bridge.models.model_provider import ModelProviderMixin, get_model
 from megatron.bridge.peft.lora import LoRA
 from megatron.bridge.training import fault_tolerance
 from megatron.bridge.training.checkpointing import (
@@ -225,11 +225,18 @@ def _sync_distrib_opt(distrib_opt):
 
 from nemo_rl.algorithms.logits_sampling_utils import TrainingSamplingParams
 from nemo_rl.distributed.named_sharding import NamedSharding
+from nemo_rl.models.generation.megatron.config import (
+    dedicated_inference_megatron_cfg,
+)
 from nemo_rl.models.megatron.community_import import (
     import_model_from_hf_name,
     iter_vlm_config_overrides,
 )
-from nemo_rl.models.megatron.config import ModelAndOptimizerState, RuntimeConfig
+from nemo_rl.models.megatron.config import (
+    ColocatedReshardPlan,
+    ModelAndOptimizerState,
+    RuntimeConfig,
+)
 from nemo_rl.models.megatron.draft.utils import (
     build_draft_model,
     find_draft_owner_chunk,
@@ -1512,6 +1519,7 @@ def main_thread_only_enter(self):
 def build_inference_model(
     policy_cfg: PolicyConfig,
     megatron_cfg: ConfigContainer,
+    initial_model_provider: ModelProviderMixin,
 ) -> MegatronModule:
     """Build a second, inference-layout model for colocated Megatron refit.
 
@@ -1520,19 +1528,12 @@ def build_inference_model(
     Args:
         policy_cfg: The inference config
         megatron_cfg: The training config
+        initial_model_provider: Pre-wrap provider snapshot taken by `setup_model_and_optimizer`.
 
     Returns:
         The inference model module (single element; not DDP-wrapped, no optimizer).
     """
-    if megatron_cfg.dist.use_torch_fsdp2:
-        raise ValueError(
-            "A dedicated inference model (reshard) is not supported with use_torch_fsdp2 training: "
-            "DP inference disables the training model's forward pre-hooks, "
-            "which requires Megatron-core DistributedDataParallel."
-        )
-    # Derive the inference provider from the initial snapshot taken by setup_model_and_optimizer.
-    inference_provider = megatron_cfg._initial_model_provider
-    del megatron_cfg._initial_model_provider
+    inference_provider = initial_model_provider
     train_pipeline_model_parallel_size = inference_provider.pipeline_model_parallel_size
     _apply_parallelism_config(inference_provider, policy_cfg)
     _apply_moe_config(inference_provider, policy_cfg)
@@ -1744,13 +1745,43 @@ def freeze_moe_router(megatron_model):
     megatron_cfg.peft = peft
 
     # Snapshot the provider before any runtime state is added onto it.
+    colocated_reshard_plan: Optional[ColocatedReshardPlan] = None
     generation_cfg = policy_cfg.get("generation")
     if (
-        generation_cfg is not None
+        load_optimizer
+        and generation_cfg is not None
         and generation_cfg.get("backend") == "megatron"
         and generation_cfg.get("colocated", {}).get("enabled", False)
     ):
-        megatron_cfg._initial_model_provider = copy.deepcopy(megatron_cfg.model)
+        inference_megatron_cfg = dedicated_inference_megatron_cfg(policy_cfg)
+    else:
+        inference_megatron_cfg = None
+    if inference_megatron_cfg is not None:
+        if megatron_cfg.dist.use_torch_fsdp2:
+            raise ValueError(
+                "MCore colocated reshard is not supported with use_torch_fsdp2 training: "
+                "DP inference disables the training model's forward pre-hooks, "
+                "which requires Megatron-core DistributedDataParallel."
+            )
+        vpp_size = megatron_cfg.model.virtual_pipeline_model_parallel_size
+        if vpp_size not in (None, 1):
+            raise NotImplementedError(
+                "MCore colocated reshard is not supported with virtual pipeline parallelism > 1. "
+                f"(virtual_pipeline_model_parallel_size={vpp_size})"
+            )
+        if peft is not None:
+            raise NotImplementedError(
+                "MCore colocated reshard is not supported with PEFT."
+            )
+        draft_cfg = policy_cfg.get("draft")
+        if draft_cfg is not None and draft_cfg.get("enabled"):
+            raise NotImplementedError(
+                "MCore colocated reshard is not supported with draft models."
+            )
+        colocated_reshard_plan = ColocatedReshardPlan(
+            initial_model_provider=copy.deepcopy(megatron_cfg.model),
+            inference_megatron_cfg=copy.deepcopy(inference_megatron_cfg),
+        )
 
     if megatron_cfg.peft is not None:
         pre_peft_hook = _create_peft_pre_wrap_hook(megatron_cfg, state)
@@ -1874,6 +1905,7 @@ def composed_peft_hook(model: list[MegatronModule]) -> list[MegatronModule]:
         checkpointing_context,
         param_sync_func,
         draft_model=draft_model,
+        colocated_reshard_plan=colocated_reshard_plan,
     )
 
 
diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py
index 362a01ea8cc..14882a1670b 100644
--- a/nemo_rl/models/policy/workers/megatron_policy_worker.py
+++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py
@@ -54,9 +54,6 @@
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.distributed.named_sharding import NamedSharding
 from nemo_rl.models.generation.interfaces import GenerationDatumSpec
-from nemo_rl.models.generation.megatron.config import (
-    merged_inference_megatron_cfg,
-)
 from nemo_rl.models.generation.megatron.megatron_worker import (
     MegatronGenerationMixin,
     MegatronGenerationRefitMixin,
@@ -542,6 +539,7 @@ def __init__(
         self.checkpointing_context = model_and_optimizer_state.checkpointing_context
         param_sync_func = model_and_optimizer_state.param_sync_func
         self.draft_model = model_and_optimizer_state.draft_model
+        self._colocated_reshard_plan = model_and_optimizer_state.colocated_reshard_plan
         log_gpu_memory_diagnostics(
             label="after_model_setup", worker_type="MegatronPolicyWorker"
         )
@@ -626,15 +624,6 @@ def __init__(
         self.inference_model = None
         self._colocated_reshard_plan_ready = False
         self._inference_model_offloaded = False
-        self._colocated_inference_model_checked = False
-        gen_cfg = config.get("generation")
-        # The build itself is deferred to the first prepare_for_generation.
-        self._colocated_reshard_eligible = (
-            init_optimizer
-            and self.is_generation_colocated
-            and gen_cfg is not None
-            and gen_cfg.get("backend") == "megatron"
-        )
 
         # vars used for refit
         ## will be initialized in prepare_refit_info
@@ -2756,50 +2745,29 @@ def prepare_for_lp_inference(self):
         gc.collect()
         torch.cuda.empty_cache()
 
-    def _maybe_build_colocated_inference_model(self, config: PolicyConfig) -> None:
-        """Build a separate inference-layout model when the colocated layout differs."""
-        # Resolve the inference layout the same way the non-colocated generation policy does:
-        # overlay the sparse mcore_generation_config onto a copy of megatron_cfg.
+    def _build_colocated_inference_model(self, config: PolicyConfig) -> None:
+        """Build the dedicated inference-layout model planned at setup."""
+        plan = self._colocated_reshard_plan
+        inference_mcfg = plan.inference_megatron_cfg
         inference_config = copy.deepcopy(config)
-        inference_config["megatron_cfg"] = merged_inference_megatron_cfg(
-            inference_config
-        )
-        # Inference never uses CP: pin CP=1, so CP>1 training builds a separate inference model.
-        inference_config["megatron_cfg"]["context_parallel_size"] = 1
-
-        train_mcfg = config["megatron_cfg"]
-        inf_mcfg = inference_config["megatron_cfg"]
-        layout_keys = (
-            "tensor_model_parallel_size",
-            "pipeline_model_parallel_size",
-            "expert_model_parallel_size",
-            "expert_tensor_parallel_size",
-            "context_parallel_size",
-        )
-        layout_differs = any(inf_mcfg[k] != train_mcfg[k] for k in layout_keys)
-        impl_differs = inf_mcfg.get("transformer_impl") != train_mcfg.get(
-            "transformer_impl"
-        )
-        if not (layout_differs or impl_differs):
-            return
+        inference_config["megatron_cfg"] = inference_mcfg
 
-        peft_cfg = train_mcfg.get("peft")
-        if peft_cfg is not None and peft_cfg.get("enabled"):
-            raise NotImplementedError(
-                "Colocated generation with a differing inference parallel layout is not "
-                "supported with PEFT. Use a matched layout or non-colocated generation."
-            )
-        draft_cfg = config.get("draft")
-        if draft_cfg is not None and draft_cfg.get("enabled"):
-            raise NotImplementedError(
-                "Colocated generation with a differing inference parallel layout is not "
-                "supported with a speculative draft model."
-            )
+        print(
+            "[colocated-reshard] building dedicated inference model "
+            f"(inference TP={inference_mcfg['tensor_model_parallel_size']} "
+            f"PP={inference_mcfg['pipeline_model_parallel_size']} "
+            f"EP={inference_mcfg['expert_model_parallel_size']} "
+            f"CP={inference_mcfg['context_parallel_size']} "
+            f"impl={inference_mcfg.get('transformer_impl')})",
+            flush=True,
+        )
         # Built inside the first prepare_for_generation, immediately before the
         # reshard needs it resident; finish_generation offloads it afterwards.
         self.inference_model = build_inference_model(
-            inference_config, self.megatron_cfg
+            inference_config, self.megatron_cfg, plan.initial_model_provider
         )
+        # The plan is consumed (provider mutated to the inference layout); release it.
+        self._colocated_reshard_plan = None
 
     def prepare_for_training(self, *args, **kwargs):
         # onload models and optimizer state to cuda
diff --git a/tests/functional/grpo_megatron_generation_colocated_reshard.sh b/tests/functional/grpo_megatron_generation_colocated_reshard.sh
index 9aea0df5c39..9af8da7f47d 100755
--- a/tests/functional/grpo_megatron_generation_colocated_reshard.sh
+++ b/tests/functional/grpo_megatron_generation_colocated_reshard.sh
@@ -50,3 +50,11 @@ uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
 
 uv run tests/check_metrics.py $JSON_METRICS \
     'max(data["train/token_mult_prob_error"]) < 1.05'
+
+# Check that colocated reshard actually took place.
+# A matched layout that requires no reshard would pass all tests.
+# While the configs we are passing guarantee reshard, this is still a useful sanity check.
+if ! grep -q "\[colocated-reshard\] building dedicated inference model" $RUN_LOG; then
+    echo "FAIL: colocated-reshard marker not found; dedicated inference model was never built"
+    exit 1
+fi

From 6f6312073380ec844afe1ad4674e855d4f2cd1b9 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Mon, 17 Aug 2026 17:46:21 -0500
Subject: [PATCH 23/25] Address reviewer comment

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 nemo_rl/models/generation/megatron/megatron_worker.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py
index c2420ade784..093c68c10e2 100644
--- a/nemo_rl/models/generation/megatron/megatron_worker.py
+++ b/nemo_rl/models/generation/megatron/megatron_worker.py
@@ -414,6 +414,8 @@ def prepare_for_generation(self, tags=None, **kwargs) -> None:
             self._build_colocated_inference_model(self.cfg)
 
         gen_model = self._gen_model()
+        # `flash_decode` selects Megatron Inference's deprecated static-batching decode path,
+        # which would cause an assertion error if taken.
         gen_model.config.flash_decode = False
         if self.is_generation_colocated and self.inference_model is None:
             self.model = self.move_model(

From 4a21615e92e85c3e378cac9fbae46f437e7bc876 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Mon, 17 Aug 2026 18:08:35 -0500
Subject: [PATCH 24/25] Fix default value of argument

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 examples/configs/distillation_math.yaml                | 1 -
 examples/configs/grpo_math_1B.yaml                     | 3 +--
 examples/configs/ppo_math_1B.yaml                      | 1 -
 nemo_rl/models/megatron/setup.py                       | 2 +-
 tests/unit/models/megatron/test_megatron_setup.py      | 2 +-
 tests/unit/reference_configs/distillation_math.yaml    | 1 -
 tests/unit/reference_configs/grpo_math_1B.yaml         | 3 +--
 tests/unit/reference_configs/ppo_math_1B_megatron.yaml | 1 -
 8 files changed, 4 insertions(+), 10 deletions(-)

diff --git a/examples/configs/distillation_math.yaml b/examples/configs/distillation_math.yaml
index eff54539827..59f7616e0e1 100644
--- a/examples/configs/distillation_math.yaml
+++ b/examples/configs/distillation_math.yaml
@@ -44,7 +44,6 @@ policy: &POLICY_BASE
     logprob_chunk_size: null
 
     offload_optimizer_for_logprob: false
-    offload_optimizer_for_refit: true # Offload the optimizer state to CPU around each colocated-generation refit; configs that do not inherit this default keep the optimizer on GPU (needs the extra headroom)
 
     dtensor_cfg: &DTENSOR_BASE
         enabled: true
diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml
index 05825750e37..62e94d7b0b4 100644
--- a/examples/configs/grpo_math_1B.yaml
+++ b/examples/configs/grpo_math_1B.yaml
@@ -136,8 +136,7 @@ policy:
   max_total_sequence_length: 512
   precision: "bfloat16"
   logprob_chunk_size: null
-  offload_optimizer_for_logprob: false # Only useful for non-colocated generation since colocated generation offloads the optimizer before refit (see offload_optimizer_for_refit)
-  offload_optimizer_for_refit: true # Offload the optimizer state to CPU around each colocated-generation refit; configs that do not inherit this default keep the optimizer on GPU (needs the extra headroom)
+  offload_optimizer_for_logprob: false # Only useful for non-colocated generation since colocated generation will always offload optimizer to cuda before refit
 
   dtensor_cfg:
     _v2: true
diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml
index e99ee8301f4..97c5f9688de 100644
--- a/examples/configs/ppo_math_1B.yaml
+++ b/examples/configs/ppo_math_1B.yaml
@@ -102,7 +102,6 @@ policy:
   precision: "bfloat16"
   logprob_chunk_size: null
   offload_optimizer_for_logprob: false
-  offload_optimizer_for_refit: true # Offload the optimizer state to CPU around each colocated-generation refit; configs that do not inherit this default keep the optimizer on GPU (needs the extra headroom)
 
   dtensor_cfg:
     _v2: true
diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py
index 866242f9e1a..fe940e066d6 100644
--- a/nemo_rl/models/megatron/setup.py
+++ b/nemo_rl/models/megatron/setup.py
@@ -359,7 +359,7 @@ def validate_and_set_config(
     # Optimizer configuration
     optimizer_cpu_offload = config["megatron_cfg"]["optimizer"]["optimizer_cpu_offload"]
     offload_optimizer_for_logprob = config["offload_optimizer_for_logprob"]
-    offload_optimizer_for_refit = bool(config.get("offload_optimizer_for_refit"))
+    offload_optimizer_for_refit = bool(config.get("offload_optimizer_for_refit", True))
 
     # Reward models are not yet supported with Megatron.
     if "reward_model_cfg" in config and config["reward_model_cfg"]["enabled"]:
diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py
index 8b7d2a9375b..afc01daddc2 100644
--- a/tests/unit/models/megatron/test_megatron_setup.py
+++ b/tests/unit/models/megatron/test_megatron_setup.py
@@ -2150,7 +2150,7 @@ def test_generation_colocation_detection(self):
                 )
 
                 assert runtime_config.is_generation_colocated is True
-                assert runtime_config.offload_optimizer_for_refit is False
+                assert runtime_config.offload_optimizer_for_refit is True
 
 
 @pytest.mark.mcore
diff --git a/tests/unit/reference_configs/distillation_math.yaml b/tests/unit/reference_configs/distillation_math.yaml
index 4ac6b56b997..5f5f2d8ddfb 100644
--- a/tests/unit/reference_configs/distillation_math.yaml
+++ b/tests/unit/reference_configs/distillation_math.yaml
@@ -44,7 +44,6 @@ policy: &POLICY_BASE
     logprob_chunk_size: null
 
     offload_optimizer_for_logprob: false
-    offload_optimizer_for_refit: true # Offload the optimizer state to CPU around each colocated-generation refit; configs that do not inherit this default keep the optimizer on GPU (needs the extra headroom)
 
     dtensor_cfg: &DTENSOR_BASE
         enabled: true
diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml
index 22c602f5dfe..d5738e8c762 100644
--- a/tests/unit/reference_configs/grpo_math_1B.yaml
+++ b/tests/unit/reference_configs/grpo_math_1B.yaml
@@ -142,8 +142,7 @@ policy:
   max_total_sequence_length: 512
   precision: "bfloat16"
   logprob_chunk_size: null
-  offload_optimizer_for_logprob: false # Only useful for non-colocated generation since colocated generation offloads the optimizer before refit (see offload_optimizer_for_refit)
-  offload_optimizer_for_refit: true # Offload the optimizer state to CPU around each colocated-generation refit; configs that do not inherit this default keep the optimizer on GPU (needs the extra headroom)
+  offload_optimizer_for_logprob: false # Only useful for non-colocated generation since colocated generation will always offload optimizer to cuda before refit
 
   dtensor_cfg:
     _v2: true
diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml
index 2d6620c4ef1..9858962d779 100644
--- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml
+++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml
@@ -101,7 +101,6 @@ policy:
   precision: "bfloat16"
   logprob_chunk_size: null
   offload_optimizer_for_logprob: false
-  offload_optimizer_for_refit: true # Offload the optimizer state to CPU around each colocated-generation refit; configs that do not inherit this default keep the optimizer on GPU (needs the extra headroom)
 
   dtensor_cfg:
     _v2: true

From 5f1f29470a8a71d3c7dfebb184d74309e0697774 Mon Sep 17 00:00:00 2001
From: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Date: Mon, 17 Aug 2026 18:32:21 -0500
Subject: [PATCH 25/25] Fix extra issues

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
---
 .../megatron/megatron_generation.py           | 22 ++++++++++++++-----
 .../generation/megatron/megatron_worker.py    |  4 ++--
 nemo_rl/models/megatron/setup.py              |  3 +--
 .../policy/workers/megatron_policy_worker.py  |  5 ++---
 4 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/nemo_rl/models/generation/megatron/megatron_generation.py b/nemo_rl/models/generation/megatron/megatron_generation.py
index 4f59b1ee1da..d21c40bb52a 100644
--- a/nemo_rl/models/generation/megatron/megatron_generation.py
+++ b/nemo_rl/models/generation/megatron/megatron_generation.py
@@ -27,6 +27,7 @@
 )
 from nemo_rl.models.generation.megatron.config import (
     MCoreGenerationConfig,
+    dedicated_inference_megatron_cfg,
     merged_inference_megatron_cfg,
 )
 from nemo_rl.models.policy import PolicyConfig
@@ -53,13 +54,22 @@ def effective_megatron_cfg(config: PolicyConfig) -> dict[str, Any]:
 
     @classmethod
     def nvlink_domain_span(cls, config: PolicyConfig) -> int:
-        """Largest GPU group requiring full NVLink connectivity."""
-        megatron_cfg = cls.effective_megatron_cfg(config)
+        """Largest GPU group requiring full NVLink connectivity.
+
+        Colocated reshard hosts a second, inference-layout model on the same ranks.
+        """
+        layouts = [cls.effective_megatron_cfg(config)]
+        if config["generation"]["colocated"]["enabled"]:
+            inference_mcfg = dedicated_inference_megatron_cfg(config)
+            if inference_mcfg is not None:
+                layouts.append(inference_mcfg)
         return max(
-            megatron_cfg["tensor_model_parallel_size"]
-            * megatron_cfg["context_parallel_size"],
-            megatron_cfg.get("expert_tensor_parallel_size", 1)
-            * megatron_cfg.get("expert_model_parallel_size", 1),
+            max(
+                mcfg["tensor_model_parallel_size"] * mcfg["context_parallel_size"],
+                mcfg.get("expert_tensor_parallel_size", 1)
+                * mcfg.get("expert_model_parallel_size", 1),
+            )
+            for mcfg in layouts
         )
 
     @classmethod
diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py
index 093c68c10e2..2bbd1014a84 100644
--- a/nemo_rl/models/generation/megatron/megatron_worker.py
+++ b/nemo_rl/models/generation/megatron/megatron_worker.py
@@ -907,7 +907,7 @@ def _reshard_into_inference_model(self) -> None:
             self._disable_forward_pre_hook_until_next_train_step(param_sync=True)
 
         # Build + cache the same-rank reshard plan once, before the first CUDA-graph capture.
-        if not self._colocated_reshard_plan_ready:
+        if not self._swap_weights_plan_prepared:
             prepare_swap_model_weights(
                 src_model=self.model,
                 target_model=inference_model,
@@ -915,7 +915,7 @@ def _reshard_into_inference_model(self) -> None:
                 src_rank_offset=0,
                 dst_rank_offset=0,
             )
-            self._colocated_reshard_plan_ready = True
+            self._swap_weights_plan_prepared = True
 
         swap_model_weights(
             self.model,
diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py
index fe940e066d6..25003ce5e5d 100644
--- a/nemo_rl/models/megatron/setup.py
+++ b/nemo_rl/models/megatron/setup.py
@@ -1773,8 +1773,7 @@ def freeze_moe_router(megatron_model):
             raise NotImplementedError(
                 "MCore colocated reshard is not supported with PEFT."
             )
-        draft_cfg = policy_cfg.get("draft")
-        if draft_cfg is not None and draft_cfg.get("enabled"):
+        if draft_enabled:
             raise NotImplementedError(
                 "MCore colocated reshard is not supported with draft models."
             )
diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py
index 14882a1670b..b6ae9a23a95 100644
--- a/nemo_rl/models/policy/workers/megatron_policy_worker.py
+++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py
@@ -622,7 +622,7 @@ def __init__(
 
         # Colocated reshard: build a dedicated inference-layout model container.
         self.inference_model = None
-        self._colocated_reshard_plan_ready = False
+        self._swap_weights_plan_prepared = False
         self._inference_model_offloaded = False
 
         # vars used for refit
@@ -2749,8 +2749,7 @@ def _build_colocated_inference_model(self, config: PolicyConfig) -> None:
         """Build the dedicated inference-layout model planned at setup."""
         plan = self._colocated_reshard_plan
         inference_mcfg = plan.inference_megatron_cfg
-        inference_config = copy.deepcopy(config)
-        inference_config["megatron_cfg"] = inference_mcfg
+        inference_config = {**config, "megatron_cfg": inference_mcfg}
 
         print(
             "[colocated-reshard] building dedicated inference model "