Repository: volcengine/verl
Pull request #553: [rollout] feat: support sampling in validation stage
Author: login: PeterSH6
Created: 2025-03-12T06:10:24Z | Merged: 2025-03-13T02:40:41Z

Currently, eager mode is applied in the validation stage. However, in some reasoning tasks, we may need to generate n times and average the scores.

In this PR, we support using non-eager sampling parameters during validation by specifying the `val_kwargs` in `actor_rollout_ref.rollout` config field.


**Future work**
- [ ] Merge `vllm_rollout_spmd.py` and `vllm_rollout.py` into one file.


Diff:
diff --git a/verl/trainer/config/ppo_megatron_trainer.yaml b/verl/trainer/config/ppo_megatron_trainer.yaml
index 74ddb192ea9..1859a05ab31 100644
--- a/verl/trainer/config/ppo_megatron_trainer.yaml
+++ b/verl/trainer/config/ppo_megatron_trainer.yaml
@@ -84,6 +84,13 @@ actor_rollout_ref:
       gate_proj_layer_name: gate_up
     # number of responses (i.e. num sample times)
     n: 1
+    val_kwargs:
+      # sampling parameters for validation
+      top_k: -1 # 0 for hf rollout, -1 for vllm rollout
+      top_p: 1.0
+      temperature: 0
+      n: 1
+      do_sample: False # default eager for validation
 
 critic:
   strategy: megatron
diff --git a/verl/trainer/config/ppo_trainer.yaml b/verl/trainer/config/ppo_trainer.yaml
index 40bcd0e39b5..bf486279bcb 100644
--- a/verl/trainer/config/ppo_trainer.yaml
+++ b/verl/trainer/config/ppo_trainer.yaml
@@ -91,6 +91,13 @@ actor_rollout_ref:
     do_sample: True
     # number of responses (i.e. num sample times)
     n: 1 # > 1 for grpo
+    val_kwargs:
+      # sampling parameters for validation
+      top_k: -1 # 0 for hf rollout, -1 for vllm rollout
+      top_p: 1.0
+      temperature: 0
+      n: 1
+      do_sample: False # default eager for validation
 
 critic:
   strategy: fsdp
diff --git a/verl/trainer/ppo/ray_trainer.py b/verl/trainer/ppo/ray_trainer.py
index 23678d43289..190e20f2af2 100644
--- a/verl/trainer/ppo/ray_trainer.py
+++ b/verl/trainer/ppo/ray_trainer.py
@@ -533,6 +533,11 @@ def check_mutually_exclusive(mbs, mbs_per_gpu, name: str):
                 f"WARNING: val_batch_size is deprecated. Validation datasets are sent to inference engines as a whole batch, which will schedule the memory themselves."
             )
 
+        # check eval config
+        if config.actor_rollout_ref.rollout.val_kwargs.do_sample:
+            assert config.actor_rollout_ref.rollout.temperature > 0, \
+                "validation gen temperature should be greater than 0 when enabling do_sample"
+
         print("[validate_config] All configuration checks passed successfully!")
 
     def _create_dataloader(self):
@@ -661,9 +666,10 @@ def _validate(self):
                 'eos_token_id': self.tokenizer.eos_token_id,
                 'pad_token_id': self.tokenizer.pad_token_id,
                 'recompute_log_prob': False,
-                'do_sample': False,
+                'do_sample': self.config.actor_rollout_ref.rollout.val_kwargs.do_sample,
                 'validate': True,
             }
+            print(f'test_gen_batch meta info: {test_gen_batch.meta_info}')
 
             # pad to be divisible by dp_size
             test_gen_batch_padded, pad_size = pad_dataproto_to_divisor(test_gen_batch, self.actor_rollout_wg.world_size)
