Arctic RL: VeRL integration

examples/arctic_rl/run_gsm8k_grpo_arl_zorro_yes.sh

@@ -0,0 +1,87 @@
+#!/bin/bash
+
+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
+# we want to make sure this runs on non-gpu client
+export CUDA_VISIBLE_DEVICES=
+
+BSZ=1024
+UBS=32
+ROLL_N=5
+PROMPT_LENGTH=512
+RESPONSE_LENGTH=1024
+MAX_STEPS=100
+
+
+BSZ=8
+UBS=2
+ROLL_N=2
+PROMPT_LENGTH=512
+RESPONSE_LENGTH=1024
+MAX_STEPS=4
+
+
+experiment_name="qwen3-0.6B_arctic_gsm8k_grpo"
+
+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=$BSZ \
+    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=4 \
+    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=$BSZ \
+    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=$UBS \
+    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_3 \
+    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=$UBS \
+    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=$ROLL_N \
+    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=$UBS \
+    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.use_arctic_rl=True \
+    trainer.critic_warmup=0 \
+    trainer.logger=console \
+    trainer.experiment_name=$experiment_name \
+    trainer.project_name=arctic_gsm8k_grpo \
+    trainer.val_before_train=False \
+    trainer.n_gpus_per_node=1 \
+    trainer.nnodes=1 \
+    trainer.save_freq=-1 \
+    trainer.test_freq=-1 \
+    trainer.total_epochs=15 \
+    trainer.total_training_steps=${MAX_STEPS} \
+    arctic_rl.colocate=False \
+    arctic_rl.training_gpus=2\
+    arctic_rl.sampling_gpus=2\
+    arctic_rl.log_prob_gpus=0\
+    arctic_rl.use_zorro=True \
+    "$@" 2>&1 | tee $experiment_name.log

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/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,8 @@ 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 +229,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}"

verl/trainer/config/ppo_trainer.yaml

@@ -203,6 +203,9 @@ trainer:
   #  mode: "auto", "enable", or "disable"
   use_legacy_worker_impl: auto
 
+  # whether to use arctic rl
+  use_arctic_rl: False
+
 # profiler configs
 global_profiler:
 
@@ -310,3 +313,16 @@ ray_kwargs:
 
   # Path to save Ray timeline JSON for performance profiling
   timeline_json_file: null
+
+
+# config for arctic rl
+arctic_rl:
+
+  # whether to use colocate mode
+  colocate: False
+
+  training_gpus: 1
+  sampling_gpus: 1
+  log_prob_gpus: 1
+
+  use_zorro: False

verl/trainer/main_ppo.py

@@ -134,6 +134,10 @@ def add_actor_rollout_worker(self, config):
             actor_rollout_cls = ActorRolloutRefWorker
             ray_worker_group_cls = RayWorkerGroup
 
+            if config.trainer.get("use_arctic_rl", False):
+                from verl.workers.arctic_workers import ActorRolloutRefWorker
+                actor_rollout_cls = ActorRolloutRefWorker
+
             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)
@@ -340,7 +344,9 @@ def run(self, config):
         train_sampler = create_rl_sampler(config.data, train_dataset)
 
         # Initialize the PPO trainer.
-        trainer = RayPPOTrainer(
+        from verl.trainer.ppo.arctic_trainer import ArcticPPOTrainer
+        ppo_trainer_cls = RayPPOTrainer if not config.trainer.use_arctic_rl else ArcticPPOTrainer
+        trainer = ppo_trainer_cls(
             config=config,
             tokenizer=tokenizer,
             processor=processor,
@@ -356,7 +362,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):