diff --git a/torchtitan/experiments/rl/actors/generator.py b/torchtitan/experiments/rl/actors/generator.py index fed2096b8d..8da9dc3790 100644 --- a/torchtitan/experiments/rl/actors/generator.py +++ b/torchtitan/experiments/rl/actors/generator.py @@ -644,7 +644,7 @@ class VLLMGenerator(Actor, Configurable): A weight sync rides the same loop: `pull_model_state_dict` queues a `LoopDecision(LoopAction.PULL_MODEL_STATE_DICT)` applied between step bursts. The engine does NOT drain in-flight requests first ("hotswap"). This behavior can be changed - on the controller side, by blocking new requests until the engine is drained. + in the inter-generator router, by blocking new requests until the engine is drained. Args: config: Generator-specific configuration. diff --git a/torchtitan/experiments/rl/actors/inter_generator_router.py b/torchtitan/experiments/rl/actors/inter_generator_router.py new file mode 100644 index 0000000000..d39f023e44 --- /dev/null +++ b/torchtitan/experiments/rl/actors/inter_generator_router.py @@ -0,0 +1,82 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Actor wrapper for inter-generator routing.""" + +from collections.abc import Sequence +from typing import Any + +from monarch.actor import Actor, endpoint + +from torchtitan.experiments.rl.routing.inter_generator_router import ( + InterGeneratorRouter, +) +from torchtitan.experiments.rl.routing.types import RoutingContext +from torchtitan.observability import structured_logger as sl + + +class InterGeneratorRouterActor(Actor): + """Give mutable inter-generator routing state a single process owner. + + ``InterGeneratorRouter`` tracks mutable generator load and lifecycle state + and is not safe to share across threads or processes. Hosting it in a + singleton Monarch actor confines that state to one actor process and makes + other actors access it through endpoints instead of sharing the router. + """ + + def __init__( + self, + config: InterGeneratorRouter.Config, + *, + generators: Sequence[Any], + ) -> None: + self._router = InterGeneratorRouter(config, generators=generators) + + @endpoint + async def generate( + self, + prompt_token_ids: list[int], + *, + request_id: str, + routing_session_id: str | None, + sampling_config: Any | None, + metrics_prefix: str, + ) -> Any: + """Route one generation call while holding global routing state.""" + # Dispatches to the chosen generator's rank-0 intake via call_one, so + # it returns the Completion directly (no ValueMesh unwrap). + return await self._router.route( + "generate", + prompt_token_ids, + request_id=request_id, + # VLLMGenerator.generate also requires this field for its + # intra-mesh DP routing. + routing_session_id=routing_session_id, + sampling_config=sampling_config, + metrics_prefix=metrics_prefix, + # Load is measured as in-flight request count (one unit per call). + routing_ctx=RoutingContext( + estimated_cost=1, + session_id=routing_session_id, + ), + ) + + @endpoint + async def start_engine_loop(self) -> None: + await self._router.fanout("start_engine_loop") + + @endpoint + async def sync_log_step(self, step: int) -> None: + sl.set_step(step) + await self._router.fanout("sync_log_step", step) + + @endpoint + async def pull_model_state_dict(self, policy_version: int) -> None: + await self._router.pull_model_state_dict(policy_version=policy_version) + + @endpoint + async def close_generators(self) -> list[Any | BaseException]: + return await self._router.fanout("close", return_exceptions=True) diff --git a/torchtitan/experiments/rl/components/weight_sync.py b/torchtitan/experiments/rl/components/weight_sync.py index 24f989c6e2..44fd362abf 100644 --- a/torchtitan/experiments/rl/components/weight_sync.py +++ b/torchtitan/experiments/rl/components/weight_sync.py @@ -6,16 +6,22 @@ """Overlap the trainer->generator weight handoff with the next training step.""" +from __future__ import annotations + import asyncio import time +from typing import TYPE_CHECKING from torchtitan.experiments.rl.components.work_buffer import RolloutGroupWorkBuffer from torchtitan.experiments.rl.observability import metrics as m -from torchtitan.experiments.rl.routing.inter_generator_router import ( - InterGeneratorRouter, -) from torchtitan.observability import structured_logger as sl +if TYPE_CHECKING: + from torchtitan.experiments.rl.actors.inter_generator_router import ( + InterGeneratorRouterActor, + ) + from torchtitan.experiments.rl.actors.trainer import PolicyTrainer + # dummy no-op for step 0, used in WeightSyncManager async def _noop() -> None: return None @@ -47,8 +53,8 @@ class WeightSyncManager: def __init__( self, *, - trainer, # PolicyTrainer actor handle - generator_router: InterGeneratorRouter, + trainer: PolicyTrainer, + generator_router: InterGeneratorRouterActor, group_buffer: RolloutGroupWorkBuffer, num_prompts_per_train_step: int, ) -> None: @@ -112,7 +118,7 @@ async def _generator_pull_and_release_buffer_slots( await push_task with sl.log_trace_span("generator_pull_model_state_dict"): start = time.perf_counter() - await self._generator_router.pull_model_state_dict(policy_version=version) + await self._generator_router.pull_model_state_dict.call_one(version) self._last_pull_s = time.perf_counter() - start # TODO(perf): pull_model_state_dict awaits ALL generators before we release any buffer slots, # so a generator that finishes its pull early idles until the slowest one. Investigate diff --git a/torchtitan/experiments/rl/controller.py b/torchtitan/experiments/rl/controller.py index 4e11561566..304c1cd900 100644 --- a/torchtitan/experiments/rl/controller.py +++ b/torchtitan/experiments/rl/controller.py @@ -102,11 +102,15 @@ import torch # noqa: F401 import torchstore as ts import tyro -from monarch.actor import ProcMesh + +from monarch.actor import ProcMesh, this_host from monarch.spmd import setup_torch_elastic_env_async from torchtitan.config import CompileConfig, Configurable from torchtitan.experiments.rl.actors.generator import SamplingConfig, VLLMGenerator +from torchtitan.experiments.rl.actors.inter_generator_router import ( + InterGeneratorRouterActor, +) from torchtitan.experiments.rl.actors.trainer import PolicyTrainer from torchtitan.experiments.rl.components.batcher import Batcher from torchtitan.experiments.rl.components.training_sample_builder import ( @@ -134,7 +138,6 @@ from torchtitan.experiments.rl.routing.inter_generator_router import ( InterGeneratorRouter, ) -from torchtitan.experiments.rl.routing.types import RoutingContext from torchtitan.experiments.rl.types import Completion, TrainingBatch from torchtitan.observability import structured_logger as sl from torchtitan.protocols.model_spec import ModelSpec @@ -432,7 +435,7 @@ def __post_init__(self): def __init__(self, config: Config): self.config = config self.trainer: PolicyTrainer | None = None - self.generator_router: InterGeneratorRouter | None = None + self.generator_router: InterGeneratorRouterActor | None = None # Resume step (0 = fresh); set in setup_async from the loaded checkpoint. self.start_step = 0 self._proc_meshes = [] @@ -470,19 +473,22 @@ async def close(self): logger.exception("trainer.close failed") if self.generator_router is not None: - close_results = await self.generator_router.fanout( - "close", return_exceptions=True - ) - for idx, result in enumerate(close_results): - if isinstance(result, BaseException): - actor_name = ( - "generator" if len(close_results) == 1 else f"generator[{idx}]" - ) - logger.error( - "%s.close failed", - actor_name, - exc_info=(type(result), result, result.__traceback__), - ) + try: + close_results = await self.generator_router.close_generators.call_one() + for idx, result in enumerate(close_results): + if isinstance(result, BaseException): + actor_name = ( + "generator" + if len(close_results) == 1 + else f"generator[{idx}]" + ) + logger.error( + "%s.close failed", + actor_name, + exc_info=(type(result), result, result.__traceback__), + ) + except Exception: + logger.exception("generator_router.close_generators failed") try: self.metrics_processor.close() @@ -519,22 +525,12 @@ async def generate( routing_session_id: str | None = None, sampling_config: SamplingConfig | None = None, ) -> Completion | None: - # 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", + return await self.generator_router.generate.call_one( prompt_token_ids, request_id=request_id, - # VLLMGenerator.generate also requires this field for its - # intra-mesh DP routing. routing_session_id=routing_session_id, sampling_config=sampling_config, metrics_prefix=metrics_prefix, - # Load is measured as in-flight request count (one unit per call). - routing_ctx=RoutingContext( - estimated_cost=1, - session_id=routing_session_id, - ), ) return generate @@ -553,9 +549,9 @@ async def setup_async( weight push/pull are all ``await``-based runtime side effects that cannot run in a synchronous constructor. - The trainer and generator meshes are provisioned by the caller - (see ``create_proc_mesh``) on disjoint GPUs; this method only - spawns the actors on them and synchronizes initial weights from + The trainer and generator meshes are provisioned by the caller (see + ``spawn_proc_mesh``). The router mesh is created on the controller host. + This method spawns the actors and synchronizes initial weights from trainer to generator. Must be called before :meth:`run`. Args: @@ -605,8 +601,9 @@ async def setup_async( # provisioner logic. Pull a PerHostProvisioner.spawn_meshes(...) helper and # shrink this span to a single call. with sl.log_trace_span("mesh_spawn"): + router_mesh = this_host().spawn_procs(per_host={"cpus": 1}) # Store proc meshes for cleanup - self._proc_meshes = [trainer_mesh, *generator_meshes] + self._proc_meshes = [router_mesh, trainer_mesh, *generator_meshes] await setup_torch_elastic_env_async(trainer_mesh) for generator_mesh in generator_meshes: @@ -641,7 +638,12 @@ async def setup_async( output_dir=config.dump_folder, ) generators.append(generator) - self.generator_router = config.generator_router.build(generators=generators) + self.generator_router = router_mesh.spawn( + "generator_router", + InterGeneratorRouterActor, + config.generator_router, + generators=generators, + ) # Initialize TorchStore for weight sync between trainer and generator. # StorageVolumes are spawned on the trainer mesh so they are colocated @@ -667,15 +669,13 @@ async def setup_async( # 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") + await self.generator_router.start_engine_loop.call_one() # 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() with sl.log_trace_span("generator_pull_model_state_dict"): - await self.generator_router.pull_model_state_dict( - policy_version=self.start_step - ) + await self.generator_router.pull_model_state_dict.call_one(self.start_step) # TODO: fold validation into a Validator(Configurable) the controller attaches, instead of 4 methods. @sl.log_trace_span("_collect_validation_rollouts") @@ -1063,7 +1063,7 @@ async def _trainer_loop( sl.set_step(step) # propagate the step counter to the actors with sl.log_trace_span("sync_log_step"): await self.trainer.sync_log_step.call(step) - await self.generator_router.fanout("sync_log_step", step) + await self.generator_router.sync_log_step.call_one(step) step_timer = MetricsTimer() with sl.log_trace_span("train_step"), step_timer.record( diff --git a/torchtitan/experiments/rl/routing/inter_generator_router.py b/torchtitan/experiments/rl/routing/inter_generator_router.py index d248faccfd..f1247302e2 100644 --- a/torchtitan/experiments/rl/routing/inter_generator_router.py +++ b/torchtitan/experiments/rl/routing/inter_generator_router.py @@ -32,7 +32,7 @@ class _GeneratorState(Enum): @dataclass(kw_only=True, slots=True) class _GeneratorHandle(RoutingCandidate): - """Controller-side metadata for one generator mesh.""" + """Router-side metadata for one generator mesh.""" actor: Any """Monarch actor handle for the full generator mesh. Used for fan-out calls diff --git a/torchtitan/experiments/rl/routing/intra_generator_router.py b/torchtitan/experiments/rl/routing/intra_generator_router.py index 9b21d432eb..a30dc67e3e 100644 --- a/torchtitan/experiments/rl/routing/intra_generator_router.py +++ b/torchtitan/experiments/rl/routing/intra_generator_router.py @@ -32,7 +32,7 @@ class _DPRankHandle(RoutingCandidate): class IntraGeneratorRouter(Configurable): """Router that partitions requests across the DP ranks within one generator. - This is layer 2 of the two-layer routing design: the controller-side + This is layer 2 of the two-layer routing design: the router-actor-side ``InterGeneratorRouter`` routes a call across generators; this router then routes each request across the DP ranks within one generator. diff --git a/torchtitan/experiments/rl/routing/strategies.py b/torchtitan/experiments/rl/routing/strategies.py index d44759d86a..3c6898d699 100644 --- a/torchtitan/experiments/rl/routing/strategies.py +++ b/torchtitan/experiments/rl/routing/strategies.py @@ -11,7 +11,7 @@ ``reserved_load`` field plus object identity -- so the same strategy classes serve both routing layers in the RL generator: -- Layer 1: ``InterGeneratorRouter`` (controller side) routes a call across +- Layer 1: ``InterGeneratorRouter`` (router actor side) routes a call across generator *meshes* (replicas). See ``inter_generator_router.py``. - Layer 2: ``IntraGeneratorRouter`` (in-mesh, rank-0 side) routes a request across the *data-parallel ranks* within one generator mesh. See diff --git a/torchtitan/experiments/rl/tests/test_shutdown.py b/torchtitan/experiments/rl/tests/test_shutdown.py index 52dc2735cc..a1fef75f63 100644 --- a/torchtitan/experiments/rl/tests/test_shutdown.py +++ b/torchtitan/experiments/rl/tests/test_shutdown.py @@ -28,7 +28,12 @@ def __init__(self, config=None): self.setup_generator_meshes = None self.instances.append(self) - async def setup_async(self, *, trainer_mesh=None, generator_meshes=None): + async def setup_async( + self, + *, + trainer_mesh=None, + generator_meshes=None, + ): self.events.append("setup") self.setup_trainer_mesh = trainer_mesh self.setup_generator_meshes = generator_meshes @@ -197,9 +202,10 @@ def stub_mesh_provisioning(monkeypatch): monkeypatch.setattr(train, "_compute_generator_world_size", lambda p: 1) def _spawn_proc_mesh(*args, num_generators=1, **kwargs): - return "trainer_mesh", [ - f"generator_mesh_{idx}" for idx in range(num_generators) - ] + return ( + "trainer_mesh", + [f"generator_mesh_{idx}" for idx in range(num_generators)], + ) monkeypatch.setattr(train, "spawn_proc_mesh", _spawn_proc_mesh) @@ -295,6 +301,19 @@ async def call(self): raise RuntimeError(f"{self._name} failed") +class _RouterCloseEndpoint: + def __init__(self, router): + self._router = router + + async def call_one(self): + return await self._router.fanout("close", return_exceptions=True) + + +class _StubRouterActor: + def __init__(self, router): + self.close_generators = _RouterCloseEndpoint(router) + + class _StubActor: def __init__(self, name, events, raises=False): self.close = _StubEndpoint(name, events, raises) @@ -321,9 +340,11 @@ async def stop(self): def _set_generator_router(rl_trainer, generators): - rl_trainer.generator_router = InterGeneratorRouter( - InterGeneratorRouter.Config(), - generators=generators, + rl_trainer.generator_router = _StubRouterActor( + InterGeneratorRouter( + InterGeneratorRouter.Config(), + generators=generators, + ) ) diff --git a/torchtitan/experiments/rl/tests/test_weight_sync.py b/torchtitan/experiments/rl/tests/test_weight_sync.py index 1ba0bfb729..943c849a8c 100644 --- a/torchtitan/experiments/rl/tests/test_weight_sync.py +++ b/torchtitan/experiments/rl/tests/test_weight_sync.py @@ -40,10 +40,16 @@ class _FakeRouter: def __init__(self, on_pull): self._on_pull = on_pull self.pulled_versions: list[int] = [] + self.pull_model_state_dict = _PullEndpoint(self) - async def pull_model_state_dict(self, *, policy_version): - self.pulled_versions.append(policy_version) - await self._on_pull() + +class _PullEndpoint: + def __init__(self, router): + self._router = router + + async def call_one(self, policy_version): + self._router.pulled_versions.append(policy_version) + await self._router._on_pull() class _FakeBuffer: diff --git a/torchtitan/experiments/rl/train.py b/torchtitan/experiments/rl/train.py index 2a5109e97a..b349e33c36 100644 --- a/torchtitan/experiments/rl/train.py +++ b/torchtitan/experiments/rl/train.py @@ -32,10 +32,7 @@ # imports transitively importing torch. os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") -# TODO: Remove `_src` after monarch cuts a new release. This is already a public -# API in monarch nightly. https://github.com/meta-pytorch/monarch/pull/4327 -from monarch._src.actor.host_mesh import default_bootstrap_cmd -from monarch.actor import HostMesh, ProcMesh, this_host +from monarch.actor import default_bootstrap_cmd, HostMesh, ProcMesh, this_host from torchtitan.config import ConfigManager, ParallelismConfig from torchtitan.experiments.rl.controller import Controller