fix(megatron): shard Megatron rollouts across DP coordinator leaders

nemo_rl/models/generation/megatron/megatron_generation.py

@@ -121,6 +121,7 @@ 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]] = []
+        self.current_generate_dp_shard_idx = 0
 
         if policy is not None:
             # Reuse the existing training policy.
@@ -189,14 +190,19 @@ def update_weights_from_collective(self) -> list[ray.ObjectRef]:
         """Receive updated weights from the training cluster via collective communication."""
         return self._policy.swap_weights_via_reshard(is_source=False)
 
+    @property
+    def sharding_annotations(self):
+        """Parallelism layout for the underlying Megatron policy workers."""
+        return self._policy.sharding_annotations
+
     def generate(
         self, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False
     ) -> BatchedDataDict[GenerationOutputSpec]:
         """Generate a batch of data using the Megatron generation backend.
 
-        mcore's data-parallel coordinator only accepts requests from DP rank 0 —
-        the other workers' engine loops drain the coordinator queue but never
-        receive a Python-side call. So we dispatch straight to worker 0.
+        Prompts are sharded across DP coordinator leaders (mirroring the vLLM
+        backend). Each leader submits its shard to its replica's mcore
+        coordinator; non-leader TP/PP ranks are not invoked.
 
         Args:
             data: BatchedDataDict containing input_ids and input_lengths.
@@ -205,25 +211,72 @@ def generate(
         Returns:
             BatchedDataDict conforming to GenerationOutputSpec.
         """
-        future = self._policy.worker_group.run_single_worker_single_data(
-            method_name="generate",
-            worker_idx=0,
-            data=data,
-            greedy=greedy,
+        assert isinstance(data, BatchedDataDict), (
+            f"data must be a BatchedDataDict, got type: {type(data)}"
+        )
+        assert "input_ids" in data and "input_lengths" in data, (
+            "input_ids and input_lengths are required in data for Megatron generation"
+        )
+
+        dp_size = self.sharding_annotations.get_axis_size("data_parallel")
+        sharded_data = data.shard_by_batch_size(dp_size, allow_uneven_shards=True)
+        future_bundle = self._policy.worker_group.run_all_workers_sharded_data(
+            "generate",
+            data=sharded_data,
+            in_sharded_axes=["data_parallel"],
+            replicate_on_axes=None,
+            output_is_replicated=None,
+            common_kwargs={"greedy": greedy},
         )
-        return ray.get(future)
+        results = self._policy.worker_group.get_all_worker_results(future_bundle)
+        combined: BatchedDataDict[GenerationOutputSpec] = BatchedDataDict.from_batches(
+            results,
+            pad_value_dict={"output_ids": self.cfg["_pad_token_id"]},
+        )
+
+        required_keys = [
+            "output_ids",
+            "generation_lengths",
+            "unpadded_sequence_lengths",
+            "logprobs",
+        ]
+        missing_keys = [key for key in required_keys if key not in combined]
+        if missing_keys:
+            raise ValueError(
+                f"Missing required keys for GenerationOutputSpec: {missing_keys}"
+            )
+
+        return combined
 
     async def generate_async(
         self, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False
     ) -> AsyncGenerator[tuple[int, BatchedDataDict[GenerationOutputSpec]], None]:
-        """Generate asynchronously, yielding `(index, batch)` tuples as they complete."""
-        worker = self._policy.worker_group.workers[0]
+        """Generate asynchronously, yielding `(index, batch)` tuples as they complete.
+
+        Single-sample requests are round-robin dispatched across DP coordinator
+        leaders, matching the vLLM async rollout path.
+        """
+        assert isinstance(data, BatchedDataDict), (
+            f"data must be a BatchedDataDict, got type: {type(data)}"
+        )
+        assert data.size == 1, (
+            "Megatron generate_async handles one sample per call; batch outside this method."
+        )
+
+        leader_worker_idx = self._policy.worker_group.get_dp_leader_worker_idx(
+            self.current_generate_dp_shard_idx
+        )
+        worker = self._policy.worker_group.workers[leader_worker_idx]
         futures = worker.generate_async.options(num_returns="streaming").remote(
             data=data, greedy=greedy
         )
+        self.current_generate_dp_shard_idx = (
+            self.current_generate_dp_shard_idx + 1
+        ) % self._policy.worker_group.dp_size
+
         async for result_ref in futures:
             index, result_batch = await result_ref
-            result_batch["gen_leader_worker_idx"] = [0]
+            result_batch["gen_leader_worker_idx"] = [leader_worker_idx]
             yield index, result_batch
 
     def prepare_for_generation(self, *args: Any, **kwargs: Any) -> bool:
@@ -240,9 +293,15 @@ def prepare_for_generation(self, *args: Any, **kwargs: Any) -> bool:
             not self.dp_openai_server_base_urls
             and self.cfg["mcore_generation_config"]["expose_http_server"]
         ):
-            url_futures = self._policy.worker_group.run_all_workers_single_data(
-                "report_dp_openai_server_base_url"
-            )
+            url_futures = []
+            for dp_idx in range(self._policy.worker_group.dp_size):
+                worker_idx = self._policy.worker_group.get_dp_leader_worker_idx(dp_idx)
+                url_futures.append(
+                    self._policy.worker_group.run_single_worker_single_data(
+                        "report_dp_openai_server_base_url",
+                        worker_idx=worker_idx,
+                    )
+                )
             self.dp_openai_server_base_urls = [
                 url for url in ray.get(url_futures) if url is not None
             ]

nemo_rl/models/generation/megatron/megatron_worker.py

@@ -22,6 +22,7 @@
 
 import requests
 import torch
+from megatron.core import parallel_state
 from megatron.core.inference.config import (
     InferenceConfig,
     KVCacheManagementMode,
@@ -59,6 +60,11 @@ class MegatronGenerationMixin:
      - is_generation_colocated: Whether colocated or distributed.
     """
 
+    @staticmethod
+    def _is_dp_inference_client_rank() -> bool:
+        """True on the rank that may talk to this replica's mcore DP coordinator."""
+        return parallel_state.get_data_parallel_rank() == 0
+
     def _init_inference_engine_state(self) -> None:
         """Reset all inference-engine attributes to their uninitialized state."""
         self.dynamic_inference_engine = None
@@ -204,7 +210,7 @@ async def _start_inference_coordinator(self):
             inference_coordinator_port=None,
             launch_inference_coordinator=True,
         )
