Merge branch 'verl-project:main' into main

Author: SchumiDing

tests/experimental/reward_loop/test_agent_reward_loop_colocate.py

@@ -18,6 +18,7 @@
 from torchdata.stateful_dataloader import StatefulDataLoader
 from transformers import AutoTokenizer
 
+from verl.checkpoint_engine import CheckpointEngineManager
 from verl.experimental.agent_loop import AgentLoopManager
 from verl.experimental.reward_loop import RewardLoopManager
 from verl.protocol import DataProto
@@ -97,6 +98,13 @@ def test_agent_loop_reward_manager():
     actor_rollout_wg.init_model()
 
     agent_loop_manager = AgentLoopManager(config, worker_group=actor_rollout_wg)
+    # sleep rollout replicas
+    checkpoint_manager = CheckpointEngineManager(
+        backend=config.actor_rollout_ref.rollout.checkpoint_engine.backend,
+        trainer=actor_rollout_wg,
+        replicas=agent_loop_manager.rollout_replicas,
+    )
+    checkpoint_manager.sleep_replicas()
     reward_loop_manager = RewardLoopManager(config, rm_resource_pool=resource_pool)
 
     # 2. init test data
@@ -143,8 +151,11 @@ def _get_gen_batch(batch: DataProto) -> DataProto:
 
         return gen_batch
 
+    # wake up rollout replicas via update_weight
+    checkpoint_manager.update_weights()
     gen_batch = _get_gen_batch(batch)
     gen_batch = agent_loop_manager.generate_sequences(gen_batch)
+    checkpoint_manager.sleep_replicas()
 
     batch = batch.union(gen_batch)
     rm_outputs = reward_loop_manager.compute_rm_score(batch)

verl/trainer/ppo/ray_trainer.py

@@ -627,6 +627,15 @@ def _validate(self, merged: bool = False):
             sample_inputs.extend(input_texts)
             sample_uids.extend(test_batch.non_tensor_batch["uid"])
 
+            # compute reward model score if needed (similar to training loop)
+            if self.use_rm and "rm_scores" not in test_batch.batch.keys():
+                if not self.use_reward_loop:
+                    reward_tensor = self.rm_wg.compute_rm_score(test_batch)
+                else:
+                    assert self.reward_loop_manager is not None, "RewardLoopManager is None"
+                    reward_tensor = self.reward_loop_manager.compute_rm_score(test_batch)
+                test_batch = test_batch.union(reward_tensor)
+
             # evaluate using reward_function
             reward_tensor, reward_extra_info = self._compute_or_extract_reward(
                 test_batch, reward_fn=self.val_reward_fn, reward_for_val=True
@@ -1648,7 +1657,11 @@ def fit(self):
                     if esi_close_to_expiration:
                         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._save_checkpoint()
+                        # wake replicas to avoid OOM during checkpoint saving
+                        self.checkpoint_manager.update_weights()
 
                 with marked_timer("stop_profile", timing_raw):
                     next_step_profile = (