diff --git a/verl/workers/rollout/vllm_rollout/vllm_rollout.py b/verl/workers/rollout/vllm_rollout/vllm_rollout.py
index 4af7eca12fc..3a696e75f9e 100644
--- a/verl/workers/rollout/vllm_rollout/vllm_rollout.py
+++ b/verl/workers/rollout/vllm_rollout/vllm_rollout.py
@@ -119,7 +119,7 @@ def __init__(self, actor_module: nn.Module, config: DictConfig, tokenizer, model
 
         kwargs = dict(
             n=1,
-            logprobs=1,  # can be set to 0 and let actor to recompute
+            logprobs=0,  # can be set to 0 and let actor to recompute
             max_tokens=config.response_length,
         )
 
@@ -175,6 +175,7 @@ def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto:
             idx_list.append(_pre_process_inputs(self.pad_token_id, idx[i]))
 
         do_sample = prompts.meta_info.get('do_sample', True)
+        is_validate = prompts.meta_info.get('validate', False)
         if not do_sample:
             kwargs = {
                 'best_of': 1,
@@ -184,6 +185,14 @@ def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto:
                 'temperature': 0,
                 'n': 1  # if greedy, only 1 response
             }
+        elif is_validate:
+            # TODO: try **
+            kwargs = {
+                'top_k': self.config.val_kwargs.top_k,
+                'top_p': self.config.val_kwargs.top_p,
+                'temperature': self.config.val_kwargs.temperature,
+                'n': self.config.val_kwargs.n,
+            }
 
         # users can customize different sampling_params at different run
         with self.update_sampling_params(**kwargs):
@@ -193,21 +202,22 @@ def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto:
                 prompt_token_ids=idx_list,
                 use_tqdm=False)
 
-        # TODO(sgm): disable logprob when recompute_log_prob is enable
-        # if n = 1: (bs, response_length) ; if n > 1: (bs * n, response_length)
-        response = output[0].to(idx.device)
-        log_probs = output[1].to(idx.device)
-
-        if response.shape[1] < self.config.response_length:
-            response = pad_sequence_to_length(response, self.config.response_length, self.pad_token_id)
-            log_probs = pad_sequence_to_length(log_probs, self.config.response_length, self.pad_token_id)
-
-        if self.config.n > 1 and do_sample:
-            idx = idx.repeat_interleave(self.config.n, dim=0)
-            attention_mask = attention_mask.repeat_interleave(self.config.n, dim=0)
-            position_ids = position_ids.repeat_interleave(self.config.n, dim=0)
-            batch_size = batch_size * self.config.n
-        seq = torch.cat([idx, response], dim=-1)
+            # TODO(sgm): disable logprob when recompute_log_prob is enable
+            # if n = 1: (bs, response_length) ; if n > 1: (bs * n, response_length)
+            response = output[0].to(idx.device)
+            # log_probs = output[1].to(idx.device)
+
+            if response.shape[1] < self.config.response_length:
+                response = pad_sequence_to_length(response, self.config.response_length, self.pad_token_id)
+                # log_probs = pad_sequence_to_length(log_probs, self.config.response_length, self.pad_token_id)
+
+            # utilize current sampling params
+            if self.sampling_params.n > 1 and do_sample:
+                idx = idx.repeat_interleave(self.sampling_params.n, dim=0)
+                attention_mask = attention_mask.repeat_interleave(self.sampling_params.n, dim=0)
+                position_ids = position_ids.repeat_interleave(self.sampling_params.n, dim=0)
+                batch_size = batch_size * self.sampling_params.n
+            seq = torch.cat([idx, response], dim=-1)
 
         response_length = response.size(1)
         delta_position_id = torch.arange(1, response_length + 1, device=position_ids.device)
diff --git a/verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py b/verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py
index df52f0786cc..8e1de199fb6 100644
--- a/verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py
+++ b/verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py
@@ -188,6 +188,7 @@ def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto:
             } for raw_prompt_ids in non_tensor_batch.pop('raw_prompt_ids')]
 
         do_sample = prompts.meta_info.get('do_sample', True)