-        if torch.distributed.get_rank() == 0:
+        if self._is_dp_inference_client_rank():
             from megatron.core.inference.inference_client import InferenceClient
 
             self.inference_client = InferenceClient(
@@ -229,11 +235,11 @@ def _sleep(self) -> None:
         print(f"[Rank {self.rank}] paused inference engine")
 
     async def _sleep_engine(self):
-        if torch.distributed.get_rank() == 0:
+        if self._is_dp_inference_client_rank():
             self.inference_client.pause_engines()
         await self.dynamic_inference_engine.wait_until(EngineState.PAUSED)
 
-        if torch.distributed.get_rank() == 0:
+        if self._is_dp_inference_client_rank():
             self.inference_client.suspend_engines()
         await self.dynamic_inference_engine.wait_until(EngineState.SUSPENDED)
 
@@ -250,11 +256,11 @@ def _wake(self) -> None:
         print(f"[Rank {self.rank}] resumed inference engine")
 
     async def _wake_engine(self):
-        if torch.distributed.get_rank() == 0:
+        if self._is_dp_inference_client_rank():
             self.inference_client.resume_engines()
         await self.dynamic_inference_engine.wait_until(EngineState.RESUMED)
 
-        if torch.distributed.get_rank() == 0:
+        if self._is_dp_inference_client_rank():
             self.inference_client.unpause_engines()
         await self.dynamic_inference_engine.wait_until(EngineState.RUNNING)
 
@@ -334,7 +340,7 @@ def _run_async_coordinator_start(self):
 
         if (
             self.cfg["generation"]["mcore_generation_config"]["expose_http_server"]
-            and torch.distributed.get_rank() == 0
+            and self._is_dp_inference_client_rank()
         ):
             print(f"[Rank {torch.distributed.get_rank()}] Starting HTTP Server")
             self.base_url = self._setup_openai_api_server()
@@ -644,8 +650,8 @@ async def _generate_with_persistent_engine(
         from megatron.core.inference.inference_request import DynamicInferenceRequest
 
         dist_rank = torch.distributed.get_rank()
-        assert dist_rank == 0, (
-            "Only rank 0 creates a client to communicate with the coordinator"
+        assert self._is_dp_inference_client_rank(), (
+            "Only the DP-coordinator leader rank may submit inference requests"
         )
 
         print(

tests/unit/models/generation/test_megatron_generation.py

@@ -245,15 +245,18 @@ def _assert_valid_generation_output(outputs, input_data, require_generation=True
 
 async def _generate_async(mg, tokenizer, test_input_data, greedy=False):
     """Drive ``generate_async`` over single-sample microbatches and reassemble in order."""
+    dp_leaders = set(mg._policy.worker_group.dp_leader_worker_indices)
     collected = []
     for single_item_input in test_input_data.make_microbatch_iterator(
         microbatch_size=1
     ):
         async for original_idx, single_item_output in mg.generate_async(
             single_item_input, greedy=greedy
         ):
-            # The mcore coordinator only accepts requests on DP rank 0.
-            assert single_item_output["gen_leader_worker_idx"] == [0]
+            leader_idx = single_item_output["gen_leader_worker_idx"][0]
+            assert leader_idx in dp_leaders, (
+                f"gen_leader_worker_idx {leader_idx} is not a DP leader"
+            )
             collected.append((original_idx, single_item_output))
 
     collected.sort(key=lambda x: x[0])

tests/unit/models/generation/test_megatron_generation_dp_shard.py

@@ -0,0 +1,118 @@
+# 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.
+
+"""CPU tests for MegatronGeneration DP prompt sharding (driver-side dispatch)."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+import numpy as np
+
+from nemo_rl.distributed.batched_data_dict import BatchedDataDict
+from nemo_rl.distributed.named_sharding import NamedSharding
+from nemo_rl.models.generation.megatron.megatron_generation import MegatronGeneration
+
+
+def _make_megatron_generation(dp_size: int) -> MegatronGeneration:
+    mg = MegatronGeneration.__new__(MegatronGeneration)
+    mg.cfg = {"_pad_token_id": 0}
+    mg.current_generate_dp_shard_idx = 0
+    mg._policy = SimpleNamespace(
+        sharding_annotations=NamedSharding(
+            layout=np.arange(dp_size).reshape(1, dp_size, 1, 1),
+            names=[
+                "pipeline_parallel",
+                "data_parallel",
+                "context_parallel",
+                "tensor_parallel",
+            ],
+        ),
+        worker_group=MagicMock(),
+    )
+    mg._policy.worker_group.dp_size = dp_size
+    mg._policy.worker_group.dp_leader_worker_indices = list(range(0, dp_size * 2, 2))
+    mg._policy.worker_group.get_dp_leader_worker_idx.side_effect = (
+        lambda dp_idx: dp_idx * 2
+    )
+    return mg
+
+
+@pytest.mark.mcore
+def test_generate_shards_batch_across_dp_leaders():
+    """MegatronGeneration.generate must fan out like vLLM, not pin worker 0."""
+    mg = _make_megatron_generation(dp_size=2)
+    data = BatchedDataDict(
+        {
+            "input_ids": torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8]]),
+            "input_lengths": torch.tensor([2, 2, 2, 2]),
+        }
+    )
+    shard_a = BatchedDataDict(
+        {
+            "input_ids": torch.tensor([[1, 2], [3, 4]]),
+            "input_lengths": torch.tensor([2, 2]),
+        }
+    )
+    shard_b = BatchedDataDict(
+        {
+            "input_ids": torch.tensor([[5, 6], [7, 8]]),
+            "input_lengths": torch.tensor([2, 2]),
+        }
+    )
+    out_a = BatchedDataDict(
+        {
+            "output_ids": torch.tensor([[1, 2, 9], [3, 4, 10]]),
+            "generation_lengths": torch.tensor([1, 1]),
+            "unpadded_sequence_lengths": torch.tensor([3, 3]),
+            "logprobs": torch.zeros(2, 3),
+        }
+    )
+    out_b = BatchedDataDict(
+        {
+            "output_ids": torch.tensor([[5, 6, 11], [7, 8, 12]]),
+            "generation_lengths": torch.tensor([1, 1]),
+            "unpadded_sequence_lengths": torch.tensor([3, 3]),
+            "logprobs": torch.zeros(2, 3),
+        }
+    )
+
+    with patch.object(
+        BatchedDataDict,
+        "shard_by_batch_size",
+        return_value=[shard_a, shard_b],
+    ) as shard_mock:
+        future_bundle = object()
+        mg._policy.worker_group.run_all_workers_sharded_data.return_value = (
+            future_bundle
+        )
+        mg._policy.worker_group.get_all_worker_results.return_value = [out_a, out_b]
+
+        result = mg.generate(data, greedy=False)
+
+    shard_mock.assert_called_once_with(2, allow_uneven_shards=True)
+    mg._policy.worker_group.run_all_workers_sharded_data.assert_called_once_with(
+        "generate",
+        data=[shard_a, shard_b],
+        in_sharded_axes=["data_parallel"],
+        replicate_on_axes=None,
+        output_is_replicated=None,
+        common_kwargs={"greedy": False},
+    )
+    mg._policy.worker_group.get_all_worker_results.assert_called_once_with(
+        future_bundle
+    )
+    assert result.size == 4
+    assert torch.equal(result["output_ids"], torch.tensor([[1, 2, 9], [3, 4, 10], [5, 6, 11], [7, 8, 12]]))