Repository: volcengine/verl
Pull request #871: Batch Rewards
Author: login: mertunsall
Created: 2025-04-01T22:45:01Z | Merged: 2025-04-07T03:15:49Z



Diff:
diff --git a/tests/sanity/check_license.py b/tests/sanity/check_license.py
index 076ca03f748..99dc1d048f6 100644
--- a/tests/sanity/check_license.py
+++ b/tests/sanity/check_license.py
@@ -16,8 +16,8 @@
 license_head_bytedance_25 = "Copyright 2025 Bytedance Ltd. and/or its affiliates"
 # Add custom license headers below
 license_head_prime = "Copyright 2024 PRIME team and/or its affiliates"
-
-license_headers = [license_head_bytedance, license_head_bytedance_25, license_head_prime]
+license_head_individual = "Copyright 2025 Individual Contributor:"
+license_headers = [license_head_bytedance, license_head_bytedance_25, license_head_prime, license_head_individual]
 
 from pathlib import Path
 from argparse import ArgumentParser
diff --git a/verl/trainer/main_eval.py b/verl/trainer/main_eval.py
index 40e871cf86a..3868e5a61a2 100644
--- a/verl/trainer/main_eval.py
+++ b/verl/trainer/main_eval.py
@@ -27,8 +27,7 @@
 
 
 def get_custom_reward_fn(config):
-    import importlib.util, os
-
+    import importlib.util, os, sys
     reward_fn_config = config.get("custom_reward_function") or {}
     file_path = reward_fn_config.get("path")
     if not file_path:
@@ -40,18 +39,24 @@ def get_custom_reward_fn(config):
     spec = importlib.util.spec_from_file_location("custom_module", file_path)
     module = importlib.util.module_from_spec(spec)
     try:
+        sys.modules["custom_module"] = module
         spec.loader.exec_module(module)
     except Exception as e:
         raise RuntimeError(f"Error loading module from '{file_path}': {e}")
 
     function_name = reward_fn_config.get("name")
-
     if not hasattr(module, function_name):
         raise AttributeError(f"Reward function '{function_name}' not found in '{file_path}'.")
 
     print(f"using customized reward function '{function_name}' from '{file_path}'")
+    raw_fn = getattr(module, function_name)
+
+    reward_kwargs = dict(reward_fn_config.get("reward_kwargs", {}))
+
+    def wrapped_fn(*args, **kwargs):
+        return raw_fn(*args, **kwargs, **reward_kwargs)
 
-    return getattr(module, function_name)
+    return wrapped_fn
 
 
 @ray.remote
diff --git a/verl/trainer/main_ppo.py b/verl/trainer/main_ppo.py
index ae4abffab8e..773f230aabb 100644
--- a/verl/trainer/main_ppo.py
+++ b/verl/trainer/main_ppo.py
@@ -22,8 +22,7 @@
 
 
 def get_custom_reward_fn(config):
-    import importlib.util, os
-
+    import importlib.util, sys
     reward_fn_config = config.get("custom_reward_function") or {}
     file_path = reward_fn_config.get("path")
     if not file_path:
@@ -35,18 +34,24 @@ def get_custom_reward_fn(config):
     spec = importlib.util.spec_from_file_location("custom_module", file_path)
     module = importlib.util.module_from_spec(spec)
     try:
+        sys.modules["custom_module"] = module
         spec.loader.exec_module(module)
     except Exception as e:
         raise RuntimeError(f"Error loading module from '{file_path}': {e}")
 
     function_name = reward_fn_config.get("name")
-
     if not hasattr(module, function_name):
         raise AttributeError(f"Reward function '{function_name}' not found in '{file_path}'.")
 
     print(f"using customized reward function '{function_name}' from '{file_path}'")
+    raw_fn = getattr(module, function_name)
+
+    reward_kwargs = dict(reward_fn_config.get("reward_kwargs", {}))
+
+    def wrapped_fn(*args, **kwargs):
+        return raw_fn(*args, **kwargs, **reward_kwargs)
 
-    return getattr(module, function_name)
+    return wrapped_fn
 
 
 @hydra.main(config_path='config', config_name='ppo_trainer', version_base=None)
