Controller sends request to generator rank0 when applicable (#3870)

`VLLMGenerator` uses its rank0 to coordinate the SPMD job running on it,
so for the following 2 endpoints, controller only needs to send requests
to generator's rank0:

* generate
* pull_model_state_dict

However, currently controller still does a Monarch call, which sends the
requests to all ranks. Just non-0 ranks, the request is a NOOP, except
the one time `start_engine_loop` step in `generate`.

This PR changes that to sending the request to rank 0 only for these 2
endpoints. This change makes the semantics more explicit, in terms of
rank0 is the one talks to controller.

Author: pzhan9

torchtitan/experiments/rl/actors/generator.py

@@ -906,7 +906,7 @@ def __init__(
 
         self._pull_model_state_dict_future: asyncio.Future[int] | None = None
 
-        # Background asyncio.Task running _engine_loop; None until the first generate/pull starts it.
+        # Background asyncio.Task running _engine_loop; None until start_engine_loop starts it.
         self._engine_loop_task: asyncio.Task | None = None
 
         logger.info("Generator initialized with vLLM engine")
@@ -940,6 +940,21 @@ async def sync_log_step(self, step: int, relative_step: int | None = None) -> No
         """Sync the structured-logger step counter from the controller."""
         sl.set_step(step, relative_step=relative_step)
 
+    @endpoint
+    async def start_engine_loop(self) -> None:
+        """Start the background engine loop on every rank (one-time, idempotent)."""
+        if self._engine_loop_task is None:
+            self._engine_loop_task = asyncio.create_task(self._engine_loop())
+
+    def _rank0_check_engine_loop_running(self, endpoint_name: str) -> None:
+        """Guard for the rank-0-only endpoints"""
+        assert self._rank == 0, f"{endpoint_name} must be routed to rank 0 only"
+        if self._engine_loop_task is None:
+            raise RuntimeError(
+                "engine loop not started; call start_engine_loop on all ranks "
+                f"before {endpoint_name}"
+            )
+
     @endpoint
     @sl.log_trace_span("generate")
     async def generate(
@@ -950,12 +965,13 @@ async def generate(
         routing_session_id: str,
         sampling_config: SamplingConfig | None = None,
         metrics_prefix: str = "generator",
-    ) -> Completion | None:
+    ) -> Completion:
         """Generates one completion for one prompt.
 
-        Returns the `Completion` on rank 0 and `None` on followers. The completion carries its
-        own per-generation metrics (`Completion.metrics`), which the controller attaches to the
-        rollout turn.
+        Can be accepted by rank 0 only (rank 0 owns the queue + futures and
+        drives the followers through the engine loop). Returns the `Completion`,
+        which carries its own per-generation metrics (`Completion.metrics`) that
+        the controller attaches to the rollout turn.
 
         Args:
             prompt_token_ids: One tokenized prompt `[token_ids]`.
@@ -969,20 +985,11 @@ async def generate(
 
         Example:
 
-            completion = await generator.generate.call(
+            completion = await generator.slice(hosts=0, gpus=0).generate.call_one(
                 [1, 2, 3], request_id="step=3/group=0/sample=0/turn=0",
             )
-            # rank 0 -> Completion(token_ids=[...], metrics=[Metric("generator/queue_time_ms", ...)]);
-            # followers -> None
         """
-
-        # Starting requires asyncio, which isn't available in the sync __init__.
-        # Start on first call; no-op after.
-        await self._ensure_engine_loop()
-
-        # Only rank 0 owns the queue + futures moving forward.
-        if self._rank != 0:
-            return None
+        self._rank0_check_engine_loop_running("generate")
 
         sampling = (
             sampling_config if sampling_config is not None else self.config.sampling
@@ -1010,11 +1017,6 @@ async def generate(
         # Await outside the lock so other generate / pull calls can proceed meanwhile.
         return await generation_future
 
-    async def _ensure_engine_loop(self) -> None:
-        """Start the single background engine loop on first use (idempotent); runs until `close()`."""
-        if self._engine_loop_task is None:
-            self._engine_loop_task = asyncio.create_task(self._engine_loop())
-
     @sl.log_trace_span("engine_loop")
     async def _engine_loop(self) -> None:
         """Non-stop loop running on all ranks to produce new tokens.
@@ -1199,13 +1201,7 @@ async def pull_model_state_dict(self, version: int) -> None:
         # TODO: if an incoming request is received while another pull request is queued
         # we should drop the older request and pull the latest version instead
 
-        # Starting requires asyncio, which isn't available in the sync __init__.
-        # Start on first call; no-op after.
-        await self._ensure_engine_loop()
-
-        # Only rank 0 owns the queue + futures moving forward.
-        if self._rank != 0:
-            return
+        self._rank0_check_engine_loop_running("pull_model_state_dict")
 
         # A placeholder future for the engine loop to resolve once the pull has been applied.
         pull_model_state_dict_future: asyncio.Future[

torchtitan/experiments/rl/controller.py

@@ -438,7 +438,9 @@ async def generate(
             routing_session_id: str | None = None,
             sampling_config: SamplingConfig | None = None,
         ) -> Completion | None:
-            result = await self.generator_router.route(
+            # Dispatches to the chosen generator's rank-0 intake via call_one, so
+            # it returns the Completion directly (no ValueMesh unwrap).
+            return await self.generator_router.route(
                 "generate",
                 prompt_token_ids,
                 request_id=request_id,
@@ -453,7 +455,6 @@ async def generate(
                     session_id=routing_session_id,
                 ),
             )
-            return self._get_rank_0_value(result)
 
         return generate
 
@@ -591,6 +592,12 @@ async def setup_async(
                 generators.append(generator)
             self.generator_router = config.generator_router.build(generators=generators)
 
+        # Start each generator's engine loop on all ranks once, before any
+        # rank-0-only generate / pull (rank 0 drives the followers through this
+        # loop, so every rank must be running it first).
+        with sl.log_trace_span("generator_start_engine_loop"):
+            await self.generator_router.fanout("start_engine_loop")
+
         # Initial weight sync: only the trainer loads weights; generators pull at start_step.
         with sl.log_trace_span("trainer_push_model_state_dict"):
             await self.trainer.push_model_state_dict.call()

torchtitan/experiments/rl/routing/inter_generator_router.py

@@ -35,7 +35,12 @@ class _GeneratorHandle(RoutingCandidate):
     """Controller-side metadata for one generator mesh."""
 
     actor: Any
-    """Monarch actor handle for the generator mesh."""
+    """Monarch actor handle for the full generator mesh. Used for fan-out calls
+    that every rank must run."""
+
+    rank0_actor: Any
+    """Cached rank-0 slice of ``actor``. Used for calls that only rank 0 needs
+    to run."""
 
     reserved_load: int = 0
     """Controller-side estimate of in-flight routed generation work."""
@@ -87,7 +92,11 @@ def __init__(
     ):
         self._config = config
         self._generators = [
-            _GeneratorHandle(actor=generator) for generator in generators
+            _GeneratorHandle(
+                actor=generator,
+                rank0_actor=generator.flatten("rank").slice(rank=0),
+            )
+            for generator in generators
         ]
         if not self._generators:
             raise ValueError("InterGeneratorRouter requires at least one generator")
@@ -143,15 +152,16 @@ async def route(
         routing_ctx: RoutingContext,
         **kwargs,
     ) -> Any:
-        """Dispatch one call to a strategy-chosen serving generator; return its result."""
-
+        """Dispatch one call to a strategy-chosen serving generator's rank 0;
+        return its result.
+        """
         await self._serving.wait()
         candidates = self._candidates()
         assert candidates, "serving event was set with no serving generators"
         h = self._strategy.choose(routing_ctx, candidates)
         self._reserve(h, routing_ctx.estimated_cost)
         try:
-            return await getattr(h.actor, method).call(*args, **kwargs)
+            return await getattr(h.rank0_actor, method).call_one(*args, **kwargs)
         finally:
             self._release(h, routing_ctx.estimated_cost)
 
@@ -194,15 +204,15 @@ async def _pull_one(h: _GeneratorHandle) -> None:
                 # Hot swap: pull concurrently with in-flight generation, without
                 # draining. Whether the pull is genuinely concurrent and safe is
                 # up to the generator's implementation.
-                await h.actor.pull_model_state_dict.call(policy_version)
+                await h.rank0_actor.pull_model_state_dict.call_one(policy_version)
             else:
                 # Drain: stop routing to this generator and wait for in-flight
                 # work to finish before pulling, then re-admit it.
                 self._set_state(h, _GeneratorState.SYNCING)
                 try:
                     with sl.log_trace_span("router_drain_wait"):
                         await h.idle.wait()
-                    await h.actor.pull_model_state_dict.call(policy_version)
+                    await h.rank0_actor.pull_model_state_dict.call_one(policy_version)
                 finally:
                     self._set_state(h, _GeneratorState.SERVING)
 

torchtitan/experiments/rl/tests/test_inter_generator_router.py

@@ -30,7 +30,7 @@ def __init__(self, value=None, *, wait: bool = False, raises: bool = False):
         if not wait:
             self.release.set()
 
-    async def call(self, *args, **kwargs):
+    async def call_one(self, *args, **kwargs):
         self.calls.append((args, kwargs))
         self.started.set()
         await self.release.wait()
@@ -40,6 +40,8 @@ async def call(self, *args, **kwargs):
 
 
 class _Actor:
+    """A single-rank generator mesh fake."""
+
     def __init__(
         self,
         name: str,
@@ -51,6 +53,15 @@ def __init__(
         self.generate = _Endpoint(name, wait=wait_generate)
         self.pull_model_state_dict = _Endpoint(None, wait=wait_pull, raises=raises_pull)
 
+    def flatten(self, *args, **kwargs):
+        return self
+
+    def slice(self, **kwargs):
+        return self
+
+    def __len__(self):
+        return 1
+
 
 def _router(actors, *, strategy=None, hot_swap=False) -> InterGeneratorRouter:
     return InterGeneratorRouter(
@@ -339,7 +350,8 @@ async def _run():
         actors[1].pull_model_state_dict.release.set()
         assert (
             await asyncio.wait_for(
-                router.route("generate", routing_ctx=RoutingContext()), timeout=1.0
+                router.route("generate", routing_ctx=RoutingContext()),
+                timeout=1.0,
             )
             == "gen1"
         )

torchtitan/experiments/rl/tests/test_shutdown.py

@@ -233,6 +233,17 @@ class _StubActor:
     def __init__(self, name, events, raises=False):
         self.close = _StubEndpoint(name, events, raises)
 
+    def flatten(self, *args, **kwargs):
+        # The router builds a rank-0 handle via flatten("rank").slice(rank=0);
+        # a single-rank stub collapses to itself for both.
+        return self
+
+    def slice(self, **kwargs):
+        return self
+
+    def __len__(self):
+        return 1
+
 
 class _StubMesh:
     def __init__(self, name, events):