diff --git a/examples/arctic_rl/run_gsm8k_grpo_arl_zorro_yes.sh b/examples/arctic_rl/run_gsm8k_grpo_arl_zorro_yes.sh index c196cb2add1..fc6cff538ab 100755 --- a/examples/arctic_rl/run_gsm8k_grpo_arl_zorro_yes.sh +++ b/examples/arctic_rl/run_gsm8k_grpo_arl_zorro_yes.sh @@ -65,14 +65,15 @@ python3 -m verl.trainer.main_ppo \ algorithm.use_kl_in_reward=False \ trainer.use_legacy_worker_impl=disable \ trainer.remote_backend=arctic \ - remote_backend.arctic.colocate=False \ - remote_backend.arctic.training_gpus=1 \ - remote_backend.arctic.sampling_gpus=1 \ - remote_backend.arctic.log_prob_gpus=0 \ - remote_backend.arctic.zero_optimization.stage=2 \ - remote_backend.arctic.zero_optimization.offload_optimizer.device=none \ - remote_backend.arctic.zero_optimization.offload_param.device=none \ - remote_backend.arctic.use_zorro=True \ + remote_backend=arctic \ + remote_backend.colocate=False \ + remote_backend.training_gpus=1 \ + remote_backend.sampling_gpus=1 \ + remote_backend.log_prob_gpus=0 \ + remote_backend.zero_optimization.stage=2 \ + remote_backend.zero_optimization.offload_optimizer.device=none \ + remote_backend.zero_optimization.offload_param.device=none \ + remote_backend.use_zorro=True \ trainer.critic_warmup=0 \ trainer.logger="['console']" \ trainer.experiment_name=gsm8k_grpo_qwen3_0p6b_ngpu1_gbs16_rolln5_zorroTrue \ diff --git a/verl/remote_backend/__init__.py b/verl/remote_backend/__init__.py index 7e0db7380c8..479d1ed6d4d 100644 --- a/verl/remote_backend/__init__.py +++ b/verl/remote_backend/__init__.py @@ -8,19 +8,24 @@ Pieces: -* :class:`RemoteBackend` (``base.py``) — the all-abstract contract - every backend implements. -* :class:`RemoteBackendRegistry` (``base.py``) — name → class lookup so - ``trainer.remote_backend=""`` resolves to a concrete adapter. -* :class:`RemoteBackendTrainer` (``trainer.py``) — `RayPPOTrainer` subclass - that creates the backend on the driver and threads its reconnect handle - to every worker. -* :class:`RemoteBackendActorRolloutRefWorker` (``worker.py``) — the - backend-agnostic CPU forwarder. -* ``worker_utils.py`` — small generic tensor / metric helpers shared by - the forwarder and backend adapters. +* :class:`RemoteBackend` (``base.py``) — minimal ABC: lifecycle + (``from_config`` / ``reconnect_handle`` / ``destroy``) + weight sync + + checkpoint + a single-forwarder parallelism flag. Compute/update + op signatures intentionally live on the per-backend adapter, not + here. +* :class:`RemoteBackendRegistry` (``base.py``) — name → class lookup + populated by explicit adapter imports (no lazy MODULES table). +* :class:`RemoteBackendTrainer` (``trainer.py``) — `RayPPOTrainer` + subclass that creates the backend on the driver and threads its + reconnect handle to every worker. +* ``workers//`` — per-backend forwarder worker (Arctic + ships as ``workers/arctic_rl/`` -> :class:`ArcticRLActorRolloutRefWorker`). +* ``worker_utils.py`` — small generic tensor / metric helpers shared + across per-backend workers. -See :mod:`verl.trainer.ppo.arctic_rl_client` for a reference adapter. +Adapter modules (the concrete :class:`RemoteBackend` implementations) +live under :mod:`verl.workers.remote_client`. See +:mod:`verl.workers.remote_client.arctic_rl` for a reference adapter. """ from verl.remote_backend.base import RemoteBackend, RemoteBackendRegistry diff --git a/verl/remote_backend/base.py b/verl/remote_backend/base.py index dce089ee1d5..7b27a389751 100644 --- a/verl/remote_backend/base.py +++ b/verl/remote_backend/base.py @@ -1,22 +1,27 @@ """`RemoteBackend` ABC + `RemoteBackendRegistry`. -The ABC is intentionally minimal: every method is abstract, so each -backend declares its own implementation explicitly. There are no -default behaviours to silently inherit. +The ABC is intentionally minimal — it only enforces the *lifecycle* +contract that verl's trainer side needs to know about. Compute/update +op signatures (``compute_log_prob``, ``update_actor``, ``generate``) +intentionally live on the concrete per-backend adapter and its +matching per-backend worker, not on the ABC, so different backends can +shape those calls however suits them (one backend for training and +another for sampling, different payload schemas, etc.) without growing +this base class. -What the abstraction is responsible for (what verl drives): +What the ABC owns (what verl drives): * Lifecycle: ``from_config`` (sole constructor, takes an optional ``handle=`` for re-attach) / ``reconnect_handle`` / ``destroy``. -* Core RL ops: ``compute_log_prob`` / ``update_actor`` / - ``generate`` / ``update_weights`` / ``save_checkpoint``. +* Weight sync + checkpoint: ``update_weights`` / ``save_checkpoint`` + (called from ``ONE_TO_ALL`` worker hooks). * Parallelism contract: ``requires_single_forwarder`` (read by :class:`verl.remote_backend.trainer.RemoteBackendTrainer` to decide whether to assert ``n_gpus_per_node × nnodes == 1``). -What the abstraction is NOT responsible for: payload schemas, wire -formats, loss-function plumbing — those live entirely inside each -backend. +What the ABC does NOT own: payload schemas, wire formats, loss-function +plumbing, compute/update method signatures — those live entirely on the +concrete backend + its per-backend worker. """ from __future__ import annotations @@ -24,7 +29,6 @@ import abc from typing import Any, Callable -import torch from omegaconf import DictConfig @@ -83,68 +87,8 @@ def destroy(self) -> None: """ # ------------------------------------------------------------------ # - # Core RL operations called by `RemoteBackendActorRolloutRefWorker` + # Weight sync + checkpoint (called from ONE_TO_ALL worker hooks). # ------------------------------------------------------------------ # - # Inputs are verl `TensorDict`s; outputs are plain Python dicts that - # the worker wraps / augments with FlopsCounter MFU before returning a - # `TensorDict` to the trainer. Backends pick their own wire format - # below this layer. - - @abc.abstractmethod - async def compute_log_prob( - self, - data, - *, - ref: bool, - calculate_entropy: bool, - rollout_n: int, - temperature, - pad_token_id: int, - ) -> dict[str, Any]: - """Forward-only pass; either the actor (``ref=False``) or the - reference model (``ref=True``). - - ``async`` so the underlying RPC (typically a Ray actor call) can be - awaited without blocking the forwarder's event loop. - - Returns: - ``{"model_output": {"log_probs": Tensor, - "entropy": Tensor?}, "metrics": dict}``. The generic worker - reconstructs nested-jagged tensors from - ``model_output["log_probs"]`` / ``["entropy"]`` and attaches - MFU. - """ - - @abc.abstractmethod - async def update_actor( - self, - data, - *, - actor_config, - pad_token_id: int, - rollout_n: int, - temperature, - ) -> dict[str, Any]: - """Forward-backward + optimizer step on a global batch. - - ``async`` so the underlying RPC (typically a Ray actor call) can be - awaited without blocking the forwarder's event loop. - - Returns: - ``{"loss": float|list[float], "metrics": dict, - "global_token_num": list[int]}``. ``metrics["loss"]`` may - also be present; the generic worker computes MFU from - ``global_token_num`` and aggregates metrics. - """ - - @abc.abstractmethod - async def generate( - self, - prompt_ids: torch.Tensor, - sampling_params: dict[str, Any], - ) -> list: - """Sample a rollout. Called from the rollout server that owns - this backend (e.g. ``ArcticLLMEngine``).""" @abc.abstractmethod async def update_weights(self) -> dict[str, Any]: @@ -172,9 +116,8 @@ def requires_single_forwarder(self) -> bool: With more than one forwarder worker, ``ONE_TO_ALL`` calls (``save_checkpoint``, ``update_weights``, ``to``, ``set_loss_fn``) get duplicated against the single backend, and mesh-dispatched - calls (``compute_log_prob``, ``update_actor``) fragment the - global batch across forwarders that each forward the whole - batch downstream. + compute/update calls fragment the global batch across forwarders + that each forward the whole batch downstream. Returning ``True`` enables the assert; returning ``False`` opts out (the backend takes responsibility for validating its own @@ -185,28 +128,24 @@ def requires_single_forwarder(self) -> bool: class RemoteBackendRegistry: """Process-wide registry of name → :class:`RemoteBackend` class. - Backends register themselves at import time:: + Registration is explicit and happens when the user (or their entry + script) imports the adapter module they want to use:: + + # in main_ppo.py, conditioned on `trainer.remote_backend == "arctic"`: + from verl.workers.remote_client import arctic_rl # noqa: F401 - @RemoteBackendRegistry.register("arctic") - class ArcticBackend(RemoteBackend): ... + # arctic_rl decorates its class with + # @RemoteBackendRegistry.register("arctic"), so by import time the + # name is available to `get()` / `create()`. - Callers don't need to know which module to import for a given - backend name — :meth:`create` and :meth:`get` lazy-import the - adapter module listed in :attr:`MODULES`. To plug in a new - backend (e.g. ``"tinker"``), add an entry to ``MODULES`` and - decorate the class with ``@RemoteBackendRegistry.register("tinker")``. + There is intentionally no eager `MODULES` table that pre-imports + every known adapter — that would force the process to take on the + transitive deps (vLLM, arctic-training, tinker, ...) of every + backend even when only one is in use. """ _backends: dict[str, type[RemoteBackend]] = {} - # Backend name → dotted module path to import on first use. Ray - # child procs (and the driver) only need to know the name; the - # registry resolves to the right adapter module without forcing - # `main_ppo.py` to grow a per-backend `import` line. - MODULES: dict[str, str] = { - "arctic": "verl.trainer.ppo.arctic_rl_client", - } - @classmethod def register(cls, name: str) -> Callable[[type[RemoteBackend]], type[RemoteBackend]]: def _decorator(backend_cls: type[RemoteBackend]) -> type[RemoteBackend]: @@ -223,20 +162,13 @@ def _decorator(backend_cls: type[RemoteBackend]) -> type[RemoteBackend]: @classmethod def get(cls, name: str) -> type[RemoteBackend]: - """Resolve ``name`` to a registered backend class, lazy-importing - the adapter module if needed.""" if name not in cls._backends: - module_path = cls.MODULES.get(name) - if module_path is None: - raise KeyError( - f"Unknown remote backend '{name}'. Registered: " - f"{sorted(cls._backends)}. Known modules: " - f"{sorted(cls.MODULES)}. Either add an entry to " - f"RemoteBackendRegistry.MODULES or import the module " - "that registers the backend before calling get()." - ) - import importlib - importlib.import_module(module_path) + raise KeyError( + f"Unknown remote backend '{name}'. Registered: " + f"{sorted(cls._backends)}. Import the adapter module " + "(e.g. `from verl.workers.remote_client import arctic_rl`) " + "before calling `get()`." + ) return cls._backends[name] @classmethod diff --git a/verl/remote_backend/worker_utils.py b/verl/remote_backend/worker_utils.py index 818da5a51f1..ebf37e051bc 100644 --- a/verl/remote_backend/worker_utils.py +++ b/verl/remote_backend/worker_utils.py @@ -1,14 +1,15 @@ -"""Backend-agnostic tensor + metric helpers used by -:class:`verl.remote_backend.worker.RemoteBackendActorRolloutRefWorker`. +"""Backend-agnostic tensor + metric helpers shared across per-backend +forwarder workers under :mod:`verl.remote_backend.workers`. -Lifted here (out of the worker module) so that backends and other -auxiliary modules can reuse them without importing the heavy Worker class -or its single-controller dispatch decorators. +Lifted here (out of any per-backend module) so each backend's worker +can reuse them without re-implementing nested-jagged reconstruction or +metric normalization, and without importing one another's Worker +classes or single-controller dispatch decorators. -The forwarder worker stays backend-agnostic by routing every payload- -encoding decision through ``RemoteBackend.compute_log_prob`` / -``RemoteBackend.update_actor``; these helpers cover the small set of -generic transformations the worker needs around those calls: +Each per-backend worker keeps its own payload-encoding decisions on +its adapter's compute/update methods (the ABC intentionally doesn't +fix those signatures); these helpers cover the small set of generic +transformations a worker needs around those calls: * :func:`make_njt` — reconstruct a nested-jagged tensor from a dense ``[B, L]`` tensor returned by a backend, using the offsets carried in diff --git a/verl/remote_backend/workers/__init__.py b/verl/remote_backend/workers/__init__.py new file mode 100644 index 00000000000..c51dbee98ad --- /dev/null +++ b/verl/remote_backend/workers/__init__.py @@ -0,0 +1,6 @@ +"""Per-backend worker implementations. + +Each sub-package houses the worker class that `RemoteBackendTrainer` +picks up when a given remote backend is selected. The split mirrors +the per-backend split of adapters under :mod:`verl.workers.remote_client`. +""" diff --git a/verl/remote_backend/workers/arctic_rl/__init__.py b/verl/remote_backend/workers/arctic_rl/__init__.py new file mode 100644 index 00000000000..0dde2365a15 --- /dev/null +++ b/verl/remote_backend/workers/arctic_rl/__init__.py @@ -0,0 +1,9 @@ +"""Arctic-RL backend worker. + +Used when `trainer.remote_backend=arctic`. Pairs with the adapter at +:mod:`verl.workers.remote_client.arctic_rl`. +""" + +from verl.remote_backend.workers.arctic_rl.worker import ArcticRLActorRolloutRefWorker + +__all__ = ["ArcticRLActorRolloutRefWorker"] diff --git a/verl/remote_backend/worker.py b/verl/remote_backend/workers/arctic_rl/worker.py similarity index 80% rename from verl/remote_backend/worker.py rename to verl/remote_backend/workers/arctic_rl/worker.py index 22fee1e0493..e820d1da9d8 100644 --- a/verl/remote_backend/worker.py +++ b/verl/remote_backend/workers/arctic_rl/worker.py @@ -1,14 +1,17 @@ -"""Backend-agnostic forwarder worker. +"""Arctic-RL forwarder worker. +Per-backend worker that drives the Arctic adapter +(:class:`verl.workers.remote_client.arctic_rl.ArcticRLClientWrapper`). Owns single-controller dispatch annotations, FlopsCounter / MFU, and -metric aggregation. Every payload encoding decision (loss, wire format, -parallelism) lives behind ``RemoteBackend.compute_log_prob`` and -``RemoteBackend.update_actor`` — the worker stays backend-agnostic. +metric aggregation. Payload encoding (loss, wire format, parallelism) +lives on the adapter's ``compute_log_prob`` / ``update_actor`` — those +are Arctic-specific concrete methods, intentionally not on the +:class:`verl.remote_backend.RemoteBackend` ABC (which only enforces +lifecycle + checkpoint + weight-sync). Generic tensor / metric helpers live in -:mod:`verl.remote_backend.worker_utils` so backends and other modules can -import them without pulling in the Worker class or its dispatch -decorators. +:mod:`verl.remote_backend.worker_utils` so other per-backend worker +modules can reuse them without pulling in this class or its decorators. """ from __future__ import annotations @@ -21,6 +24,13 @@ from tensordict import TensorDict from verl.remote_backend.base import RemoteBackend, RemoteBackendRegistry + +# Eager import so the adapter registers with `RemoteBackendRegistry` in +# every process that loads this worker — including Ray child procs, which +# do not inherit the driver's import side-effects. The driver also +# imports this adapter explicitly (see `verl.trainer.main_ppo`). +from verl.workers.remote_client import arctic_rl # noqa: F401 + from verl.remote_backend.worker_utils import make_njt, normalize_backend_metrics from verl.single_controller.base import Worker from verl.single_controller.base.decorator import ( @@ -38,11 +48,25 @@ # ---------------------------------------------------------------------- # -# Generic forwarder worker +# Arctic-RL forwarder worker # ---------------------------------------------------------------------- # -class RemoteBackendActorRolloutRefWorker(Worker, DistProfilerExtension): - """CPU-only forwarder; assumes ``data["input_ids"]`` is nested-jagged.""" +class ArcticRLActorRolloutRefWorker(Worker, DistProfilerExtension): + """CPU-only forwarder for the Arctic backend. + + Assumes ``data["input_ids"]`` is nested-jagged. + + NOTE: per @zw0610's review on verl-project/verl#6422, the long-term + direction is to inherit from + :class:`verl.workers.engine_workers.ActorRolloutRefWorker` and only + override ``init_model`` + the compute/update methods so we don't + drift from the canonical worker. That requires decoupling + ``ActorRolloutRefWorker.__init__`` from megatron-specific config + fields (e.g. ``config.actor.megatron.router_replay``) which the + Arctic config tree doesn't carry. Left as a follow-up; the current + class parents (``Worker``, ``DistProfilerExtension``) match the + minimum required by the dispatch decorators below. + """ def __init__(self, config: DictConfig, role: str, **kwargs): Worker.__init__(self) @@ -58,14 +82,17 @@ def __init__(self, config: DictConfig, role: str, **kwargs): backend_name = self.main_config.trainer.get("remote_backend") if backend_name is None: raise ValueError( - "RemoteBackendActorRolloutRefWorker requires " + "ArcticRLActorRolloutRefWorker requires " "main_config.trainer.remote_backend to be set." ) - # `RemoteBackendRegistry.get` lazy-imports the adapter module - # (Ray child procs don't inherit the driver's import side-effects). - # `from_config(handle=...)` is the sole public constructor; the - # handle is what makes this a re-attach rather than a fresh init. + # Registry has no lazy `MODULES` table any more (per @zw0610): the + # adapter module is imported explicitly by `main_ppo.py` when the + # corresponding backend is selected, which decorates the class + # with `@RemoteBackendRegistry.register(...)`. Here we only look + # it up. `from_config(handle=...)` is the sole public constructor; + # the handle is what makes this a re-attach rather than a fresh + # init. backend_cls = RemoteBackendRegistry.get(backend_name) self.backend: RemoteBackend = backend_cls.from_config( self.main_config, handle=backend_handle diff --git a/verl/trainer/config/ppo_trainer.yaml b/verl/trainer/config/ppo_trainer.yaml index 60e89995938..f6487e1cbad 100644 --- a/verl/trainer/config/ppo_trainer.yaml +++ b/verl/trainer/config/ppo_trainer.yaml @@ -29,6 +29,12 @@ defaults: # Critic model config. - critic@critic: ${model_engine}_critic + # Remote backend config (optional). Opt in with `remote_backend=arctic` + # (loads `trainer/config/remote_backend/arctic.yaml` into the + # `remote_backend` key). Default leaves the field unset so the regular + # FSDP/Megatron actor/rollout/ref worker stack stays in charge. + - optional remote_backend@remote_backend: null + # legacy reward impl config, for backward compatibility - legacy_reward_impl @@ -318,38 +324,9 @@ ray_kwargs: timeline_json_file: null -# Per-backend configuration for `verl.remote_backend` adapters. The backend -# is selected by `trainer.remote_backend` (a string registered in -# `RemoteBackendRegistry`); each backend nests its config under -# `remote_backend.`. To plug in a new backend (e.g. "tinker"), -# add a sibling block here and register the adapter — no changes needed -# to the rest of this file or `main_ppo.py`. -remote_backend: - - # Arctic adapter (`verl.trainer.ppo.arctic_rl_client`). Active when - # `trainer.remote_backend=arctic`. - arctic: - - # whether to use colocate mode - colocate: False - - training_gpus: 1 - sampling_gpus: 1 - log_prob_gpus: 1 - sampling_tp_size: 1 - - use_zorro: False - - # ray or http: ray would be much faster for payload comms in the en-prem use case - comm_protocol: ray - - zero_optimization: - # 2 for optim + grad sharding - # 3 for optim + grad +weight sharding - stage: 0 - - offload_optimizer: - device: none - - offload_param: - device: none +# Per-backend `remote_backend` config lives in its own file under +# `trainer/config/remote_backend/.yaml` (loaded via the Hydra +# defaults entry above). To plug in a new backend (e.g. "tinker"), add +# `trainer/config/remote_backend/tinker.yaml` and opt in with +# `+remote_backend@remote_backend=tinker` — no changes needed to this +# file or `main_ppo.py`. diff --git a/verl/trainer/config/remote_backend/arctic.yaml b/verl/trainer/config/remote_backend/arctic.yaml new file mode 100644 index 00000000000..a80cbe2b90e --- /dev/null +++ b/verl/trainer/config/remote_backend/arctic.yaml @@ -0,0 +1,49 @@ +# Format checks enforced on CI: +# 1. Comments must appear above each field. +# 2. There must be a blank line between each field. +# 3. Inline comments (after a field on the same line) are not allowed. +# 4. Indentation level is respected for nested fields. + +# Per-backend block loaded into `config.remote_backend` via Hydra +# defaults (`+remote_backend@remote_backend=arctic`). Companion adapter: +# `verl.workers.remote_client.arctic_rl.ArcticRLClientWrapper`. +# Selected at runtime by `trainer.remote_backend=arctic`. Per +# @sfc-gh-truwase (PR #4): kept flat — the file name (`arctic.yaml`) +# already names the backend, so no extra top-level `arctic:` nesting. + +# whether to use colocate mode +colocate: True + +# number of GPUs allocated to the training engine +training_gpus: 1 + +# number of GPUs allocated to the sampling/inference engine +sampling_gpus: 1 + +# number of GPUs allocated to the log-prob engine; 0 means colocated with training +log_prob_gpus: 1 + +# tensor-parallel size for the sampling engine +sampling_tp_size: 1 + +# whether to use zorro packing for off-policy log-probs +use_zorro: False + +# ray or http: ray would be much faster for payload comms in the on-prem use case +comm_protocol: ray + +# DeepSpeed ZeRO-style sharding for the training engine +zero_optimization: + + # 2 for optim + grad sharding; 3 also shards weights + stage: 0 + + offload_optimizer: + + # `none` keeps the optimizer state on GPU; `cpu` offloads it + device: none + + offload_param: + + # `none` keeps params on GPU; `cpu` offloads them + device: none diff --git a/verl/trainer/main_ppo.py b/verl/trainer/main_ppo.py index 3517fbb154a..90ceb4aff18 100644 --- a/verl/trainer/main_ppo.py +++ b/verl/trainer/main_ppo.py @@ -134,16 +134,28 @@ def add_actor_rollout_worker(self, config): actor_rollout_cls = ActorRolloutRefWorker ray_worker_group_cls = RayWorkerGroup - if config.trainer.get("remote_backend"): - # Generic forwarder worker — payload encoding lives in - # the registered backend's `compute_log_prob` / - # `update_actor`. The trainer side lazy-imports the - # adapter via `RemoteBackendRegistry.get(name)`; nothing - # backend-specific is imported here. - from verl.remote_backend.worker import ( - RemoteBackendActorRolloutRefWorker as ActorRolloutRefWorker, - ) - actor_rollout_cls = ActorRolloutRefWorker + backend_name = config.trainer.get("remote_backend") + if backend_name: + # Per @zw0610: each remote backend ships its own + # per-backend worker (under `verl/remote_backend/workers/ + # /`) and adapter (under `verl/workers/ + # remote_client/`). We import the adapter explicitly here + # so it registers with `RemoteBackendRegistry`; the + # registry no longer carries a lazy `MODULES` table. + if backend_name == "arctic": + from verl.workers.remote_client import arctic_rl # noqa: F401 + from verl.remote_backend.workers.arctic_rl import ( + ArcticRLActorRolloutRefWorker, + ) + actor_rollout_cls = ArcticRLActorRolloutRefWorker + else: + raise ValueError( + f"Unknown trainer.remote_backend={backend_name!r}. " + "Known: 'arctic'. Plug in a new backend by adding " + "verl/workers/remote_client/.py + " + "verl/remote_backend/workers//worker.py and " + "wiring it here." + ) lora_rank = config.actor_rollout_ref.model.get("lora", {}).get("rank", 0) if lora_rank <= 0: diff --git a/verl/workers/remote_client/__init__.py b/verl/workers/remote_client/__init__.py new file mode 100644 index 00000000000..531eba1ae97 --- /dev/null +++ b/verl/workers/remote_client/__init__.py @@ -0,0 +1,7 @@ +"""Adapter clients for remote RL backends. + +Each module here implements a ``RemoteBackend`` (see :mod:`verl.remote_backend`) +that talks to an out-of-process training+rollout cluster owned by a +third-party library (e.g. ``arctic_training``). Importing the module +registers the adapter with :class:`verl.remote_backend.RemoteBackendRegistry`. +""" diff --git a/verl/trainer/ppo/arctic_rl_client.py b/verl/workers/remote_client/arctic_rl.py similarity index 95% rename from verl/trainer/ppo/arctic_rl_client.py rename to verl/workers/remote_client/arctic_rl.py index 21e2b8a6053..023a4920aaf 100644 --- a/verl/trainer/ppo/arctic_rl_client.py +++ b/verl/workers/remote_client/arctic_rl.py @@ -1,4 +1,3 @@ -import asyncio import os from typing import Any @@ -95,13 +94,13 @@ class ArcticRLClientWrapper(RemoteBackend): it via ``config.trainer.remote_backend = "arctic"``. """ - # Key under `main_config.remote_backend.` where this adapter's - # config lives. Matches the registered backend name above. - _BACKEND_CONFIG_KEY = "arctic" - def __init__(self, config, reconnect_job_config: dict = None, rl_server_state: ArcticRLRayServerState = None): self.config = config - self._backend_config = config.remote_backend[self._BACKEND_CONFIG_KEY] + # Per @sfc-gh-truwase (PR #4): the per-backend yaml + # (`trainer/config/remote_backend/arctic.yaml`) is loaded into + # `config.remote_backend` as a flat block — the file name already + # names the backend, so no extra `arctic:` nesting is needed. + self._backend_config = config.remote_backend self._client = None self.use_zorro = self._backend_config.use_zorro self.use_liger = self.config.actor_rollout_ref.model.use_liger @@ -159,7 +158,7 @@ def requires_single_forwarder(self) -> bool: return True # ------------------------------------------------------------------ # - # Core RL ops — called by `RemoteBackendActorRolloutRefWorker` + # Core RL ops — called by `ArcticRLActorRolloutRefWorker` # ------------------------------------------------------------------ # # Build Arctic's payload shape (dense padded batch, ``meta`` dict, # ``processing`` pipeline) and dispatch through the private @@ -422,16 +421,13 @@ async def generate(self, prompt_ids, sampling_params) -> list: async def _send_compute_ref_log_prob(self, payload: dict): payload["processing"] = {"post": ["compute_entropy_and_logprobs"], "loss_fn": None} - # Offload the blocking RPC to a worker thread so the forwarder's - # asyncio loop (and the Ray actor that owns it) stays responsive - # to concurrent calls. - response = await asyncio.to_thread(self._client.fwd_no_grad, payload, reference_model=True) + response = await self._client.fwd_no_grad(payload, reference_model=True) response["batch"]["log_probs"] = response["batch"].pop("logprobs") return response async def _send_compute_log_prob(self, payload: dict): payload["processing"] = {"post": ["compute_entropy_and_logprobs"], "loss_fn": None} - response = await asyncio.to_thread(self._client.fwd_no_grad, payload, reference_model=False) + response = await self._client.fwd_no_grad(payload, reference_model=False) response["batch"]["log_probs"] = response["batch"].pop("logprobs") return response @@ -455,8 +451,8 @@ def _left_pad(t: torch.Tensor, seq_len: int) -> torch.Tensor: payload["batch"][name] = _left_pad(payload["batch"][name], seq_len) payload["batch"]["loss_mask"] = payload["batch"]["response_mask"] - fwd_bwd_response = await asyncio.to_thread(self._client.fwd_bwd, payload) - step_response = await asyncio.to_thread(self._client.step) + fwd_bwd_response = await self._client.fwd_bwd(payload) + step_response = await self._client.step() step_response["metrics"].update(**fwd_bwd_response["metrics"]) return step_response diff --git a/verl/workers/rollout/remote_rollout/__init__.py b/verl/workers/rollout/remote_rollout/__init__.py new file mode 100644 index 00000000000..e3e92480adf --- /dev/null +++ b/verl/workers/rollout/remote_rollout/__init__.py @@ -0,0 +1,7 @@ +"""Rollout servers that delegate `generate` to a remote RL backend. + +Each sub-package wraps an in-process rollout server (e.g. +``vLLMHttpServer``) and routes prompts to the backend's `generate` +endpoint instead of running the model locally. Used by adapters that +co-train sampling and training in a separate cluster. +""" diff --git a/verl/workers/rollout/arctic_rollout/__init__.py b/verl/workers/rollout/remote_rollout/arctic_rollout/__init__.py similarity index 100% rename from verl/workers/rollout/arctic_rollout/__init__.py rename to verl/workers/rollout/remote_rollout/arctic_rollout/__init__.py diff --git a/verl/workers/rollout/arctic_rollout/arctic_rollout.py b/verl/workers/rollout/remote_rollout/arctic_rollout/arctic_rollout.py similarity index 99% rename from verl/workers/rollout/arctic_rollout/arctic_rollout.py rename to verl/workers/rollout/remote_rollout/arctic_rollout/arctic_rollout.py index 99927e795be..481713d93eb 100644 --- a/verl/workers/rollout/arctic_rollout/arctic_rollout.py +++ b/verl/workers/rollout/remote_rollout/arctic_rollout/arctic_rollout.py @@ -4,7 +4,7 @@ import argparse from typing import Any, Optional -from verl.trainer.ppo.arctic_rl_client import ArcticRLClientWrapper +from verl.workers.remote_client.arctic_rl import ArcticRLClientWrapper from collections.abc import AsyncGenerator import ray diff --git a/verl/workers/rollout/replica.py b/verl/workers/rollout/replica.py index 2557eb74d7a..89477f800af 100644 --- a/verl/workers/rollout/replica.py +++ b/verl/workers/rollout/replica.py @@ -349,10 +349,17 @@ def _load_trtllm(): return TRTLLMReplica def _load_arctic(): - from verl.workers.rollout.arctic_rollout.arctic_rollout import ArcticReplica - + from verl.workers.rollout.remote_rollout.arctic_rollout.arctic_rollout import ArcticReplica + return ArcticReplica + +# TODO(@zw0610): lazy-init each option for `RolloutReplicaRegistry`. For most +# cases there will be only 1 rollout backend used, but importing all of them +# eagerly could surface dependency conflicts (e.g. pulling in vLLM + SGLang + +# TRT-LLM simultaneously). The `_load_*` callables above already defer imports +# to call-time; the registration calls below could be moved into a per-backend +# `register_()` helper that the user calls explicitly at startup. # Register built-in types RolloutReplicaRegistry.register("vllm", _load_vllm) RolloutReplicaRegistry.register("sglang", _load_sglang)