diff --git a/verl/experimental/agent_loop/agent_loop.py b/verl/experimental/agent_loop/agent_loop.py index 8879960f128..4f251594bad 100644 --- a/verl/experimental/agent_loop/agent_loop.py +++ b/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) ] diff --git a/verl/remote_backend/__init__.py b/verl/remote_backend/__init__.py new file mode 100644 index 00000000000..4b30c5791d5 --- /dev/null +++ b/verl/remote_backend/__init__.py @@ -0,0 +1,48 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""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``) -- 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 -> backend class + registry (populated by the adapter's ``@register`` decorator) plus a + parallel ``register_worker`` / ``get_worker`` slot for the matching + ActorRollout forwarder worker class. Populated by adapter packages + via the ``VERL_USE_EXTERNAL_MODULES`` hook. +* :class:`RemoteBackendTrainer` (``trainer.py``) -- ``RayPPOTrainer`` + subclass that creates the backend on the driver and threads its + reconnect handle to every worker. +* ``worker_utils.py`` -- small generic tensor / metric helpers shared + across per-backend workers, which live in the adapter packages. + +Verl-core carries no concrete backends. The reference implementation for +the ABC lives in ``arctic_platform.integrations.verl`` (Arctic RL); +plug it in with +``VERL_USE_EXTERNAL_MODULES=arctic_platform.integrations.verl.register``. +""" + +from verl.remote_backend.base import RemoteBackend, RemoteBackendRegistry + +__all__ = ["RemoteBackend", "RemoteBackendRegistry"] diff --git a/verl/remote_backend/base.py b/verl/remote_backend/base.py new file mode 100644 index 00000000000..3c632f2983a --- /dev/null +++ b/verl/remote_backend/base.py @@ -0,0 +1,271 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""`RemoteBackend` ABC + `RemoteBackendRegistry`. + +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 ABC owns (what verl drives): + +* Lifecycle: ``from_config`` (sole constructor, takes an optional + ``handle=`` for re-attach) / ``reconnect_handle`` / ``destroy``. +* 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 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. + +Registration model +------------------ + +Backends live in their own packages -- verl core carries none -- and +are wired in via the ``VERL_USE_EXTERNAL_MODULES`` hook. Users set:: + + VERL_USE_EXTERNAL_MODULES=my_pkg.integrations.verl.register + +That module's top level: + +1. Imports the adapter module, which at class-definition time is + decorated with ``@RemoteBackendRegistry.register("")`` and thus + inserts itself into the class registry as a side effect. +2. Calls :meth:`RemoteBackendRegistry.register_worker` with the + backend's ActorRollout(Ref) forwarder worker class. The trainer + ``main_ppo`` reads this back via + :meth:`RemoteBackendRegistry.get_worker` to select + ``actor_rollout_cls`` at bootstrap time, without hard-coding a + per-backend if-branch. +3. Registers the rollout replica class with + :class:`verl.workers.rollout.replica.RolloutReplicaRegistry` (which + uses its own lazy-loader signature for vLLM/SGLang/... parity). +""" + +from __future__ import annotations + +import abc +from typing import Any, Callable + +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.``; 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()``. + """ + + # ------------------------------------------------------------------ # + # Weight sync + checkpoint (called from ONE_TO_ALL worker hooks). + # ------------------------------------------------------------------ # + + @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 + async def save_checkpoint(self) -> dict[str, Any]: + """Persist current model + optimizer state. + + ``async`` so the underlying RPC (typically a Ray actor call) can be + awaited without blocking the forwarder's event loop. + """ + + # ------------------------------------------------------------------ # + # 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 + 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 + worker-group config). + """ + + +class RemoteBackendRegistry: + """Process-wide registry of name -> (:class:`RemoteBackend` class, + ActorRollout forwarder worker class). + + Backend classes register themselves via the + ``@RemoteBackendRegistry.register(name)`` decorator at class + definition time. Forwarder worker classes are registered + imperatively by the same plugin's entry-point module, via + :meth:`register_worker`; that keeps the decorator on the backend + class simple, and lets the worker module retain its own eager + imports without having to know about registry mechanics. + + 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]] = {} + _worker_loaders: dict[str, Callable[[], type]] = {} + _resolved_workers: dict[str, type] = {} + + # -- Backend class registry ------------------------------------------- + + @classmethod + def register(cls, name: str) -> Callable[[type[RemoteBackend]], type[RemoteBackend]]: + """Decorator: register the decorated class as backend ``name``. + + Duplicate registrations of the same name with the identical class + object are a no-op (so a re-import of the plugin module during + test teardown / hot-reload doesn't blow up); different classes + under the same name raise, so the collision surfaces at import + time. + """ + + 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]: + if name not in cls._backends: + raise KeyError( + f"Unknown remote backend '{name}'. Registered: " + f"{sorted(cls._backends)}. Wire the adapter package in via " + "VERL_USE_EXTERNAL_MODULES=.integrations.verl.register " + "before starting verl." + ) + 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) + + # -- ActorRollout forwarder worker registry --------------------------- + + @classmethod + def register_worker(cls, name: str, loader: Callable[[], type]) -> None: + """Register a lazy loader for the ActorRollout forwarder worker + class matching backend ``name``. + + ``loader`` is a zero-arg callable returning the concrete worker + class; it is invoked once on the driver at first + :meth:`get_worker` and its result cached. Keeps this symmetric + with :class:`verl.workers.rollout.replica.RolloutReplicaRegistry` + (also lazy-loader) so an adapter plugin's ``register.py`` never + forces an import of vLLM / DeepSpeed / tensordict just to wire a + name into the registry. + + Duplicate registrations of the same name with the same loader + object are a no-op; different loaders raise, so the collision + surfaces at import time. + """ + existing = cls._worker_loaders.get(name) + if existing is not None and existing is not loader: + raise ValueError( + f"Remote backend '{name}' worker loader already registered to " + f"{existing!r}; cannot re-register to {loader!r}." + ) + cls._worker_loaders[name] = loader + + @classmethod + def get_worker(cls, name: str) -> type | None: + """Return the ActorRollout forwarder worker class for ``name``, + or ``None`` if the backend didn't register one (in which case + ``main_ppo`` falls back to verl's stock ``ActorRolloutRefWorker`` + -- only correct for backends whose payload/loss shape matches + the stock worker). + """ + if name in cls._resolved_workers: + return cls._resolved_workers[name] + loader = cls._worker_loaders.get(name) + if loader is None: + return None + cls._resolved_workers[name] = loader() + return cls._resolved_workers[name] diff --git a/verl/remote_backend/trainer.py b/verl/remote_backend/trainer.py new file mode 100644 index 00000000000..ed090ee52d6 --- /dev/null +++ b/verl/remote_backend/trainer.py @@ -0,0 +1,148 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""`RayPPOTrainer` subclass that drives a `RemoteBackend` instead of an +in-process engine. Resolves the backend by name, enforces single-forwarder +constraints when required, threads the reconnect handle into `wg_kwargs`. +""" + +from __future__ import annotations + +from typing import Optional + +from torch.utils.data import Dataset, Sampler + +from verl.remote_backend.base import RemoteBackend, RemoteBackendRegistry +from verl.single_controller.ray import RayWorkerGroup, ResourcePoolManager +from verl.trainer.ppo.ray_trainer import RayPPOTrainer +from verl.trainer.ppo.utils import Role, WorkerType + + +class RemoteBackendTrainer(RayPPOTrainer): + """PPO trainer that delegates train/rollout/log-prob/sync/ckpt to a + :class:`RemoteBackend`. + """ + + def __init__( + self, + config, + tokenizer, + role_worker_mapping: dict[Role, WorkerType], + resource_pool_manager: ResourcePoolManager, + ray_worker_group_cls: type[RayWorkerGroup] = RayWorkerGroup, + processor=None, + train_dataset: Optional[Dataset] = None, + val_dataset: Optional[Dataset] = None, + collate_fn=None, + train_sampler: Optional[Sampler] = None, + device_name=None, + backend: Optional[RemoteBackend] = None, + ): + super().__init__( + config=config, + tokenizer=tokenizer, + processor=processor, + role_worker_mapping=role_worker_mapping, + resource_pool_manager=resource_pool_manager, + ray_worker_group_cls=ray_worker_group_cls, + train_dataset=train_dataset, + val_dataset=val_dataset, + collate_fn=collate_fn, + train_sampler=train_sampler, + device_name=device_name, + ) + + if backend is None: + backend_name = config.trainer.get("remote_backend") + if not backend_name: + raise ValueError( + "RemoteBackendTrainer requires trainer.remote_backend " + "to be set (e.g. trainer.remote_backend=arctic)." + ) + backend = RemoteBackendRegistry.create(backend_name, config) + self.backend: RemoteBackend = backend + + self._enforce_single_forwarder_if_required() + + self.use_gpu = False + self.wg_kwargs["main_config"] = config + self.wg_kwargs["backend_handle"] = self.backend.reconnect_handle() + + # ------------------------------------------------------------------ # + # Helpers + # ------------------------------------------------------------------ # + + def _enforce_single_forwarder_if_required(self) -> None: + """Honour the backend's :meth:`RemoteBackend.requires_single_forwarder` + declaration: assert ``n_gpus_per_node × nnodes == 1`` and a single + rollout replica. See the ABC docstring for the rationale.""" + if not self.backend.requires_single_forwarder(): + return + + n_gpus = self.config.trainer.n_gpus_per_node + nnodes = self.config.trainer.nnodes + if n_gpus * nnodes != 1: + raise AssertionError( + f"Remote backend {type(self.backend).__name__!r} requires a " + "single forwarder worker, but the verl-side worker group " + f"would have n_gpus_per_node={n_gpus} × nnodes={nnodes} = " + f"{n_gpus * nnodes} workers. With more than one, ONE_TO_ALL " + "calls duplicate against the backend and mesh-dispatched " + "calls fragment the global batch. Set " + "trainer.n_gpus_per_node=1 and trainer.nnodes=1 (the backend " + "owns its own GPUs and parallelism) or override " + "RemoteBackend.requires_single_forwarder() in your backend " + "and validate the config yourself." + ) + + rollout_replicas = self._inferred_rollout_replica_count() + if rollout_replicas != 1: + raise AssertionError( + f"Remote backend {type(self.backend).__name__!r} requires a " + f"single rollout replica until the multi-replica path is " + f"validated; got {rollout_replicas}. Set " + "actor_rollout_ref.rollout.agent.num_workers=1 or override " + "RemoteBackend.requires_single_forwarder() in your backend." + ) + + def _inferred_rollout_replica_count(self) -> int: + """Best-effort introspection of the configured rollout-replica count. + + Different rollout backends use different config keys; we look at the + few we know about and default to 1. + """ + rollout_cfg = self.config.actor_rollout_ref.rollout + for key_path in (("agent", "num_workers"), ("num_workers",), ("replicas",)): + cur = rollout_cfg + for part in key_path: + cur = cur.get(part) if hasattr(cur, "get") else None + if cur is None: + break + if isinstance(cur, int): + return cur + return 1 + + # ------------------------------------------------------------------ # + # Lifecycle + # ------------------------------------------------------------------ # + + def destroy(self) -> None: + if getattr(self, "backend", None) is not None: + # `RemoteBackend.destroy` is async (so adapters can await + # remote RPCs without blocking event loops). The trainer's + # `destroy` is called from synchronous shutdown paths, so + # bridge with `asyncio.run`. + import asyncio + + asyncio.run(self.backend.destroy()) + self.backend = None diff --git a/verl/remote_backend/worker_utils.py b/verl/remote_backend/worker_utils.py new file mode 100644 index 00000000000..c92c53eb4c1 --- /dev/null +++ b/verl/remote_backend/worker_utils.py @@ -0,0 +1,93 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Backend-agnostic tensor + metric helpers shared across per-backend +forwarder workers. + +Per-backend forwarder workers ship in adapter packages (verl core carries +none). Any such worker can pull generic nested-jagged reconstruction / +metric normalization from this module without importing another +backend's Worker class or its single-controller dispatch decorators. + +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 + ``data["input_ids"]``. +* :func:`normalize_backend_metrics` — coerce whatever shape a backend + put in its ``metrics`` dict (a :class:`Metric`, a list of them, a list + of bare scalars, or a bare scalar) into the canonical verl + :class:`Metric` form (or a pass-through list when the trainer needs + to aggregate later). +""" + +from __future__ import annotations + +from typing import Any + +import torch +from tensordict import TensorDict +from torch import Tensor + +from verl.utils.metric import AggregationType, Metric + + +def make_njt(data: TensorDict, tensor: Tensor) -> Tensor: + """Reconstruct a nested-jagged tensor from a dense ``[B, L]`` tensor + and the offsets carried in ``data["input_ids"]``.""" + cu_seqlens = data["input_ids"].offsets() + seq_lengths = cu_seqlens.diff() + starts = data["attention_mask"].long().argmax(dim=1) + pieces = [tensor[b, starts[b].item() : starts[b].item() + seq_lengths[b].item()] for b in range(tensor.shape[0])] + flat = torch.cat(pieces, dim=0) + return torch.nested.nested_tensor_from_jagged(flat, cu_seqlens) + + +def shift_nested_response_aligned_to_predict_next(njt: Tensor) -> Tensor: + """Convert a response-aligned nested tensor back to the legacy + "predict-next" convention by left-shifting each sequence by one. + + The Arctic zorro fast path emits model output already aligned with the + response token at each slot (slot i holds the log prob of token[i]). Shift it + left so slot i holds log P(token[i+1]); ``no_padding_2_padding`` can then apply + a single uniform ``shift=-1`` slice for both the zorro and standard paths, + instead of branching the slice on a per-call ``shift`` argument.""" + v = njt.values() + v_new = torch.empty_like(v) + v_new[:-1] = v[1:] + v_new[-1] = 0 # trailing slot of last seq; never read by the shift=-1 slice + return torch.nested.nested_tensor_from_jagged(v_new, njt.offsets()) + + +def normalize_backend_metrics(metrics: dict[str, Any]) -> dict[str, Any]: + """Normalise backend metrics into verl's :class:`Metric` type. + + Shapes handled: a `Metric` (pass through), a `list[Metric]` (call + `Metric.aggregate_dp`), a `list[scalar]` (pass through; trainer + aggregates), or a bare scalar (wrap as MEAN). + """ + out: dict[str, Any] = {} + for key, val in metrics.items(): + if isinstance(val, Metric): + out[key] = val + elif isinstance(val, list): + if val and isinstance(val[0], Metric): + out[key] = Metric.aggregate_dp(val) + else: + out[key] = val + else: + out[key] = Metric(value=val, aggregation=AggregationType.MEAN) + return out diff --git a/verl/single_controller/ray/base.py b/verl/single_controller/ray/base.py index 2f6ee47064f..ab5e4b0177f 100644 --- a/verl/single_controller/ray/base.py +++ b/verl/single_controller/ray/base.py @@ -187,8 +187,9 @@ class ResourcePoolManager: resource_pool_spec: dict[str, list[int]] mapping: dict[int, str] resource_pool_dict: dict[str, RayResourcePool] = field(default_factory=dict) + gpu_resource_pool_dict: dict[str, RayResourcePool] = field(default_factory=dict) - def create_resource_pool(self): + def create_resource_pool(self, use_gpu: bool = True): """Create Ray resource pools for distributed training. Initializes resource pools based on the resource pool specification, @@ -202,10 +203,11 @@ def create_resource_pool(self): # For Megatron backend, we recommend using max_colocate_count>1 # that can utilize different WorkerGroup for differnt models resource_pool = RayResourcePool( - process_on_nodes=process_on_nodes, use_gpu=True, max_colocate_count=3, name_prefix=resource_pool_name + process_on_nodes=process_on_nodes, use_gpu=use_gpu, max_colocate_count=3, name_prefix=resource_pool_name ) self.resource_pool_dict[resource_pool_name] = resource_pool - + if use_gpu: + self.gpu_resource_pool_dict[resource_pool_name] = resource_pool self._check_resource_available() def get_resource_pool(self, role) -> RayResourcePool: @@ -214,7 +216,12 @@ def get_resource_pool(self, role) -> RayResourcePool: def get_n_gpus(self) -> int: """Get the number of gpus in this cluster.""" - return sum([n_gpus for process_on_nodes in self.resource_pool_spec.values() for n_gpus in process_on_nodes]) + process_on_gpu_nodes = [ + process_on_nodes + for pool_name, process_on_nodes in self.resource_pool_spec.items() + if pool_name in self.gpu_resource_pool_dict + ] + return sum([n_gpus for process_on_nodes in process_on_gpu_nodes for n_gpus in process_on_nodes]) def _check_resource_available(self): """Check if the resource pool can be satisfied in this ray cluster.""" @@ -226,9 +233,7 @@ def _check_resource_available(self): # check total required gpus can be satisfied total_available_gpus = sum(node_available_gpus.values()) - total_required_gpus = sum( - [n_gpus for process_on_nodes in self.resource_pool_spec.values() for n_gpus in process_on_nodes] - ) + total_required_gpus = self.get_n_gpus() if total_available_gpus < total_required_gpus: raise ValueError( f"Total available GPUs {total_available_gpus} is less than total desired GPUs {total_required_gpus}" diff --git a/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml b/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml index 9459721df44..ebf30773fa5 100644 --- a/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml @@ -645,6 +645,7 @@ trainer: ray_wait_register_center_timeout: 300 device: cuda use_legacy_worker_impl: auto + remote_backend: null global_profiler: _target_: verl.utils.profiler.ProfilerConfig tool: null diff --git a/verl/trainer/config/_generated_ppo_trainer.yaml b/verl/trainer/config/_generated_ppo_trainer.yaml index 79782a09b85..0ab981b9553 100644 --- a/verl/trainer/config/_generated_ppo_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_trainer.yaml @@ -687,6 +687,7 @@ trainer: ray_wait_register_center_timeout: 300 device: cuda use_legacy_worker_impl: auto + remote_backend: null global_profiler: _target_: verl.utils.profiler.ProfilerConfig tool: null diff --git a/verl/trainer/config/_generated_ppo_veomni_trainer.yaml b/verl/trainer/config/_generated_ppo_veomni_trainer.yaml index b02e95348c4..ab06976c736 100644 --- a/verl/trainer/config/_generated_ppo_veomni_trainer.yaml +++ b/verl/trainer/config/_generated_ppo_veomni_trainer.yaml @@ -622,6 +622,7 @@ trainer: ray_wait_register_center_timeout: 300 device: cuda use_legacy_worker_impl: auto + remote_backend: null global_profiler: _target_: verl.utils.profiler.ProfilerConfig tool: null diff --git a/verl/trainer/config/ppo_trainer.yaml b/verl/trainer/config/ppo_trainer.yaml index fd9b59862ae..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 @@ -203,6 +209,12 @@ trainer: # mode: "auto", "enable", or "disable" use_legacy_worker_impl: auto + # Remote-backend selector. Set to a registered backend name + # (e.g. "arctic") to route training/rollout/log-prob/sync/ckpt through + # `verl.remote_backend.RemoteBackend` instead of an in-process engine. + # Leave null for the standard in-process verl path. + remote_backend: null + # profiler configs global_profiler: @@ -310,3 +322,11 @@ ray_kwargs: # Path to save Ray timeline JSON for performance profiling timeline_json_file: null + + +# 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/main_ppo.py b/verl/trainer/main_ppo.py index 2c84374d245..aa06a92fdf7 100644 --- a/verl/trainer/main_ppo.py +++ b/verl/trainer/main_ppo.py @@ -134,6 +134,28 @@ def add_actor_rollout_worker(self, config): actor_rollout_cls = ActorRolloutRefWorker ray_worker_group_cls = RayWorkerGroup + backend_name = config.trainer.get("remote_backend") + if backend_name: + # The adapter package is loaded during `import verl` via + # VERL_USE_EXTERNAL_MODULES; its side effects register + # both the RemoteBackend class (decorator on the adapter) + # and the matching ActorRollout forwarder worker class + # (RemoteBackendRegistry.register_worker call in the + # plugin's register.py). Here we only look the worker up. + from verl.remote_backend.base import RemoteBackendRegistry + + worker_cls = RemoteBackendRegistry.get_worker(backend_name) + if worker_cls is None: + raise ValueError( + f"Remote backend {backend_name!r} did not register an " + "ActorRollout forwarder worker. Ensure the adapter " + "package's register.py calls " + "`RemoteBackendRegistry.register_worker(name, cls)`, " + "and that VERL_USE_EXTERNAL_MODULES points at it. " + f"Registered backends: {RemoteBackendRegistry.list()!r}." + ) + actor_rollout_cls = worker_cls + lora_rank = config.actor_rollout_ref.model.get("lora", {}).get("rank", 0) if lora_rank <= 0: lora_rank = config.actor_rollout_ref.model.get("lora_rank", 0) @@ -339,8 +361,16 @@ def run(self, config): ) train_sampler = create_rl_sampler(config.data, train_dataset) - # Initialize the PPO trainer. - trainer = RayPPOTrainer( + # Pick the trainer: `RemoteBackendTrainer` when a remote backend + # is selected via `trainer.remote_backend = ""`, otherwise + # the standard in-process `RayPPOTrainer`. + if config.trainer.get("remote_backend"): + from verl.remote_backend.trainer import RemoteBackendTrainer + + ppo_trainer_cls = RemoteBackendTrainer + else: + ppo_trainer_cls = RayPPOTrainer + trainer = ppo_trainer_cls( config=config, tokenizer=tokenizer, processor=processor, @@ -356,7 +386,12 @@ def run(self, config): trainer.init_workers() # Start the training process. - trainer.fit() + try: + trainer.fit() + finally: + # Ensure remote services shutdown gracefully + if hasattr(trainer, "destroy"): + trainer.destroy() def create_rl_dataset(data_paths, data_config, tokenizer, processor, is_train=True, max_samples: int = -1): diff --git a/verl/trainer/ppo/ray_trainer.py b/verl/trainer/ppo/ray_trainer.py index e178ffc143d..478bfd07908 100644 --- a/verl/trainer/ppo/ray_trainer.py +++ b/verl/trainer/ppo/ray_trainer.py @@ -309,6 +309,9 @@ def __init__( self.checkpoint_manager = None + self.wg_kwargs = {} + self.use_gpu = True + def _create_dataloader(self, train_dataset, val_dataset, collate_fn, train_sampler: Optional[Sampler]): """ Creates the train and validation dataloaders. @@ -682,7 +685,7 @@ def init_workers(self): 1. Ray resource pools from configuration 2. Worker groups for each role (actor, critic, etc.) """ - self.resource_pool_manager.create_resource_pool() + self.resource_pool_manager.create_resource_pool(use_gpu=self.use_gpu) self.resource_pool_to_cls = {pool: {} for pool in self.resource_pool_manager.resource_pool_dict.values()} @@ -694,6 +697,7 @@ def init_workers(self): cls=self.role_worker_mapping[actor_role], config=self.config.actor_rollout_ref, role=str(actor_role), + **self.wg_kwargs, ) self.resource_pool_to_cls[actor_rollout_resource_pool][str(actor_role)] = actor_rollout_cls else: @@ -840,6 +844,7 @@ def init_workers(self): worker_group=self.actor_rollout_wg, rollout_resource_pool=actor_rollout_resource_pool, reward_loop_worker_handles=reward_loop_worker_handles, + **self.wg_kwargs, ) checkpoint_engine_config = omega_conf_to_dataclass(self.config.actor_rollout_ref.rollout.checkpoint_engine) self.checkpoint_manager = CheckpointEngineManager( @@ -1589,7 +1594,8 @@ def fit(self): metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw)) # TODO: implement actual tflpo and theoretical tflpo n_gpus = self.resource_pool_manager.get_n_gpus() - metrics.update(compute_throughout_metrics(batch=batch, timing_raw=timing_raw, n_gpus=n_gpus)) + # To support serverless/tinker-like training, we need to support 0 GPUs training + metrics.update(compute_throughout_metrics(batch=batch, timing_raw=timing_raw, n_gpus=max(n_gpus, 1))) # compute variance proxy metrics gradient_norm = metrics.get("actor/grad_norm", None) metrics.update(compute_variance_proxy_metrics(batch=batch, gradient_norm=gradient_norm)) diff --git a/verl/utils/profiler/profile.py b/verl/utils/profiler/profile.py index 8e3145a66bb..f871a5da501 100644 --- a/verl/utils/profiler/profile.py +++ b/verl/utils/profiler/profile.py @@ -13,6 +13,7 @@ # limitations under the License. import functools +import inspect from typing import Callable, Optional from ..memory_utils import MemorySnapshotSampler, enable_memory_visualize @@ -161,6 +162,15 @@ def annotate( **kwargs_outer, ) -> Callable: def decorator(func): + if inspect.iscoroutinefunction(func): + + @functools.wraps(func) + async def async_wrapper(self_instance, *args, **kwargs_inner): + # Nested profiler annotate paths assume sync callables; async methods run uninstrumented. + return await func(self_instance, *args, **kwargs_inner) + + return async_wrapper + @functools.wraps(func) def wrapper(self_instance, *args, **kwargs_inner): profiler = getattr(self_instance, "profiler", None) diff --git a/verl/workers/engine_workers.py b/verl/workers/engine_workers.py index d0c065e4dfd..d7d458f4264 100644 --- a/verl/workers/engine_workers.py +++ b/verl/workers/engine_workers.py @@ -666,8 +666,9 @@ async def update_weights(self, global_steps: int = None): log_gpu_memory_usage("After update_weights", logger=logger) - # 3. offload model to cpu - self.actor.engine.to("cpu", model=True, optimizer=False, grad=False) + # 3. offload model to cpu (only when param offload is enabled, else params stay on CPU for the next forward) + if self.actor.engine.is_param_offload_enabled: + self.actor.engine.to("cpu", model=True, optimizer=False, grad=False) aggressive_empty_cache(force_sync=True) # 4. resume kv_cache diff --git a/verl/workers/rollout/replica.py b/verl/workers/rollout/replica.py index 969c6208083..45869436025 100644 --- a/verl/workers/rollout/replica.py +++ b/verl/workers/rollout/replica.py @@ -349,10 +349,18 @@ def _load_trtllm(): return TRTLLMReplica +# 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) RolloutReplicaRegistry.register("trtllm", _load_trtllm) +# Out-of-tree rollout backends (e.g. "arctic") register themselves via the +# VERL_USE_EXTERNAL_MODULES hook; see arctic_platform/integrations/verl. # Original function for backward compatibility