feat: colocated Megatron reshard

examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.yaml

@@ -0,0 +1,47 @@
+defaults: ../../grpo_math_1B.yaml
+grpo:
+  num_prompts_per_step: 2
+  num_generations_per_prompt: 8
+  max_num_steps: 500
+checkpointing:
+  enabled: false
+  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
+  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
+      moe_router_dtype: fp32
+      activation_checkpointing: false
+      tensor_model_parallel_size: 4
+      expert_model_parallel_size: 4
+      sequence_parallel: true
+logger:
+  log_dir: logs/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard
+  wandb_enabled: true
+  tensorboard_enabled: true
+  wandb:
+    project: nemo-rl
+    name: grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard
+cluster:
+  gpus_per_node: 4
+  num_nodes: 4
+  segment_size: 2

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(
@@ -1452,8 +1445,25 @@ 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()
+        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()
+            setup_timing_metrics.generation_init_load_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 +1477,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 +2322,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 +2377,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 +2396,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 {}
 
 

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, Optional, TypedDict, cast
 
 from nemo_rl.models.generation.interfaces import GenerationConfig
+from nemo_rl.models.policy import PolicyConfig
 
 
 class MCoreGenerationSpecificArgs(TypedDict):
@@ -69,3 +70,46 @@ class MCoreGenerationConfig(GenerationConfig):
     """Generation config for Megatron Inference."""
 
     mcore_generation_config: MCoreGenerationSpecificArgs
+
+
+def merged_inference_megatron_cfg(policy_config: PolicyConfig) -> dict[str, Any]:
+    """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

nemo_rl/models/generation/megatron/megatron_generation.py

@@ -25,8 +25,13 @@
     GenerationInterface,
     GenerationOutputSpec,
 )
-from nemo_rl.models.generation.megatron.config import MCoreGenerationConfig
+from nemo_rl.models.generation.megatron.config import (
+    MCoreGenerationConfig,
+    dedicated_inference_megatron_cfg,
+    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
@@ -43,23 +48,28 @@ 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:
-        """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
@@ -121,6 +131,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.
@@ -137,8 +149,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(