@@ -152,6 +157,9 @@ def run(self, config):
         elif reward_manager_name == 'prime':
             from verl.workers.reward_manager import PrimeRewardManager
             reward_manager_cls = PrimeRewardManager
+        elif reward_manager_name == 'batch':
+            from verl.workers.reward_manager import BatchRewardManager
+            reward_manager_cls = BatchRewardManager
         elif reward_manager_name == 'dapo':
             from verl.workers.reward_manager import DAPORewardManager
             reward_manager_cls = DAPORewardManager
@@ -160,10 +168,12 @@ def run(self, config):
             raise NotImplementedError
 
         compute_score = get_custom_reward_fn(config)
+        reward_kwargs = dict(config.reward_model.get("reward_kwargs", {}))
         reward_fn = reward_manager_cls(tokenizer=tokenizer,
                                        num_examine=0,
                                        compute_score=compute_score,
-                                       reward_fn_key=config.data.reward_fn_key)
+                                       reward_fn_key=config.data.reward_fn_key,
+                                       **reward_kwargs)
 
         # Note that we always use function-based RM for validation
         val_reward_fn = reward_manager_cls(tokenizer=tokenizer,
diff --git a/verl/utils/reward_score/math_batch.py b/verl/utils/reward_score/math_batch.py
new file mode 100644
index 00000000000..7c5425a335f
--- /dev/null
+++ b/verl/utils/reward_score/math_batch.py
@@ -0,0 +1,25 @@
+# Copyright 2025 Individual Contributor: Mert Unsal
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .math import compute_score
+
+
+def compute_score_batched(data_sources, solution_strs, ground_truths, extra_infos):
+    """
+    This is a demonstration of how the batched reward function should look like.
+    Typically, you want to use batched reward to speed up the process with parallelization
+    """
+    return [
+        compute_score(solution_str, ground_truth) for solution_str, ground_truth in zip(solution_strs, ground_truths)
+    ]
diff --git a/verl/workers/reward_manager/__init__.py b/verl/workers/reward_manager/__init__.py
index 6e85a388cbe..04300d94974 100644
--- a/verl/workers/reward_manager/__init__.py
+++ b/verl/workers/reward_manager/__init__.py
@@ -14,4 +14,5 @@
 
 from .naive import NaiveRewardManager
 from .prime import PrimeRewardManager
+from .batch import BatchRewardManager
 from .dapo import DAPORewardManager
diff --git a/verl/workers/reward_manager/batch.py b/verl/workers/reward_manager/batch.py
new file mode 100644
index 00000000000..ed92f9c4f91
--- /dev/null
+++ b/verl/workers/reward_manager/batch.py
@@ -0,0 +1,107 @@
+# Copyright 2025 Individual Contributor: Mert Unsal
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from verl import DataProto
+from collections import defaultdict
+
+
+class BatchRewardManager:
+
+    def __init__(self, tokenizer, num_examine, compute_score, reward_fn_key='data_source', **reward_kwargs):
+        self.tokenizer = tokenizer
+        self.num_examine = num_examine
+        self.compute_score = compute_score
+        self.reward_fn_key = reward_fn_key
+        self.reward_kwargs = reward_kwargs
+
+    def verify(self, data):
+        prompt_ids = data.batch['prompts']
+        response_ids = data.batch['responses']
+        attention_mask = data.batch['attention_mask']
+
+        prompt_len = prompt_ids.shape[-1]
+        valid_response_lengths = attention_mask[:, prompt_len:].sum(dim=-1)
+
+        responses_str = []
+        for i in range(len(data)):
+            valid_len = valid_response_lengths[i]
+            valid_response_ids = response_ids[i][:valid_len]
+            response_str = self.tokenizer.decode(valid_response_ids, skip_special_tokens=True)
+            responses_str.append(response_str)
+
+        ground_truths = [item.non_tensor_batch['reward_model'].get('ground_truth', None) for item in data]
+        data_sources = data.non_tensor_batch[self.reward_fn_key]
+        extras = data.non_tensor_batch.get('extra_info', [None] * len(data))
+
+        scores = self.compute_score(data_sources=data_sources,
+                                    solution_strs=responses_str,
+                                    ground_truths=ground_truths,
+                                    extra_infos=extras,
+                                    **self.reward_kwargs)
+
+        return scores
+
+    def __call__(self, data: DataProto, return_dict=False):
+
+        # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn
+        if 'rm_scores' in data.batch.keys():
+            if return_dict:
+                return {"reward_tensor": data.batch['rm_scores']}
+            else:
+                return data.batch['rm_scores']
+
+        reward_tensor = torch.zeros_like(data.batch['responses'], dtype=torch.float32)
+        reward_extra_info = defaultdict(list)
+        prompt_ids = data.batch['prompts']
+        prompt_len = prompt_ids.shape[-1]
+        attention_mask = data.batch['attention_mask']
+        valid_response_lengths = attention_mask[:, prompt_len:].sum(dim=-1)
+        data_sources = data.non_tensor_batch[self.reward_fn_key]
+
+        scores = self.verify(data)
+        rewards = []
+        already_printed = {}
+
+        for i in range(len(data)):
+            length = valid_response_lengths[i].item()
+            score = scores[i]
+
+            if isinstance(score, dict):
+                reward = score["score"]
+                for key, value in score.items():
+                    reward_extra_info[key].append(value)
+            else:
+                reward = score
+
+            rewards.append(reward)
+            reward_tensor[i, length - 1] = reward
+
+            data_source = data_sources[i]
+            if already_printed.get(data_source, 0) < self.num_examine:
+                response_str = self.tokenizer.decode(data.batch['responses'][i][:length], skip_special_tokens=True)
+                prompt_str = self.tokenizer.decode(data.batch['prompts'][i], skip_special_tokens=True)
+                ground_truth = data[i].non_tensor_batch['reward_model'].get('ground_truth', None)
+                print("[prompt]", prompt_str)
+                print("[response]", response_str)
+                print("[ground_truth]", ground_truth)
+                print("[score]", scores[i])
+                already_printed[data_source] = already_printed.get(data_source, 0) + 1
+
+        data.batch['acc'] = torch.tensor(rewards, dtype=torch.float32, device=prompt_ids.device)
+
+        if return_dict:
+            return {"reward_tensor": reward_tensor, "reward_extra_info": reward_extra_info}
+        else:
+            return reward_tensor
diff --git a/verl/workers/reward_manager/dapo.py b/verl/workers/reward_manager/dapo.py
index f98558d998a..a5ef53e5a79 100644
--- a/verl/workers/reward_manager/dapo.py
+++ b/verl/workers/reward_manager/dapo.py
@@ -45,7 +45,7 @@ def __call__(self, data: DataProto, return_dict: bool = False):
         # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn
         if 'rm_scores' in data.batch.keys():
             if return_dict:
-                return {"reward": data.batch['rm_scores']}
+                return {"reward_tensor": data.batch['rm_scores']}
             else:
                 return data.batch['rm_scores']
 
diff --git a/verl/workers/reward_manager/naive.py b/verl/workers/reward_manager/naive.py
index 8dd1684775a..ed50206f80e 100644
--- a/verl/workers/reward_manager/naive.py
+++ b/verl/workers/reward_manager/naive.py
@@ -34,7 +34,7 @@ def __call__(self, data: DataProto, return_dict=False):
         # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn
         if 'rm_scores' in data.batch.keys():
             if return_dict:
-                return {"reward": data.batch['rm_scores']}
+                return {"reward_tensor": data.batch['rm_scores']}
             else:
                 return data.batch['rm_scores']
 


Reviews:
user: login: hiyouga
state: DISMISSED
author_association: COLLABORATOR
submitted_at: 2025-04-02T05:35:42Z
user: login: hiyouga
state: DISMISSED
author_association: COLLABORATOR
submitted_at: 2025-04-03T09:59:47Z
user: login: hiyouga
state: DISMISSED
author_association: COLLABORATOR
submitted_at: 2025-04-04T06:00:12Z
user: login: patrik-bartak
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T23:18:11Z
user: login: mertunsall
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-05T08:55:15Z
user: login: hiyouga
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-07T03:15:38Z
user: login: eric-haibin-lin
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-07T03:49:04Z
user: login: hiyouga
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-07T05:32:03Z
user: login: Django-Jiang
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-07T07:53:47Z
user: login: patrik-bartak
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-07T15:47:41Z

Review comments:
diff_hunk: @@ -35,6 +35,7 @@ def get_custom_reward_fn(config):
     spec = importlib.util.spec_from_file_location("custom_module", file_path)
     module = importlib.util.module_from_spec(spec)
     try:
+        sys.modules["custom_module"] = module
path: verl/trainer/main_ppo.py
user: login: patrik-bartak
body: This is a bugfix right? Without this line, I get errors when using the Prime reward manager + a custom reward function. I wanted to submit an issue and fix it but I see you ran into the same problem. Correct me if I am wrong.

Maybe then this line should also be added to the get_custom_reward_fn inside main_eval.py?
created_at: 2025-04-04T23:18:10Z
author_association: CONTRIBUTOR
diff_hunk: @@ -35,6 +35,7 @@ def get_custom_reward_fn(config):
     spec = importlib.util.spec_from_file_location("custom_module", file_path)
     module = importlib.util.module_from_spec(spec)
     try:
+        sys.modules["custom_module"] = module
path: verl/trainer/main_ppo.py
user: login: mertunsall
body: yes, thanks for pointing this out. I am adding some new functionalities for adding kwargs to the reward function as well, I will add this once everything is tested
created_at: 2025-04-05T08:55:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,107 @@
+# Copyright 2025 Individual Contributor: Mert Unsal
path: verl/workers/reward_manager/batch.py
user: login: eric-haibin-lin
body: Thanks for the PR. Could you explain when & how this reward manager can be used? 
created_at: 2025-04-07T03:49:04Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,107 @@
+# Copyright 2025 Individual Contributor: Mert Unsal
path: verl/workers/reward_manager/batch.py
user: login: hiyouga
body: Using batch rewards, we can use multiprocessing to accelerate the computation of reward scores
created_at: 2025-04-07T05:32:03Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,107 @@
+# Copyright 2025 Individual Contributor: Mert Unsal
path: verl/workers/reward_manager/batch.py
user: login: Django-Jiang
body: hi, I just wonder what is the different of batch reward manager compared to prime
created_at: 2025-04-07T07:53:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,107 @@
+# Copyright 2025 Individual Contributor: Mert Unsal
path: verl/workers/reward_manager/batch.py
user: login: patrik-bartak
body: There can be situations where computing the reward for a batch using one call is much faster than performing a call for each sample. Especially if your reward is implemented as a service. 
created_at: 2025-04-07T15:47:41Z
author_association: CONTRIBUTOR

PR comments:
user: login: CLAassistant
created_at: 2025-04-01T22:45:07Z
body: [![CLA assistant check](https://cla-assistant.io/pull/badge/signed)](https://cla-assistant.io/volcengine/verl?pullRequest=871) <br/>All committers have signed the CLA.
author_association: NONE
user: login: hiyouga
created_at: 2025-04-02T05:44:11Z
body: remember to run `bash scripts/format.sh`
author_association: COLLABORATOR
user: login: mertunsall
created_at: 2025-04-02T12:47:19Z
body: Finished reformatting.
author_association: CONTRIBUTOR
user: login: mertunsall
created_at: 2025-04-02T19:14:22Z
body: I synced with main as well @hiyouga 
author_association: CONTRIBUTOR
user: login: mertunsall
created_at: 2025-04-02T20:43:52Z
body: There was a test not passing due to missing license, I added it.
author_association: CONTRIBUTOR
user: login: mertunsall
created_at: 2025-04-03T08:47:27Z
body: @hiyouga can you explain why the checks are not passing?
author_association: CONTRIBUTOR
user: login: hiyouga
created_at: 2025-04-03T09:02:56Z
body: @mertunsall the license should start with https://github.com/volcengine/verl/blob/8cae42dc29736d0802ded43c5ecf67a809d56bd8/verl/__init__.py#L1
author_association: COLLABORATOR
user: login: mertunsall
created_at: 2025-04-03T09:04:22Z
body: are you sure? this one is different: https://github.com/volcengine/verl/blob/8cae42dc29736d0802ded43c5ecf67a809d56bd8/verl/workers/reward_manager/prime.py

There is also another test that doesn’t pass.
author_association: CONTRIBUTOR
user: login: hiyouga
created_at: 2025-04-03T09:07:27Z
body: @mertunsall you can check this script https://github.com/volcengine/verl/blob/8cae42dc29736d0802ded43c5ecf67a809d56bd8/tests/sanity/check_license.py
author_association: COLLABORATOR
user: login: mertunsall
created_at: 2025-04-03T09:15:06Z
body: I made a fix for this, license check seems to pass right now. I added an option for an individual contributor license.

Do you have any insights on the other test that doesn't pass? This one doesn't seem related to my changes.
author_association: CONTRIBUTOR
user: login: mertunsall
created_at: 2025-04-03T22:11:55Z
body: resolved merge conflicts.
author_association: CONTRIBUTOR
user: login: mertunsall
created_at: 2025-04-04T13:37:34Z
body: I realized some more changes were brought by DAPO release - I fixed them and also added an example.
author_association: CONTRIBUTOR
user: login: mertunsall
created_at: 2025-04-06T13:40:04Z
body: One check was cancelled for some reason - but I think this should be ready to merge @hiyouga 

author_association: CONTRIBUTOR
user: login: hiyouga
created_at: 2025-04-07T03:10:59Z
body: Yeah let's merge it
author_association: COLLABORATOR
user: login: patrik-bartak
created_at: 2025-04-08T16:28:29Z
body: @mertunsall 
Hey I just wanted to ask here in case you already know the answer before I start digging into the code, but the data_sources within a batch can be different right? For example, if we configure multiple evaluation datasets, we have to make sure that the reward function is able to handle a batch where data comes from different data sources in that batch. From what I can see, the entire evaluation dataloader is one batch. 

author_association: CONTRIBUTOR
user: login: mertunsall
created_at: 2025-04-08T17:49:06Z
body: Yes, that sounds right to me!

On Tue 8 Apr 2025 at 18:28, Patrik Bartak ***@***.***> wrote:

> @mertunsall <https://github.com/mertunsall>
> Hey I just wanted to ask here in case you already know the answer before I
> start digging into the code, but the data_sources within a batch can be
> different right? For example, if we configure multiple evaluation datasets,
> we have to make sure that the reward function is able to handle a batch
> where data comes from different data sources in that batch.
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/volcengine/verl/pull/871#issuecomment-2787020526>, or
> unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AFIEPR2ZW5E3WNV5YQW6ZML2YP2MHAVCNFSM6AAAAAB2ICJEG2VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDOOBXGAZDANJSGY>
> .
> You are receiving this because you were mentioned.Message ID:
> ***@***.***>
> *patrik-bartak* left a comment (volcengine/verl#871)
> <https://github.com/volcengine/verl/pull/871#issuecomment-2787020526>
>
> @mertunsall <https://github.com/mertunsall>
> Hey I just wanted to ask here in case you already know the answer before I
> start digging into the code, but the data_sources within a batch can be
> different right? For example, if we configure multiple evaluation datasets,
> we have to make sure that the reward function is able to handle a batch
> where data comes from different data sources in that batch.
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/volcengine/verl/pull/871#issuecomment-2787020526>, or
> unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AFIEPR2ZW5E3WNV5YQW6ZML2YP2MHAVCNFSM6AAAAAB2ICJEG2VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDOOBXGAZDANJSGY>
> .
> You are receiving this because you were mentioned.Message ID:
> ***@***.***>
>

author_association: CONTRIBUTOR
user: login: chenhaiq
created_at: 2025-05-06T07:42:57Z
body: @mertunsall Do you have an example about how to use batch rewards?
Since _default_compute_score do not accept datastore in list, it always fails in my job.
author_association: COLLABORATOR
user: login: xin-ran-w
created_at: 2025-07-08T15:21:07Z
body: How to use this reward manager in config.
author_association: NONE