Restructure remote_backend per @zw0610 review (stacked on #3)

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 \

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="<name>"`` 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/<backend_name>/`` — 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

verl/remote_backend/base.py

@@ -1,30 +1,34 @@
 """`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
 
 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

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

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`.
+"""

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"]

verl/remote_backend/worker.py → 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