[arctic-rl] align adapter with upcoming zorro_train config grouping

examples/arctic_rl/run_gsm8k_grpo_arl_zorro_yes.sh

@@ -73,7 +73,7 @@ python3 -m verl.trainer.main_ppo \
     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 \
+    remote_backend.zorro_train.enable=True \
     trainer.critic_warmup=0 \
     trainer.logger="['console']" \
     trainer.experiment_name=gsm8k_grpo_qwen3_0p6b_ngpu1_gbs16_rolln5_zorroTrue \

verl/trainer/config/remote_backend/arctic.yaml

@@ -26,12 +26,32 @@ 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
+# Grouped block for the ZoRRO training path. The arctic_training
+# server reads `zorro_train_enable` from `ds_worker_config` and
+# per-call request metas, so this adapter forwards `enable` as that
+# flat kwarg rather than passing the nested dict.
+zorro_train:
+
+  # whether to use zorro packing for off-policy log-probs
+  enable: False
+
+  # should equal actor_rollout_ref.rollout.n for best perf
+  max_rollouts: ${actor_rollout_ref.rollout.n}
 
 # ray or http: ray would be much faster for payload comms in the on-prem use case
 comm_protocol: ray
 
+# Zero-copy CUDA IPC weight sync between training and rollout engines.
+# Only effective with colocate=True (both engines on the same GPU). Avoids
+# the NCCL handshake path entirely.
+cuda_ipc_weight_sync: False
+
+# Stream one gathered param at a time during CUDA IPC weight sync so peak extra
+# GPU memory is one full param per GPU instead of the whole model (avoids OOM
+# on large models, at the cost of more round-trips). Only effective when
+# cuda_ipc_weight_sync=True and colocate=True.
+low_memory_weight_sync: False
+
 # DeepSpeed ZeRO-style sharding for the training engine
 zero_optimization:
 

verl/workers/remote_client/arctic_rl.py

@@ -15,8 +15,8 @@
 from typing import Any
 
 import torch
-from arctic_training.arctic_rl import ArcticRLClientConfig, create_arctic_rl_client
-from arctic_training.arctic_rl.ray_server import ArcticRLRayServerState
+from arctic_training.rl import ArcticRLClientConfig, create_arctic_rl_client
+from arctic_training.rl.ray_server import ArcticRLRayServerState
 from transformers import AutoTokenizer
 
 from verl.remote_backend.base import RemoteBackend, RemoteBackendRegistry
@@ -110,7 +110,12 @@ def __init__(self, config, reconnect_job_config: dict = None, rl_server_state: A
         # 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.zorro_train_enable = self._backend_config.zorro_train.enable
+        self.zorro_train_max_rollouts = self._backend_config.zorro_train.max_rollouts
+        # Weight-sync transport selection. CUDA-IPC bypasses the NCCL
+        # all_reduce path entirely; only valid when colocate=True.
+        self.cuda_ipc_weight_sync = self._backend_config.get("cuda_ipc_weight_sync", False)
+        self.low_memory_weight_sync = self._backend_config.get("low_memory_weight_sync", False)
         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
@@ -182,7 +187,8 @@ async def compute_log_prob(
     ) -> dict:
         batch, max_prompt_len, max_response_len = _prepare_padded_arctic_batch_dict(data, pad_token_id)
         meta = dict(
-            use_zorro=self.use_zorro,
+            zorro_train_enable=self.zorro_train_enable,
+            zorro_train_max_rollouts=self.zorro_train_max_rollouts,
             rollout_n=rollout_n,
             max_prompt_len=max_prompt_len,
             max_response_len=max_response_len,
@@ -241,7 +247,8 @@ async def update_actor(
             max_response_len=max_response_len,
             max_token_len_per_gpu=self._max_token_len_per_gpu,
             temperature=data["temperature"],
-            use_zorro=self.use_zorro,
+            zorro_train_enable=self.zorro_train_enable,
+            zorro_train_max_rollouts=self.zorro_train_max_rollouts,
             global_batch_size=data["global_batch_size"],
             rollout_is_weights=data.get("rollout_is_weights", None),
             batch_num_tokens=data["loss_mask"].sum(),
@@ -309,12 +316,12 @@ def _create_ds_worker_config(self):
             attn_implementation=attn_implementation,
         )
 
-        if self.use_zorro:
+        if self.zorro_train_enable:
             # XXX: can't find where it's configured
             use_unpad = True
 
             ds_worker_config.update(
-                use_zorro=True,
+                zorro_train_enable=True,
                 response_len=self.config.data.max_response_length,
                 max_token_len=self.config.actor_rollout_ref.rollout.max_num_batched_tokens,
                 rollout_n=self.config.actor_rollout_ref.rollout.n,
@@ -384,6 +391,7 @@ def _initialize_client(self, reconnect_job_config: dict = None, rl_server_state:
             },
             ds_worker_config=self._create_ds_worker_config(),
             vllm_config=vllm_config,
+            checkpoint_path=self.config.trainer.default_local_dir,
         )
 
         # ArcticRLClient is constructed as a ray remote actor with num_gpus=0,
@@ -403,10 +411,14 @@ def _initialize_client(self, reconnect_job_config: dict = None, rl_server_state:
         "max_tokens": 1024,
     }
 
-    async def generate(self, prompt_ids, sampling_params) -> list:
+    async def generate(self, prompt_ids, sampling_params, routing_key=None) -> list:
         prompts = [self.tokenizer.decode(prompt_ids)]
         merged_params = {**self._default_sampling_params, **sampling_params}
-        return await self._client.async_generate(prompts=prompts, sampling_params=merged_params)
+        return await self._client.generate(
+            prompts=prompts,
+            sampling_params=merged_params,
+            routing_key=routing_key,
+        )
 
     # ------------------------------------------------------------------ #
     # Arctic wire helpers (private; not on `RemoteBackend`)
@@ -456,7 +468,10 @@ async def save_checkpoint(self):
         return await self._client.save_checkpoint()
 
     async def update_weights(self):
-        return await self._client.sync_weights()
+        return await self._client.sync_weights(
+            cuda_ipc=self.cuda_ipc_weight_sync,
+            low_memory=self.low_memory_weight_sync,
+        )
 
     async def destroy(self) -> None:
         if self._client is not None: