bug fix: validation ckpt

Author: atlas5301

verl/trainer/config/_generated_ppo_megatron_trainer.yaml

@@ -595,6 +595,7 @@ trainer:
   resume_from_path: null
   del_local_ckpt_after_load: false
   val_before_train: true
+  validation_use_train_namespace: true
   test_freq: -1
   critic_warmup: 0
   default_hdfs_dir: null

verl/trainer/config/_generated_ppo_trainer.yaml

@@ -511,6 +511,7 @@ trainer:
   log_val_generations: 0
   rollout_data_dir: null
   validation_data_dir: null
+  validation_use_train_namespace: true
   nnodes: 1
   n_gpus_per_node: 8
   save_freq: -1

verl/trainer/config/ppo_megatron_trainer.yaml

@@ -147,6 +147,8 @@ trainer:
   resume_from_path: null
   del_local_ckpt_after_load: False
   val_before_train: True
+  # Whether to run validation using each training namespace's own policy (instead of the rollout_from provider)
+  validation_use_train_namespace: True
   test_freq: -1
   critic_warmup: 0
   default_hdfs_dir: null

verl/trainer/config/ppo_trainer.yaml

@@ -183,6 +183,9 @@ trainer:
   # Directory for logging validation data; no dump if null
   validation_data_dir: null
 
+  # Whether to run validation using each training namespace's own policy (instead of the rollout_from provider)
+  validation_use_train_namespace: True
+
   # Number of nodes used in the training
   nnodes: 1
 

verl/trainer/namespace.py

@@ -105,7 +105,18 @@ def _maybe_override_namespace_fields(cfg: DictConfig, name: str, rollout_from: s
 
 
 def _compute_spawn_roles(spec: NamespaceSpec, rollout_dependents: Iterable[str], ref_dependents: Iterable[str]):
-    provides_rollout = spec.rollout_from == spec.name or len(list(rollout_dependents)) > 0
+    prefer_val_rollout = bool(spec.config.trainer.get("validation_use_train_namespace", False))
+    wants_validation = bool(
+        spec.train
+        and prefer_val_rollout
+        and (
+            spec.config.trainer.get("val_only", False)
+            or spec.config.trainer.get("val_before_train", True)
+            or spec.config.trainer.get("test_freq", -1) > 0
+        )
+    )
+
+    provides_rollout = spec.rollout_from == spec.name or len(list(rollout_dependents)) > 0 or wants_validation
     provides_ref = (spec.needs_ref and spec.ref_from == spec.name) or len(list(ref_dependents)) > 0
 
     spec.provides_rollout = provides_rollout

verl/trainer/ppo/ray_trainer.py

@@ -808,7 +808,7 @@ def _maybe_log_val_generations(self, inputs, outputs, scores):
         # Log to each configured logger
         self.validation_generations_logger.log(self.config.trainer.logger, samples, self.global_steps)
 
-    def _get_gen_batch(self, batch: DataProto) -> DataProto:
+    def _get_gen_batch(self, batch: DataProto, async_rollout_mode: bool | None = None) -> DataProto:
         reward_model_keys = set({"data_source", "reward_model", "extra_info", "uid"}) & batch.non_tensor_batch.keys()
 
         # pop those keys for generation
@@ -820,7 +820,10 @@ def _get_gen_batch(self, batch: DataProto) -> DataProto:
         )
 
         # For agent loop, we need reward model keys to compute score.
-        if self.async_rollout_mode:
+        if async_rollout_mode is None:
+            async_rollout_mode = self.async_rollout_mode
+
+        if async_rollout_mode:
             gen_batch.non_tensor_batch.update(batch.non_tensor_batch)
 
         return gen_batch
@@ -829,6 +832,40 @@ def _validate(self):
         data_source_lst = []
         reward_extra_infos_dict: dict[str, list] = defaultdict(list)
 
