fix: Reconcile #2315 and #2612

examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym.yaml

@@ -15,8 +15,6 @@ policy:
   megatron_cfg:
     pipeline_model_parallel_size: 1
     context_parallel_size: 2
-  sequence_packing:
-    enabled: false
   generation:
     backend: megatron
     colocated:

examples/nemo_gym/grpo_nanov3.yaml

@@ -222,7 +222,6 @@ policy:
       inference_cuda_graph_scope: "block"
       activation_checkpointing: false
       mamba_inference_ssm_states_dtype: "float32"
-      inference_moe_token_dispatcher_type: "nccl"  # Fall back to NCCL for now
       inference_grouped_gemm_backend: "vllm"
       moe_router_num_groups: null  # InferenceTopKRouter requires num_groups=None
       moe_router_group_topk: null  # paired with moe_router_num_groups=null
@@ -241,7 +240,7 @@ policy:
       kv_cache_management_mode: "persist"  # KV cache lifecycle across suspend/resume. Options: "persist", "offload". To select "recompute", set grpo.async_grpo.recompute_kv_cache_after_weight_updates=true.
       materialize_only_last_token_logits: true
       num_speculative_tokens: 0
-      refit_backend: "nvshmem"  # Copy-service backend for non-colocated megatron weight refit. Options: "gloo" or "nvshmem".
+      refit_backend: "nccl"  # Copy-service backend for non-colocated megatron weight refit. Options: "gloo" or "nccl".
       async_engine: true
       expose_http_server: true
       enable_prefix_caching: true

nemo_rl/algorithms/grpo.py

@@ -703,6 +703,12 @@ def _spinup_nemo_gym(base_urls, model_name):
         )
         train_cluster = cluster
         inference_cluster = cluster
+        # Colocated generation reuses the policy's cluster; need to decide topology here.
+        if (
+            node_resource_constraints is not None
+            and generation_config["backend"] == "megatron"
+        ):
+            MegatronGeneration.init_cluster_placement_groups(cluster, policy_config)
         print(
             f"  ✓ Ray cluster for policy initialized with {policy_nodes} nodes",
             flush=True,
@@ -817,15 +823,21 @@ def _spinup_nemo_gym(base_urls, model_name):
                         flush=True,
                     )
 
-                # Inference topology: each vLLM/SGLang instance spans
+                # Inference topology: each inference instance spans
                 # nodes_per_instance nodes; keep those within one domain
                 # so cross-node all-reduce uses NVLink, not InfiniBand.
                 #
                 # For vLLM: total GPUs per instance = TP * PP (separate dimensions).
                 # For SGLang: gpus_per_server already includes all parallelism
                 #   dimensions (TP, DP-attention, PP are internal subdivisions),
                 #   so we use it directly without multiplying by pp_size.
-                if generation_config["backend"] == "vllm":
+                # For Megatron: the NVLink-domain span of the parallelism the
+                #   generation workers actually run with.
+                if generation_config["backend"] == "megatron":
+                    gpus_per_instance = MegatronGeneration.nvlink_domain_span(
+                        policy_config
+                    )
+                elif generation_config["backend"] == "vllm":
                     vllm_cfg = generation_config.get("vllm_cfg", {})
                     gpus_per_instance = vllm_cfg["tensor_parallel_size"] * vllm_cfg.get(
                         "pipeline_parallel_size", 1
@@ -903,13 +915,19 @@ def _spinup_nemo_gym(base_urls, model_name):
             node_resource_constraints=inference_node_resource_constraints,
         )
         if inference_node_resource_constraints is not None:
