diff --git a/examples/arctic_rl/run_gsm8k_grpo_arl_zorro_yes.sh b/examples/arctic_rl/run_gsm8k_grpo_arl_zorro_yes.sh index fc6cff538ab..8331cbe628f 100755 --- a/examples/arctic_rl/run_gsm8k_grpo_arl_zorro_yes.sh +++ b/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 \ diff --git a/verl/trainer/config/remote_backend/arctic.yaml b/verl/trainer/config/remote_backend/arctic.yaml index a80cbe2b90e..a62cc999c69 100644 --- a/verl/trainer/config/remote_backend/arctic.yaml +++ b/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: diff --git a/verl/workers/remote_client/arctic_rl.py b/verl/workers/remote_client/arctic_rl.py index 0ea3922e582..21a38c5f6de 100644 --- a/verl/workers/remote_client/arctic_rl.py +++ b/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: