Generic RemoteBackend abstraction + Arctic adapter (parallel to PR #1)

examples/arctic_rl/run_gsm8k_grpo_arl_zorro_yes.sh

@@ -0,0 +1,88 @@
+#!/bin/bash
+# Canonical GSM8K example for the Arctic RemoteBackend.
+#
+# Demonstrates the generic `verl.remote_backend` abstraction with the
+# Arctic adapter (`trainer.remote_backend=arctic`). Single-GPU, GRPO,
+# Qwen3-0.6B; intended as a quick convergence sanity check.
+
+set -x
+export PYTHONUNBUFFERED=1
+export HYDRA_FULL_ERROR=1
+export RAY_DEDUP_LOGS=0
+export HF_HUB_OFFLINE=1
+export HF_HOME=/checkpoint/huggingface
+export USE_ARCTIC_TRAINING_CLIENT=1
+export VLLM_BATCH_INVARIANT=1
+export VLLM_ENABLE_V1_MULTIPROCESSING=0
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
+export PYTHONPATH="$REPO_ROOT:$PYTHONPATH"
+
+gpu_name=$(nvidia-smi --query-gpu=gpu_name --format=csv,noheader -i 0)
+if   [[ $gpu_name == *"H200"* ]]; then flash_attention_v=flash_attention_3
+elif [[ $gpu_name == *"B200"* || $gpu_name == *"B300"* ]]; then flash_attention_v=flash_attention_2
+else flash_attention_v=flash_attention_2
+fi
+
+python3 -m verl.trainer.main_ppo \
+    algorithm.adv_estimator=grpo \
+    data.train_files=/code/shared/gsm8k/train.parquet \
+    data.val_files=/code/shared/gsm8k/test.parquet \
+    data.train_batch_size=16 \
+    data.max_prompt_length=512 \
+    data.max_response_length=1024 \
+    data.filter_overlong_prompts=True \
+    data.truncation='error' \
+    data.shuffle=False \
+    +data.seed=42 \
+    actor_rollout_ref.actor.data_loader_seed=42 \
+    reward.num_workers=1 \
+    actor_rollout_ref.rollout.agent.num_workers=1 \
+    actor_rollout_ref.model.path=Qwen/Qwen3-0.6B \
+    actor_rollout_ref.actor.optim.lr=1e-6 \
+    actor_rollout_ref.model.use_remove_padding=False \
+    actor_rollout_ref.actor.ppo_mini_batch_size=16 \
+    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \
+    actor_rollout_ref.actor.use_kl_loss=False \
+    actor_rollout_ref.actor.kl_loss_coef=0.001 \
+    actor_rollout_ref.actor.kl_loss_type=low_var_kl \
+    actor_rollout_ref.actor.entropy_coeff=0 \
+    actor_rollout_ref.model.enable_gradient_checkpointing=True \
+    +actor_rollout_ref.model.override_config.attn_implementation=$flash_attention_v \
+    actor_rollout_ref.actor.strategy=fsdp2 \
+    actor_rollout_ref.actor.fsdp_config.param_offload=False \
+    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
+    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \
+    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
+    actor_rollout_ref.rollout.name=arctic \
+    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \
+    actor_rollout_ref.rollout.enforce_eager=True \
+    actor_rollout_ref.rollout.n=5 \
+    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \
+    actor_rollout_ref.ref.fsdp_config.param_offload=False \
+    actor_rollout_ref.ref.strategy=fsdp2 \
+    algorithm.use_kl_in_reward=False \
+    trainer.use_legacy_worker_impl=disable \
+    trainer.remote_backend=arctic \
+    arctic_rl.colocate=False \
+    arctic_rl.training_gpus=1 \
+    arctic_rl.sampling_gpus=1 \
+    arctic_rl.log_prob_gpus=0 \
+    arctic_rl.tiled_logits_compute=True \
+    arctic_rl.zero_optimization.stage=2 \
+    arctic_rl.zero_optimization.offload_optimizer.device=none \
+    arctic_rl.zero_optimization.offload_param.device=none \
+    arctic_rl.use_zorro=True \
+    trainer.critic_warmup=0 \
+    trainer.logger="['console']" \
+    trainer.experiment_name=gsm8k_grpo_qwen3_0p6b_ngpu1_gbs16_rolln5_zorroTrue \
+    trainer.project_name=arctic_rl_gsm8k_public \
+    trainer.val_before_train=False \
+    trainer.n_gpus_per_node=1 \
+    trainer.nnodes=1 \
+    trainer.save_freq=-1 \
+    trainer.test_freq=-1 \
+    trainer.total_training_steps=80 \
+    trainer.total_epochs=15 \
+    "$@" 2>&1

verl/experimental/agent_loop/agent_loop.py

@@ -918,13 +918,14 @@ def __init__(
         worker_group: RayWorkerGroup = None,
         rollout_resource_pool: RayResourcePool = None,
         reward_loop_worker_handles: list[ray.actor.ActorHandle] = None,
-    ):
+        **kwargs,
+  ):
         self.config = config
         self.rollout_config, self.model_config = _get_rollout_and_model_config(config)
         self.worker_group = worker_group
         self.rollout_resource_pool = rollout_resource_pool
         self.reward_loop_worker_handles = reward_loop_worker_handles
-
+        self.kwargs = kwargs
         assert worker_group is not None or self.rollout_config.nnodes > 0, "nnodes must be > 0 in standalone mode"
 
         # for recipe to change
@@ -941,9 +942,10 @@ async def create(
         worker_group: RayWorkerGroup = None,
         rollout_resource_pool: RayResourcePool = None,
         reward_loop_worker_handles: list[ray.actor.ActorHandle] = None,
+        **kwargs,
     ):
         """Create agent loop manager."""
-        instance = cls(config, worker_group, rollout_resource_pool, reward_loop_worker_handles)
+        instance = cls(config, worker_group, rollout_resource_pool, reward_loop_worker_handles, **kwargs)
         await instance._initialize_llm_servers()
         await instance._init_global_load_balancer()
         await instance._init_agent_loop_workers()
@@ -968,6 +970,7 @@ async def _initialize_llm_servers(self):
                 config=self.rollout_config,
                 model_config=self.model_config,
                 gpus_per_node=self.rollout_config.n_gpus_per_node,
+                **self.kwargs,
             )
             for replica_rank in range(num_replicas)
         ]

verl/remote_backend/__init__.py

@@ -0,0 +1,28 @@
+"""Generic remote-backend abstraction for verl.
+
+Lets verl drive an out-of-process RL backend (training + rollout +
+log-prob + checkpoint) that owns its own GPUs. Verl talks to a CPU-only
+forwarder worker group; the forwarder forwards every dispatched call to
+a :class:`RemoteBackend` implementation behind a Ray actor (or any other
+RPC the backend prefers).
+
+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.
+
+See :mod:`verl.trainer.ppo.arctic_rl_client` for a reference adapter.
+"""
+
+from verl.remote_backend.base import RemoteBackend, RemoteBackendRegistry
+
+__all__ = ["RemoteBackend", "RemoteBackendRegistry"]

verl/remote_backend/base.py

@@ -0,0 +1,238 @@
+"""`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.
+
+What the abstraction is responsible for (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``.
+* 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.
+"""
+
+from __future__ import annotations
+
+import abc
+from typing import Any, Callable
+
+import torch
+from omegaconf import DictConfig
+
+
+class RemoteBackend(abc.ABC):
+    """Out-of-process RL backend that owns its own GPUs.
+
+    Created once on the driver by ``RemoteBackendTrainer`` (via
+    ``from_config(main_config)``); re-attached inside every forwarder
+    worker via ``from_config(main_config, handle=...)``.
+    """
+
+    # ------------------------------------------------------------------ #
+    # Lifecycle
+    # ------------------------------------------------------------------ #
+
+    @classmethod
+    @abc.abstractmethod
+    def from_config(
+        cls,
+        main_config: DictConfig,
+        *,
+        handle: dict[str, Any] | None = None,
+    ) -> "RemoteBackend":
+        """Sole public constructor.
+
+        Args:
+            main_config: the full verl config tree. Backend-specific knobs
+                live under ``main_config.remote_backend.<name>``; backends
+                MUST NOT read outside their own namespace plus the small set
+                of standard fields under ``main_config.{trainer, data,
+                actor_rollout_ref}``.
+            handle: when supplied, re-attach to an existing backend
+                instance described by a previous
+                :meth:`reconnect_handle` (used by forwarder workers /
+                rollout replicas that share the driver-side backend
+                instead of creating a second one). When ``None``,
+                create a fresh backend on the driver.
+        """
+
+    @abc.abstractmethod
+    def reconnect_handle(self) -> dict[str, Any]:
+        """A serializable handle that, when passed back to
+        :meth:`from_config` as ``handle=...``, yields a reference to
+        *this* backend.
+
+        Typically contains a Ray actor handle and a small config blob.
+        ``RemoteBackendTrainer`` puts this dict into ``wg_kwargs`` so each
+        forwarder worker can re-attach.
+        """
+
+    @abc.abstractmethod
+    def destroy(self) -> None:
+        """Tear down the backend cleanly. Must be idempotent.
+
+        Called from ``RemoteBackendTrainer.destroy()`` after ``fit()``.
+        """
+
+    # ------------------------------------------------------------------ #
+    # Core RL operations called by `RemoteBackendActorRolloutRefWorker`
+    # ------------------------------------------------------------------ #
+    # 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
+    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``).
+
+        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
+    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.
+
+        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]:
+        """Sync trained weights from the training engine to the rollout
+        engine. May be a no-op for colocated backends.
+        """
+
+    @abc.abstractmethod
+    def save_checkpoint(self) -> dict[str, Any]:
+        """Persist current model + optimizer state."""
+
+    # ------------------------------------------------------------------ #
+    # Parallelism contract
+    # ------------------------------------------------------------------ #
+
+    @abc.abstractmethod
+    def requires_single_forwarder(self) -> bool:
+        """Whether ``RemoteBackendTrainer`` should assert
+        ``n_gpus_per_node × nnodes == 1`` and a single rollout replica.
+
+        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.
+
+        Returning ``True`` enables the assert; returning ``False`` opts
+        out (the backend takes responsibility for validating its own
+        worker-group config).
+        """
+
+
+class RemoteBackendRegistry:
+    """Process-wide registry of name → :class:`RemoteBackend` class.
+
+    Backends register themselves at import time::
+
+        @RemoteBackendRegistry.register("arctic")
+        class ArcticBackend(RemoteBackend): ...
+
+    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")``.
+    """
+
+    _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]:
+            existing = cls._backends.get(name)
+            if existing is not None and existing is not backend_cls:
+                raise ValueError(
+                    f"Remote backend name '{name}' is already registered to "
+                    f"{existing!r}; cannot re-register to {backend_cls!r}."
+                )
+            cls._backends[name] = backend_cls
+            return backend_cls
+
+        return _decorator
+
+    @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)
+        return cls._backends[name]
+
+    @classmethod
+    def create(cls, name: str, main_config: DictConfig) -> RemoteBackend:
+        return cls.get(name).from_config(main_config)
+
+    @classmethod
+    def list(cls) -> list[str]:
+        return sorted(cls._backends)