-            {
-                "vllm": VllmGeneration,
-                "trtllm": TrtllmGeneration,
-            }[generation_config["backend"]].init_cluster_placement_groups(
-                inference_cluster,
-                generation_config,
-            )
+            if generation_config["backend"] == "megatron":
+                # Megatron inference reuses the training parallelism config.
+                MegatronGeneration.init_cluster_placement_groups(
+                    inference_cluster, policy_config
+                )
+            else:
+                {
+                    "vllm": VllmGeneration,
+                    "trtllm": TrtllmGeneration,
+                }[generation_config["backend"]].init_cluster_placement_groups(
+                    inference_cluster,
+                    generation_config,
+                )
         print(
             f"  ✓ Ray inference cluster initialized with {inference_nodes} nodes with {inference_gpus_per_node} GPUs per node",
             flush=True,

nemo_rl/models/generation/megatron/megatron_generation.py

@@ -36,7 +36,35 @@ class MegatronGeneration(GenerationInterface):
     """Generation interface backed by Megatron (colocated or non-colocated)."""
 
     @staticmethod
+    def effective_megatron_cfg(config: PolicyConfig) -> dict[str, Any]:
+        """The megatron_cfg the generation workers actually run with.
+
+        Colocated generation shares the training model, so the training
+        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", {}),
+        }
+
+    @classmethod
+    def nvlink_domain_span(cls, config: PolicyConfig) -> int:
+        """Largest GPU group requiring full NVLink connectivity."""
+        megatron_cfg = cls.effective_megatron_cfg(config)
+        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),
+        )
+
+    @classmethod
     def init_cluster_placement_groups(
+        cls,
         cluster: RayVirtualCluster,
         config: PolicyConfig,
     ) -> None:
@@ -46,16 +74,10 @@ def init_cluster_placement_groups(
             cluster: The inference `RayVirtualCluster`.
             config: The full `PolicyConfig` (megatron parallelism + colocation).
         """
-        megatron_cfg = config["megatron_cfg"]
-        model_parallel_size = (
-            megatron_cfg["tensor_model_parallel_size"]
-            * megatron_cfg["pipeline_model_parallel_size"]
-            * megatron_cfg["context_parallel_size"]
-        )
         colocated = config["generation"]["colocated"]["enabled"]
         cluster._init_placement_groups(
             strategy=None if colocated else "PACK",
-            use_unified_pg=model_parallel_size > cluster.num_gpus_per_node,
+            use_unified_pg=cls.nvlink_domain_span(config) > cluster.num_gpus_per_node,
         )
 
     def __init__(
@@ -111,7 +133,10 @@ def __init__(
 
         # Stand up a dedicated inference-only policy.
         self._owns_policy = True
-        self._policy_config["megatron_cfg"].update(self.cfg["mcore_generation_config"])
+        self._policy_config = {
+            **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.

tests/functional/L1_Functional_Tests_Megatron_4.sh

@@ -50,6 +50,7 @@ megatron_generation_supported() {
 
 if megatron_generation_supported; then
     run_test      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_async.sh
     run_test fast uv run --no-sync bash ./tests/functional/grpo_megatron_generation_colocated_async.sh

tests/functional/grpo_megatron_generation_topology.sh

@@ -0,0 +1,71 @@
+#!/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
+
+cd $PROJECT_ROOT
+
+# cluster.segment_size only engages when Ray nodes carry nvlink_domain_* labels,
+# which ray.sub probes from `nvidia-smi -q` ClusterUUID on NVLink-fabric clusters
+# (e.g. GB200 NVL72); CI runners have none. Pre-start a Ray head with a synthetic
+# domain label so init_ray() attaches to it (externally managed cluster) and the
+# topology-aware megatron placement path runs for real.
+cleanup() {
+    uv run ray stop --force || true
+}
+trap cleanup EXIT
+uv run ray stop --force || true # don't attach to a stale cluster
+uv run ray start --head --disable-usage-stats \
+    --resources='{"nvlink_domain_ci_synthetic": 1, "topo_rank": 1}'
+
+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.generation.backend=megatron \
+    cluster.gpus_per_node=2 \
+    cluster.segment_size=1 \
+    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
+
+# Guard against the silent fallback: with no (or unreadable) domain labels the run
+# would succeed without ever exercising the topology placement path under test.
+grep -q "Topology-aware allocation" $RUN_LOG || {
+    echo "ERROR: topology-aware allocation did not engage (no segment selection logged)" >&2
+    exit 1
+}
+# NOTE: `! grep` is exempt from `set -e`, hence the explicit if.
+if grep -q "no NVLink domain info" $RUN_LOG; then
+    echo "ERROR: segment_size fell back to unordered allocation" >&2
+    exit 1
+fi
+
+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'

tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym.sh

@@ -8,11 +8,9 @@ GPUS_PER_NODE=8
 STEPS_PER_RUN=8
 MAX_STEPS=8
 NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN ))  # Round up
-# ~25 min startup (30B-MoE load + CUDA-graph warmup + nemo_gym servers) plus
-# ~63 min for 8 async steps left no headroom at 90 min, so the driver finished
-# but Slurm SIGKILLed teardown/metric-dump at the wall-clock limit (CI mislabels
-# the exit-137 as OOM). 120 min leaves margin for teardown + metrics.
-NUM_MINUTES=120
+# ~25 min startup (30B-MoE load + CUDA-graph warmup + nemo_gym servers) plus ~130 min for 8 steps
+# 180 min leaves margin for teardown + metric-dump within the 4 h job allocation.
+NUM_MINUTES=180
 # ===== END CONFIG =====
 
 exit_if_max_steps_reached

tests/unit/distributed/test_megatron_placement_groups.py

@@ -0,0 +1,86 @@
+# 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 unittest.mock import MagicMock
+
+import pytest
+
+from nemo_rl.models.generation.megatron import MegatronGeneration
+
+
+def _placement_policy_config(
+    *,
+    tp: int = 1,
+    pp: int = 1,
+    cp: int = 1,
+    ep: int = 1,
+    etp: int = 1,
+    colocated: bool = False,
+    mcore_overrides: dict | None = None,
+) -> dict:
+    """Minimal PolicyConfig slice consumed by init_cluster_placement_groups."""
+    return {
+        "megatron_cfg": {
+            "tensor_model_parallel_size": tp,
+            "pipeline_model_parallel_size": pp,
+            "context_parallel_size": cp,
+            "expert_model_parallel_size": ep,
+            "expert_tensor_parallel_size": etp,
+        },
+        "generation": {
+            "colocated": {"enabled": colocated},
+            "mcore_generation_config": mcore_overrides or {},
+        },
+    }
+
+
+@pytest.mark.parametrize(
+    "config_kwargs,expected_strategy,expected_unified",
+    [
+        # cross-node span via TP alone -> one unified PG
+        (dict(tp=8), "PACK", True),
+        # PP is excluded from the NVLink span: TP*CP=4 fits a node even
+        # though the full TP*PP*CP instance would not
+        (dict(tp=2, pp=2, cp=2), "PACK", False),
+        # node-local span at the == boundary -> per-node PGs
+        (dict(tp=4), "PACK", False),
+        # cross-node MoE expert group (ETP*EP > TP*CP) -> one unified PG:
+        # the NVLS dispatcher needs the ep_group fully NVLink-connected
+        (dict(tp=2, ep=8), "PACK", True),
+        # non-colocated generation parallelism overrides the training values
+        # (mirrors MegatronGeneration's megatron_cfg merge)
+        (
+            dict(tp=8, mcore_overrides={"tensor_model_parallel_size": 2}),
+            "PACK",
+            False,
+        ),
+        # colocated reuses the training layout (overrides do not apply):
+        # no PACK strategy, span from the training config incl. its EP
+        (dict(tp=2, ep=8, colocated=True), None, True),
+    ],
+)
+def test_megatron_init_cluster_placement_groups(
+    config_kwargs, expected_strategy, expected_unified
+):
+    """The NVLink-domain span is max(TP*CP, ETP*EP) of the operative config."""
+    cluster = MagicMock(num_gpus_per_node=4)
+
+    MegatronGeneration.init_cluster_placement_groups(
+        cluster, _placement_policy_config(**config_kwargs)
+    )
+
+    cluster._init_placement_groups.assert_called_once_with(
+        strategy=expected_strategy,
+        use_unified_pg=expected_unified,
+    )