+        is_validate = prompts.meta_info.get('validate', False)
         if not do_sample:
             kwargs = {
                 'best_of': 1,
@@ -197,6 +198,14 @@ def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto:
                 'temperature': 0,
                 'n': 1  # if greedy, only 1 response
             }
+        elif is_validate:
+            # TODO: try **
+            kwargs = {
+                'top_k': self.config.val_kwargs.top_k,
+                'top_p': self.config.val_kwargs.top_p,
+                'temperature': self.config.val_kwargs.temperature,
+                'n': self.config.val_kwargs.n,
+            }
 
         # users can customize different sampling_params at different run
         with self.update_sampling_params(**kwargs):
@@ -205,27 +214,27 @@ def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto:
                 sampling_params=self.sampling_params,
                 use_tqdm=False)
 
-        # TODO(sgm): disable logprob when recompute_log_prob is enable
-        # if n = 1: (bs, response_length) ; if n > 1: (bs * n, response_length)
+            # TODO(sgm): disable logprob when recompute_log_prob is enable
+            # if n = 1: (bs, response_length) ; if n > 1: (bs * n, response_length)
 
-        response = []
-        for output in outputs:
-            for sample_id in range(len(output.outputs)):
-                response.append(output.outputs[sample_id].token_ids)
+            response = []
+            for output in outputs:
+                for sample_id in range(len(output.outputs)):
+                    response.append(output.outputs[sample_id].token_ids)
 
-        response = pad_2d_list_to_length(response, self.pad_token_id,
-                                         max_length=self.config.response_length).to(idx.device)
+            response = pad_2d_list_to_length(response, self.pad_token_id,
+                                             max_length=self.config.response_length).to(idx.device)
 
-        if self.config.n > 1 and do_sample:
-            idx = _repeat_interleave(idx, self.config.n)
-            attention_mask = _repeat_interleave(attention_mask, self.config.n)
-            position_ids = _repeat_interleave(position_ids, self.config.n)
-            batch_size = batch_size * self.config.n
-            if 'multi_modal_inputs' in non_tensor_batch.keys():
-                non_tensor_batch['multi_modal_inputs'] = _repeat_interleave(non_tensor_batch['multi_modal_inputs'],
-                                                                            self.config.n)
+            if self.sampling_params.n > 1 and do_sample:
+                idx = _repeat_interleave(idx, self.sampling_params.n)
+                attention_mask = _repeat_interleave(attention_mask, self.sampling_params.n)
+                position_ids = _repeat_interleave(position_ids, self.sampling_params.n)
+                batch_size = batch_size * self.sampling_params.n
+                if 'multi_modal_inputs' in non_tensor_batch.keys():
+                    non_tensor_batch['multi_modal_inputs'] = _repeat_interleave(non_tensor_batch['multi_modal_inputs'],
+                                                                                self.sampling_params.n)
 
-        seq = torch.cat([idx, response], dim=-1)
+            seq = torch.cat([idx, response], dim=-1)
 
         response_length = response.size(1)
         delta_position_id = torch.arange(1, response_length + 1, device=position_ids.device)


Reviews:
user: login: vermouth1992
state: DISMISSED
author_association: COLLABORATOR
submitted_at: 2025-03-12T07:24:45Z
user: login: tongyx361
body: This PR looks correct to me.
But some details can be improved. Could you / Should I make a new commit or PR?
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-12T10:38:32Z
user: login: PeterSH6
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-12T11:19:10Z
user: login: PeterSH6
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-12T11:19:51Z
user: login: PeterSH6
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-12T11:20:33Z
user: login: tongyx361
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-12T11:40:48Z
user: login: vermouth1992
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-13T02:40:40Z

Review comments:
diff_hunk: @@ -84,6 +84,13 @@ actor_rollout_ref:
       gate_proj_layer_name: gate_up
     # number of responses (i.e. num sample times)
     n: 1