+        prefer_train_rollout = bool(self.config.trainer.get("validation_use_train_namespace", False))
+        rollout_ns = self.namespace_specs[self.active_namespace].rollout_from
+        if prefer_train_rollout and self.active_namespace in self.rollout_wg_map:
+            rollout_ns = self.active_namespace
+
+        rollout_wg = self.rollout_wg_map.get(rollout_ns)
+        if rollout_wg is None and rollout_ns != self.namespace_specs[self.active_namespace].rollout_from:
+            fallback_ns = self.namespace_specs[self.active_namespace].rollout_from
+            rollout_wg = self.rollout_wg_map.get(fallback_ns)
+            rollout_ns = fallback_ns
+
+        if rollout_wg is None:
+            raise ValueError(
+                f"Rollout worker for validation namespace '{rollout_ns}' (active namespace '{self.active_namespace}') is missing"
+            )
+
+        rollout_cfg = self.namespace_specs[rollout_ns].config.actor_rollout_ref.rollout
+        async_rollout_mode = rollout_cfg.mode == "async"
+        async_rollout_manager = None
+        if async_rollout_mode:
+            if self.async_rollout_mode and rollout_wg is self.rollout_wg:
+                async_rollout_manager = self.async_rollout_manager
+            else:
+                from verl.experimental.agent_loop import AgentLoopManager
+
+                rm_resource_pool = None
+                if self.config.reward_model.enable and self.config.reward_model.enable_resource_pool:
+                    rm_resource_pool = self.resource_pool_manager.get_resource_pool(Role.RewardModel)
+                async_rollout_manager = AgentLoopManager(
+                    config=self.namespace_specs[rollout_ns].config,
+                    worker_group=rollout_wg,
+                    rm_resource_pool=rm_resource_pool,
+                )
+
         # Lists to collect samples for the table
         sample_inputs = []
         sample_outputs = []
@@ -850,7 +887,7 @@ def _validate(self):
 
             # repeat test batch
             test_batch = test_batch.repeat(
-                repeat_times=self.config.actor_rollout_ref.rollout.val_kwargs.n, interleave=True
+                repeat_times=rollout_cfg.val_kwargs.n, interleave=True
             )
 
             # we only do validation on rule-based rm
@@ -869,28 +906,26 @@ def _validate(self):
             ]
             sample_gts.extend(ground_truths)
 
-            test_gen_batch = self._get_gen_batch(test_batch)
+            test_gen_batch = self._get_gen_batch(test_batch, async_rollout_mode=async_rollout_mode)
             test_gen_batch.meta_info = {
                 "eos_token_id": self.tokenizer.eos_token_id,
                 "pad_token_id": self.tokenizer.pad_token_id,
                 "recompute_log_prob": False,
-                "do_sample": self.config.actor_rollout_ref.rollout.val_kwargs.do_sample,
+                "do_sample": rollout_cfg.val_kwargs.do_sample,
                 "validate": True,
                 "global_steps": self.global_steps,
             }
             print(f"test_gen_batch meta info: {test_gen_batch.meta_info}")
 
             # pad to be divisible by dp_size
             size_divisor = (
-                self.rollout_wg.world_size
-                if not self.async_rollout_mode
-                else self.config.actor_rollout_ref.rollout.agent.num_workers
+                rollout_wg.world_size if not async_rollout_mode else rollout_cfg.agent.num_workers
             )
             test_gen_batch_padded, pad_size = pad_dataproto_to_divisor(test_gen_batch, size_divisor)
-            if not self.async_rollout_mode:
-                test_output_gen_batch_padded = self.rollout_wg.generate_sequences(test_gen_batch_padded)
+            if not async_rollout_mode:
+                test_output_gen_batch_padded = rollout_wg.generate_sequences(test_gen_batch_padded)
             else:
-                test_output_gen_batch_padded = self.async_rollout_manager.generate_sequences(test_gen_batch_padded)
+                test_output_gen_batch_padded = async_rollout_manager.generate_sequences(test_gen_batch_padded)
 
             # unpad
             test_output_gen_batch = unpad_dataproto(test_output_gen_batch_padded, pad_size=pad_size)