fix one step off ci

Author: ArronHZG

…rkflows/stash/e2e_fully_async_policy.yml …hub/workflows/e2e_fully_async_policy.yml.github/workflows/stash/e2e_fully_async_policy.yml renamed to .github/workflows/e2e_fully_async_policy.yml .github/workflows/stash/e2e_fully_async_policy.yml renamed to .github/workflows/e2e_fully_async_policy.yml

@@ -41,21 +41,21 @@ on:
       - main
       - v0.*
     paths:
-      - "**/*.py"
+      - "../../setup.py"
       - "!**/*.md"
       - "!**/*.sh"
       # Other entrypoints
       - "!examples/*trainer*"
       - "!tests/**"
       - "!verl/trainer/main_*.py"
       - "!verl/trainer/fsdp_sft_trainer.py"
-      - "verl/experimental/one_step_off_policy"
+      - "../../verl/experimental/one_step_off_policy"
   pull_request:
     branches:
       - main
       - v0.*
     paths:
-      - "**/*.py"
+      - "../../setup.py"
       - "!**/*.md"
       - "!**/*.sh"
       # Other entrypoints
@@ -64,11 +64,11 @@ on:
       - "!verl/trainer/main_*.py"
       - "!verl/trainer/fsdp_sft_trainer.py"
       # Home
-      - "verl/experimental/one_step_off_policy"
+      - "../../verl/experimental/one_step_off_policy"
       # Entrypoints
       - ".github/workflows/e2e_one_step_off_policy.yml"
-      - "examples/data_preprocess/gsm8k.py"
-      - "tests/special_e2e/run_one_step_off_policy.sh"
+      - "../../examples/data_preprocess/gsm8k.py"
+      - "../../tests/special_e2e/run_one_step_off_policy.sh"
 
 # Cancel jobs on the same ref if a new one is triggered
 concurrency:

…kflows/stash/e2e_one_step_off_policy.yml …ub/workflows/e2e_one_step_off_policy.yml.github/workflows/stash/e2e_one_step_off_policy.yml renamed to .github/workflows/e2e_one_step_off_policy.yml .github/workflows/stash/e2e_one_step_off_policy.yml renamed to .github/workflows/e2e_one_step_off_policy.yml

@@ -170,12 +170,24 @@ def run(self, config):
         processor = hf_processor(local_path, trust_remote_code=trust_remote_code, use_fast=True)
 
         # Load the reward manager for training and validation.
-        reward_fn = load_reward_manager(
-            config, tokenizer, num_examine=0, **config.reward_model.get("reward_kwargs", {})
-        )
-        val_reward_fn = load_reward_manager(
-            config, tokenizer, num_examine=1, **config.reward_model.get("reward_kwargs", {})
-        )
+        use_reward_loop = config.reward_model.use_reward_loop
+        if not use_reward_loop:
+            print(
+                "WARNING: Init reward manager in single controller will be deprecated. "
+                "Please set config.reward_model.use_reward_loop to use distributed reward manager."
+            )
+            # Load the reward manager for training and validation.
+            reward_fn = load_reward_manager(
+                config, tokenizer, num_examine=0, **config.reward_model.get("reward_kwargs", {})
+            )
+            val_reward_fn = load_reward_manager(
+                config, tokenizer, num_examine=1, **config.reward_model.get("reward_kwargs", {})
+            )
+        else:
+            # reward_loop will use init a reward loop manager in ray_trainer
+            # and use it to compute reward score
+            reward_fn = None
+            val_reward_fn = None
 
         resource_pool_manager = create_resource_pool_manager(config, role_worker_mapping.keys())
 

verl/experimental/one_step_off_policy/main_ppo.py

@@ -182,12 +182,27 @@ def _init_models(self):
         self._create_weight_sync_group()
 
     def _init_async_rollout_manager(self):