+    val_kwargs:
path: verl/trainer/config/ppo_megatron_trainer.yaml
user: login: tongyx361
body: yaml items are naturally kwargs. Is `val` a better name?

```suggestion
    val:
```
created_at: 2025-03-12T09:23:16Z
author_association: COLLABORATOR
diff_hunk: @@ -84,6 +84,13 @@ actor_rollout_ref:
       gate_proj_layer_name: gate_up
     # number of responses (i.e. num sample times)
     n: 1
+    val_kwargs:
+      # sampling parameters for validation
+      top_k: -1 # 0 for hf rollout, -1 for vllm rollout
+      top_p: 1.0
+      temperature: 0
+      n: 1
path: verl/trainer/config/ppo_megatron_trainer.yaml
user: login: tongyx361
body: It might be helpful to support different `n` per validation dataset.

This might be simply implemented by duplicating the data and setting `n` in `SamplingParams` as `1`.
created_at: 2025-03-12T09:25:41Z
author_association: COLLABORATOR
diff_hunk: @@ -184,6 +185,14 @@ def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto:
                 'temperature': 0,
                 'n': 1  # if greedy, only 1 response
             }
+        elif is_validate and do_sample:
path: verl/workers/rollout/vllm_rollout/vllm_rollout.py
user: login: tongyx361
body: `and do_sample` is redundant. But leaving it as comment might making this branch more readable.


```suggestion
        elif is_validate: # and do_sample
```
created_at: 2025-03-12T10:07:50Z
author_association: COLLABORATOR
diff_hunk: @@ -184,6 +185,14 @@ def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto:
                 'temperature': 0,
                 'n': 1  # if greedy, only 1 response
             }
+        elif is_validate and do_sample:
path: verl/workers/rollout/vllm_rollout/vllm_rollout.py
user: login: PeterSH6
body: This `do_sample` is necessary. If the user accidentally set do_sample=False while setting `temperature>0`. Then, sampling parameters will be override
created_at: 2025-03-12T11:19:10Z
author_association: COLLABORATOR
diff_hunk: @@ -84,6 +84,13 @@ actor_rollout_ref:
       gate_proj_layer_name: gate_up
     # number of responses (i.e. num sample times)
     n: 1
+    val_kwargs:
path: verl/trainer/config/ppo_megatron_trainer.yaml
user: login: PeterSH6
body: val_kwargs is used to specify that this field relates to the kwargs variable in vllm_rollout.py
created_at: 2025-03-12T11:19:51Z
author_association: COLLABORATOR
diff_hunk: @@ -84,6 +84,13 @@ actor_rollout_ref:
       gate_proj_layer_name: gate_up
     # number of responses (i.e. num sample times)
     n: 1
+    val_kwargs:
+      # sampling parameters for validation
+      top_k: -1 # 0 for hf rollout, -1 for vllm rollout
+      top_p: 1.0
+      temperature: 0
+      n: 1
path: verl/trainer/config/ppo_megatron_trainer.yaml
user: login: PeterSH6
body: That's a good suggestion for the next PR.
created_at: 2025-03-12T11:20:33Z
author_association: COLLABORATOR
diff_hunk: @@ -184,6 +185,14 @@ def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto:
                 'temperature': 0,
                 'n': 1  # if greedy, only 1 response
             }
+        elif is_validate and do_sample:
path: verl/workers/rollout/vllm_rollout/vllm_rollout.py
user: login: tongyx361
body: I mean there is a `if not do_sample` above, so here `elif` implicitly implies `and do_sample`.
created_at: 2025-03-12T11:40:48Z
author_association: COLLABORATOR

PR comments:
user: login: vermouth1992
created_at: 2025-03-12T07:21:05Z
body: When eval using n samples, shall we also record pass@n metrics?
author_association: COLLABORATOR
user: login: PeterSH6
created_at: 2025-03-12T07:23:02Z
body: @vermouth1992 Yes, we have an implementation of pass@n logging. Will push them in the upcoming PRs
author_association: COLLABORATOR