[refactor] Rename arctic_rl config namespace to generic remote_backend.arctic

examples/arctic_rl/run_gsm8k_grpo_arl_zorro_yes.sh

@@ -65,15 +65,14 @@ python3 -m verl.trainer.main_ppo \
     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 \
+    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 \
     trainer.critic_warmup=0 \
     trainer.logger="['console']" \
     trainer.experiment_name=gsm8k_grpo_qwen3_0p6b_ngpu1_gbs16_rolln5_zorroTrue \

verl/remote_backend/base.py

@@ -91,7 +91,7 @@ def destroy(self) -> None:
     # below this layer.
 
     @abc.abstractmethod
-    def compute_log_prob(
+    async def compute_log_prob(
         self,
         data,
         *,
@@ -104,6 +104,9 @@ def compute_log_prob(
         """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
@@ -113,7 +116,7 @@ def compute_log_prob(
         """
 
     @abc.abstractmethod
-    def update_actor(
+    async def update_actor(
         self,
         data,
         *,
@@ -124,6 +127,9 @@ def update_actor(
     ) -> 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
@@ -147,8 +153,12 @@ async def update_weights(self) -> dict[str, Any]:
         """
 
     @abc.abstractmethod
-    def save_checkpoint(self) -> dict[str, Any]:
-        """Persist current model + optimizer state."""
+    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

verl/remote_backend/trainer.py

@@ -125,5 +125,11 @@ def _inferred_rollout_replica_count(self) -> int:
 
     def destroy(self) -> None:
         if getattr(self, "backend", None) is not None:
-            self.backend.destroy()
+            # `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

verl/remote_backend/worker.py

@@ -118,17 +118,17 @@ def reset(self):
     @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="ref"))
     @DistProfiler.annotate(color="olive", role="ref_compute_log_prob")
     async def compute_ref_log_prob(self, data: TensorDict) -> TensorDict:
-        return self._run_log_prob(data, ref=True, calculate_entropy=False)
+        return await self._run_log_prob(data, ref=True, calculate_entropy=False)
 
     @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="actor"))
     @DistProfiler.annotate(color="blue", role="actor_compute_log_prob")
     async def compute_log_prob(self, data: TensorDict) -> TensorDict:
-        return self._run_log_prob(data, ref=False, calculate_entropy=True)
+        return await self._run_log_prob(data, ref=False, calculate_entropy=True)
 
     @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="actor"))
     @DistProfiler.annotate(color="red", role="actor_update")
     async def update_actor(self, data: TensorDict) -> TensorDict:
-        return self._run_update_actor(data)
+        return await self._run_update_actor(data)
 
     @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
     async def update_weights(self, global_steps: int = None):
@@ -137,7 +137,7 @@ async def update_weights(self, global_steps: int = None):
     @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
     async def save_checkpoint(self, local_path, hdfs_path=None, global_step=0, max_ckpt_to_keep=None):
         assert self._is_actor, "save_checkpoint only supported on actor role"
-        self.backend.save_checkpoint()
+        await self.backend.save_checkpoint()
 
     @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
     async def sleep_replicas(self):
@@ -146,7 +146,7 @@ async def sleep_replicas(self):
     @register(dispatch_mode=Dispatch.ONE_TO_ALL)
     async def destroy(self):
         if self.backend is not None:
-            self.backend.destroy()
+            await self.backend.destroy()
             self.backend = None
 
     # ------------------------------------------------------------------ #
@@ -156,8 +156,8 @@ async def destroy(self):
     # are read directly from `self.config` by the backend at init time —
     # this forwarder no longer re-injects them into every batch.
 
-    def _run_log_prob(self, data: TensorDict, *, ref: bool, calculate_entropy: bool) -> TensorDict:
-        result = self.backend.compute_log_prob(
+    async def _run_log_prob(self, data: TensorDict, *, ref: bool, calculate_entropy: bool) -> TensorDict:
+        result = await self.backend.compute_log_prob(
             data,
             ref=ref,
             calculate_entropy=calculate_entropy,
@@ -178,7 +178,7 @@ def _run_log_prob(self, data: TensorDict, *, ref: bool, calculate_entropy: bool)
 
         return tu.get_tensordict(tensor_dict=model_output, non_tensor_dict={"metrics": metrics})
 
-    def _run_update_actor(self, data: TensorDict) -> TensorDict:
+    async def _run_update_actor(self, data: TensorDict) -> TensorDict:
         from tensordict import NonTensorData
 
         # `train_batch` style fields the backend may want to read from data.
@@ -191,7 +191,7 @@ def _run_update_actor(self, data: TensorDict) -> TensorDict:
             )
 
         with Timer(name="train_batch", logger=None) as timer:
-            result = self.backend.update_actor(
+            result = await self.backend.update_actor(
                 data,
                 actor_config=self.actor_config,
                 pad_token_id=self.pad_token_id,

verl/trainer/config/ppo_trainer.yaml

@@ -318,32 +318,38 @@ ray_kwargs:
   timeline_json_file: null
 
 
-# config for arctic rl
-arctic_rl:
+# 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.<name>`. 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:
 
-  # whether to use colocate mode
-  colocate: False
+  # Arctic adapter (`verl.trainer.ppo.arctic_rl_client`). Active when
+  # `trainer.remote_backend=arctic`.
+  arctic:
 
-  training_gpus: 1
-  sampling_gpus: 1
-  log_prob_gpus: 1
-  sampling_tp_size: 1
+    # whether to use colocate mode
+    colocate: False
 
-  use_zorro: False
+    training_gpus: 1
+    sampling_gpus: 1
+    log_prob_gpus: 1
+    sampling_tp_size: 1
 
-  # only set to True if running into OOM, otherwise it'll be faster not to tile (extra forward replay and smaller chunks to compute)
-  tiled_logits_compute: False
+    use_zorro: False
 
-  # ray or http: ray would be much faster for payload comms in the en-prem use case
-  comm_protocol: ray
+    # 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
+    zero_optimization:
+      # 2 for optim + grad sharding
+      # 3 for optim + grad +weight sharding
+      stage: 0
 
-    offload_optimizer:
-      device: none
+      offload_optimizer:
+        device: none
 
-    offload_param:
-      device: none
+      offload_param:
+        device: none

verl/trainer/ppo/arctic_rl_client.py

@@ -1,3 +1,4 @@
+import asyncio
 import os
 from typing import Any
 
@@ -94,10 +95,15 @@ class ArcticRLClientWrapper(RemoteBackend):
     it via ``config.trainer.remote_backend = "arctic"``.
     """
 
+    # Key under `main_config.remote_backend.<name>` 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]
         self._client = None
-        self.use_zorro = self.config.arctic_rl.use_zorro
+        self.use_zorro = self._backend_config.use_zorro
         self.use_liger = self.config.actor_rollout_ref.model.use_liger
         # Static, config-derived engineering value Arctic needs on every
         # `compute_log_prob` / `update_actor` call. Cached here at init
@@ -159,7 +165,7 @@ def requires_single_forwarder(self) -> bool:
     # ``processing`` pipeline) and dispatch through the private
     # `_send_*` wire helpers below.
 
-    def compute_log_prob(
+    async def compute_log_prob(
         self,
         data,
         *,
@@ -182,7 +188,7 @@ def compute_log_prob(
         )
         payload = dict(batch=batch, meta=meta)
 
-        response = (
+        response = await (
             self._send_compute_ref_log_prob(payload) if ref else self._send_compute_log_prob(payload)
         )
 
@@ -193,7 +199,7 @@ def compute_log_prob(
         metrics = dict(response.get("metrics", {}))
         return {"model_output": model_output, "metrics": metrics}
 
-    def update_actor(
+    async def update_actor(
         self,
         data,
         *,
@@ -251,7 +257,7 @@ def _safe_serialize(obj):
         meta["policy_loss_config"] = _safe_serialize(vars(policy_loss)) if policy_loss is not None else {}
 
         payload = dict(batch=batch, meta=meta)
-        response = self._send_update_actor(payload)
+        response = await self._send_update_actor(payload)
 
         metrics = dict(response["metrics"])
         loss = metrics.pop("loss")
@@ -267,7 +273,7 @@ def _safe_serialize(obj):
     def _create_ds_config(self, n_gpus: int) -> dict[str, Any]:
         actor_cfg = self.config.actor_rollout_ref.actor
         data_cfg = self.config.data
-        zero_optimization_config = self.config.arctic_rl.zero_optimization
+        zero_optimization_config = self._backend_config.zero_optimization
 
         micro_batch_size = actor_cfg.ppo_micro_batch_size_per_gpu or 1
         train_batch_size = data_cfg.train_batch_size * self.config.actor_rollout_ref.rollout.n
@@ -283,7 +289,7 @@ def _create_ds_config(self, n_gpus: int) -> dict[str, Any]:
                 "offload_optimizer": {"device": zero_optimization_config.offload_optimizer.device},
                 "offload_param": {"device": zero_optimization_config.offload_param.device},
             },
-    }
+        }
 
 
     def reconnect_config(self):
@@ -309,7 +315,7 @@ def _create_ds_worker_config(self):
                 max_token_len=self.config.actor_rollout_ref.rollout.max_num_batched_tokens,
                 rollout_n=self.config.actor_rollout_ref.rollout.n,
                 temperature=self.config.actor_rollout_ref.rollout.temperature,
-                tiled_logits_compute=self.config.arctic_rl.tiled_logits_compute,
+                tiled_logits_compute=self._backend_config.get("tiled_logits_compute", True),
                 use_unpad=use_unpad,
             )
 
@@ -324,10 +330,10 @@ def _initialize_client(self, reconnect_job_config: dict = None, rl_server_state:
             return create_arctic_rl_client(reconnect_job_config)
 
         model_name = self.config.actor_rollout_ref.model.path
-        n_training_gpus = self.config.arctic_rl.get("training_gpus", self.config.trainer.n_gpus_per_node)
-        n_sampling_gpus = self.config.arctic_rl.get("sampling_gpus", self.config.trainer.n_gpus_per_node)
-        n_log_prob_gpus = self.config.arctic_rl.get("log_prob_gpus", self.config.trainer.n_gpus_per_node)
-        colocate = self.config.arctic_rl.get("colocate", False)
+        n_training_gpus = self._backend_config.get("training_gpus", self.config.trainer.n_gpus_per_node)
+        n_sampling_gpus = self._backend_config.get("sampling_gpus", self.config.trainer.n_gpus_per_node)
+        n_log_prob_gpus = self._backend_config.get("log_prob_gpus", self.config.trainer.n_gpus_per_node)
+        colocate = self._backend_config.get("colocate", True)
         attn_implementation = self.config.actor_rollout_ref.model.override_config.get(
             'attn_implementation', 'eager'
         )
@@ -340,7 +346,7 @@ def _initialize_client(self, reconnect_job_config: dict = None, rl_server_state:
 
         rollout_cfg = self.config.actor_rollout_ref.rollout
         vllm_config = {
-            "tensor_parallel_size": self.config.arctic_rl.sampling_tp_size,
+            "tensor_parallel_size": self._backend_config.sampling_tp_size,
             "gpu_memory_utilization": rollout_cfg.gpu_memory_utilization,
             "max_model_len": rollout_cfg.get("max_model_len") or max_length,
             "max_num_seqs": rollout_cfg.max_num_seqs,
@@ -351,9 +357,9 @@ def _initialize_client(self, reconnect_job_config: dict = None, rl_server_state:
             vllm_config["quantization"] = rollout_cfg.quantization
 
         rl_config = ArcticRLClientConfig(
-            host=None if self.config.arctic_rl.comm_protocol == "ray" else "localhost",
-            port=None if self.config.arctic_rl.comm_protocol == "ray" else 7000,
-            comm_protocol=self.config.arctic_rl.comm_protocol,
+            host=None if self._backend_config.comm_protocol == "ray" else "localhost",
+            port=None if self._backend_config.comm_protocol == "ray" else 7000,
+            comm_protocol=self._backend_config.comm_protocol,
             backend="local",
             training_gpus=n_training_gpus,
             sampling_gpus=n_sampling_gpus,
@@ -414,19 +420,22 @@ async def generate(self, prompt_ids, sampling_params) -> list:
     # backends are free to pick their own wire format; verl never calls
     # these directly.
 
-    def _send_compute_ref_log_prob(self, payload: dict):
+    async def _send_compute_ref_log_prob(self, payload: dict):
         payload["processing"] = {"post": ["compute_entropy_and_logprobs"], "loss_fn": None}
-        response = self._client.fwd_no_grad(payload, reference_model=True)
+        # 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["batch"]["log_probs"] = response["batch"].pop("logprobs")
         return response
 
-    def _send_compute_log_prob(self, payload: dict):
+    async def _send_compute_log_prob(self, payload: dict):
         payload["processing"] = {"post": ["compute_entropy_and_logprobs"], "loss_fn": None}
-        response = self._client.fwd_no_grad(payload, reference_model=False)
+        response = await asyncio.to_thread(self._client.fwd_no_grad, payload, reference_model=False)
         response["batch"]["log_probs"] = response["batch"].pop("logprobs")
         return response
 
-    def _send_update_actor(self, payload: dict):
+    async def _send_update_actor(self, payload: dict):
         payload["processing"] = {
             "post": ["apply_temperature", "compute_entropy_and_logprobs"],
             "loss_fn": "verl_grpo",
@@ -446,13 +455,13 @@ 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 = self._client.fwd_bwd(payload)
-        step_response = self._client.step()
+        fwd_bwd_response = await asyncio.to_thread(self._client.fwd_bwd, payload)
+        step_response = await asyncio.to_thread(self._client.step)
         step_response["metrics"].update(**fwd_bwd_response["metrics"])
         return step_response
 
-    def save_checkpoint(self):
-        return self._client.save_checkpoint()
+    async def save_checkpoint(self):
+        return await self._client.save_checkpoint()
 
     async def update_weights(self):
         return await self._client.sync_weights()