+
+        # infrastructure overview: https://verl.readthedocs.io/en/latest/advance/reward_loop.html#architecture-design
+        # agent_reward_loop: streaming reward computation with actor rollout
+        # two conditions satisfied: (1) no reward model, or (2) reward model with extra resource pool
+        enable_agent_reward_loop = self.use_reward_loop and (
+                not self.use_rm or self.config.reward_model.enable_resource_pool
+        )
+        # if enable_agent_reward_loop, we directly pass reward_loop_workers to agent loop manager
+        # to stream reward computation with actor rollout
+        self.reward_loop_worker_handles = self.reward_loop_manager.reward_loop_workers if enable_agent_reward_loop else None
+        reward_loop_worker_handles = self.reward_loop_manager.reward_loop_workers if enable_agent_reward_loop else None
+
+
         # create async rollout manager and request scheduler
         assert self.config.actor_rollout_ref.rollout.mode == "async"
         from verl.experimental.one_step_off_policy.agent_loop import OneStepOffAgentLoopManager
 
         self.async_rollout_mode = True
-        self.async_rollout_manager = OneStepOffAgentLoopManager(config=self.config, worker_group=self.rollout_wg)
+        self.async_rollout_manager = OneStepOffAgentLoopManager(config=self.config,
+                                                                worker_group=self.rollout_wg,
+                                                                reward_loop_worker_handles=reward_loop_worker_handles)
 
     def _create_weight_sync_group(self):
         from verl.utils.device import get_nccl_backend
@@ -356,7 +371,7 @@ async def fit(self):
 
         # perform validation before training
         # currently, we only support validation using the reward_function.
-        if self.val_reward_fn is not None and self.config.trainer.get("val_before_train", True):
+        if self.config.trainer.get("val_before_train", True):
             val_metrics = self._validate()
             assert val_metrics, f"{val_metrics=}"
             pprint(f"Initial validation metrics: {val_metrics}")
@@ -390,6 +405,8 @@ async def fit(self):
         batch_data_future = asyncio.create_task(self._async_gen_next_batch(continuous_iterator))
         while batch_data_future is not None:
             batch_data_future = await self.fit_step(batch_data_future, continuous_iterator)
+            if self.is_last_step:
+                return
 
     async def fit_step(self, batch_data_future, continuous_iterator):
         """
@@ -469,19 +486,20 @@ async def _fit_generate(self, batch_data_future, continuous_iterator):
 
         # sync weights from actor to rollout
         with marked_timer("sync_rollout_weights", timing_raw, color="purple"):
-            self.sync_rollout_weights()
+            self._fit_update_weights()
             await self.async_rollout_manager.clear_kv_cache()
 
         # async next generation
         if not self.is_last_step:
             batch_data_future = asyncio.create_task(self._async_gen_next_batch(continuous_iterator))
             await asyncio.sleep(0)
+        else:
+            batch_data_future = None
 
         return batch, batch_data_future
 
 
     def _fit_update_weights(self):
         # TODO: use checkpoint engine to update weight
-        self.sync_rollout_weights()
-
-
+        # self.sync_rollout_weights()
+        pass

verl/experimental/one_step_off_policy/ray_trainer.py

@@ -32,6 +32,7 @@
 from tqdm import tqdm
 
 from verl import DataProto
+from verl.checkpoint_engine import CheckpointEngineManager
 from verl.experimental.dataset.sampler import AbstractCurriculumSampler
 from verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup, ResourcePoolManager
 from verl.single_controller.ray.base import create_colocated_worker_cls
@@ -118,8 +119,16 @@ def init_workers(self):
         self._create_worker_classes()
         self._init_worker_groups()
         self._init_models()
+        self._init_reward_loop()
         self._init_async_rollout_manager()
 
+        self.checkpoint_manager = CheckpointEngineManager(
+            backend=self.config.actor_rollout_ref.rollout.checkpoint_engine.backend,
+            trainer=self.actor_rollout_wg,
+            replicas=self.async_rollout_manager.rollout_replicas,
+        )
+
+
     def _init_resource_pools(self):
         self.resource_pool_manager.create_resource_pool()
         self.resource_pool_to_cls = {pool: {} for pool in self.resource_pool_manager.resource_pool_dict.values()}
@@ -212,6 +221,21 @@ def _init_models(self):
         self.actor_rollout_wg = self.all_wg[str(Role.ActorRollout)]
         self.actor_rollout_wg.init_model()
 
+    def _init_reward_loop(self):
+        if self.use_reward_loop:
+            # create reward loop manager
+            if self.use_reward_loop:
+                from verl.experimental.reward_loop import RewardLoopManager
+
+                # initalize reward loop manager
+                # reward model (colocate or standalone): get resource_pool
+                # no reward model: resource_pool = None
+                resource_pool = self.resource_pool_manager.get_resource_pool(Role.RewardModel) if self.use_rm else None
+                self.reward_loop_manager = RewardLoopManager(
+                    config=self.config,
+                    rm_resource_pool=resource_pool,
+                )
+
     def _init_async_rollout_manager(self):
         pass
 
@@ -247,7 +271,7 @@ def fit(self):
 
         # perform validation before training
         # currently, we only support validation using the reward_function.
-        if self.val_reward_fn is not None and self.config.trainer.get("val_before_train", True):
+        if self.config.trainer.get("val_before_train", True):
             val_metrics = self._validate()
             assert val_metrics, f"{val_metrics=}"
             pprint(f"Initial validation metrics: {val_metrics}")
@@ -279,6 +303,8 @@ def fit(self):
             for batch_dict in self.train_dataloader:
                 self.epoch = epoch
                 self.fit_step(batch_dict)
+                if self.is_last_step:
+                    return
 
     def fit_step(self, batch_dict: Any = None):
         """
@@ -367,9 +393,6 @@ def _fit_generate(self, batch: DataProto = None) -> DataProto:
             gen_batch_output.meta_info.pop("timing", None)
 
         if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX:
-            if self.reward_fn is None:
-                raise ValueError("A reward_fn is required for REMAX advantage estimation.")
-
             with marked_timer("gen_max", timing_raw, color="purple"):
                 gen_baseline_batch = deepcopy(gen_batch)
                 gen_baseline_batch.meta_info["do_sample"] = False
@@ -386,17 +409,16 @@ def _fit_generate(self, batch: DataProto = None) -> DataProto:
                 # compute reward model score on batch
                 rm_scores = None
                 if self.use_rm and "rm_scores" not in batch.batch.keys():
-                    if not self.use_reward_loop:
-                        rm_scores = self.rm_wg.compute_rm_score(batch)
-                    else:
-                        assert self.reward_loop_manager is not None, "RewardLoopManager is None"
-                        rm_scores = self.reward_loop_manager.compute_rm_score(batch)
-                    batch = batch.union(rm_scores)
+                    batch_reward = self._compute_reward_colocate(batch)
+                    batch = batch.union(batch_reward)
 
                 # Compute or extract reward for REMAX baseline
-                reward_baseline_tensor = self._compute_or_extract_reward(
-                    batch, reward_fn=self.reward_fn, sum_reward=True
-                )
+                if not self.use_reward_loop:
+                    reward_baseline_tensor = self._compute_reward_legacy(
+                        batch, reward_fn=self.reward_fn, sum_reward=True
+                    )
+                else:
+                    reward_baseline_tensor = batch.batch["rm_scores"].sum(dim=-1)
 
                 keys_to_pop = set(gen_baseline_output.batch.keys())
                 if rm_scores is not None:
@@ -435,22 +457,23 @@ def _fit_compute_reward(self, batch: DataProto) -> DataProto:
         with marked_timer("reward", timing_raw, color="yellow"):
             # compute reward model score
             if self.use_rm and "rm_scores" not in batch.batch.keys():
-                if not self.use_reward_loop:
-                    self.reward_tensor = self.rm_wg.compute_rm_score(batch)
+                batch_reward = self._compute_reward_colocate(batch)
+                batch = batch.union(batch_reward)
+
+            # Compute or extract reward_tensor and reward_extra_infos_dict for training
+            if not self.use_reward_loop:
+                if self.config.reward_model.launch_reward_fn_async:
+                    self.future_reward = compute_reward_async.remote(
+                        data=batch, config=self.config, tokenizer=self.tokenizer
+                    )
                 else:
-                    assert self.reward_loop_manager is not None
-                    self.reward_tensor = self.reward_loop_manager.compute_rm_score(batch)
-                batch = batch.union(self.reward_tensor)
-
-            # Compute or extract reward for training
-            if self.config.reward_model.launch_reward_fn_async:
-                self.future_reward = compute_reward_async.remote(
-                    data=batch, config=self.config, tokenizer=self.tokenizer
-                )
+                    self.reward_tensor, self.reward_extra_infos_dict = self._compute_reward_legacy(
+                        batch, reward_fn=self.reward_fn, reward_for_val=False
+                    )
             else:
-                self.reward_tensor, self.reward_extra_infos_dict = self._compute_or_extract_reward(
-                    batch, reward_fn=self.reward_fn, reward_for_val=False
-                )
+                self.reward_tensor = batch.batch["rm_scores"]
+                reward_extra_keys = batch.meta_info.get("reward_extra_keys", [])
+                self.reward_extra_infos_dict = {key: batch.non_tensor_batch[key] for key in reward_extra_keys}
         return batch
 
     def _fit_compute_log_prob(self, batch: DataProto) -> DataProto:
@@ -620,11 +643,8 @@ def _fit_dump_data(self, batch: DataProto):
     def _fit_validate(self):
         metrics = self.metrics
         timing_raw = self.timing_raw
-        if (
-            self.val_reward_fn is not None
-            and self.config.trainer.test_freq > 0
-            and (self.is_last_step or self.global_steps % self.config.trainer.test_freq == 0)
-        ):
+        if self.config.trainer.test_freq > 0 and (
+                self.is_last_step or self.global_steps % self.config.trainer.test_freq == 0):
             with marked_timer("testing", timing_raw, color="green"):
                 val_metrics: dict = self._validate()
                 if self.is_last_step:
@@ -652,10 +672,11 @@ def _fit_save_checkpoint(self):
                 print("Force saving checkpoint: ESI instance expiration approaching.")
             with marked_timer("save_checkpoint", timing_raw, color="green"):
                 # sleep replicas to avoid OOM during checkpoint saving
-                self.checkpoint_manager.sleep_replicas()
+                # self.checkpoint_manager.sleep_replicas()
                 self._save_checkpoint()
                 # wake replicas to avoid OOM during checkpoint saving
-                self.checkpoint_manager.update_weights()
+                # TODO: Check separation is needed.
+                # self.checkpoint_manager.update_weights()
 
     def _fit_stop_profile(self):
         timing_raw = self.timing_raw
@@ -727,4 +748,3 @@ def _fit_postprocess_step(self):
                 self.actor_rollout_wg.async_calls_finalize_fn_exec(blocking=True)
             pprint(f"Final validation metrics: {self.last_val_metrics}")
             self.progress_bar.close()
-            return

verl/trainer/ppo/ray_trainer_for_separation.py

@@ -266,7 +266,7 @@ async def launch_server(self, master_address: str = None, master_port: int = Non
             "dp_size": self.config.data_parallel_size,
             "ep_size": self.config.expert_parallel_size,
             "node_rank": self.node_rank,
-            "load_format": self.config.load_format,
+            "load_format": "auto",
             "dist_init_addr": dist_init_addr,
             "nnodes": self.nnodes,
             "trust_remote_code": self.model_config.trust_remote_code,
@@ -391,8 +391,8 @@ async def sleep(self):
             logger.info("skip sleep in standalone mode")
 
     async def clear_kv_cache(self):
-        obj = ReleaseMemoryOccupationReqInput(tags=["kv_cache"])
-        await self.tokenizer_manager.release_memory_occupation(obj, None)
+        if self.node_rank == 0:
+            await self.tokenizer_manager.flush_cache()
 
     async def generate(
         self,