🍭 Custom reward function for RLOO (#2612)

* rloo custom reward function and test

* idont even know why i did that

* removing get_reward_custom

* remove get_reward_custom test

* fix code quality check

* adding test

* end this mysery already

* fix test

Author: August-murr

tests/test_rloo_trainer.py

@@ -172,3 +172,42 @@ def test_rloo_training(self):
 
             # Check if objective/rlhf_reward is available
             self.assertIn("objective/rlhf_reward", trainer.state.log_history[-1])
+
+    def test_rloo_training_with_custom_reward(self):
+        # dummy reward function
+        def reward_function(texts):
+            # based on length of text
+            rewards = [len(text) for text in texts]
+            return rewards
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            training_args = RLOOConfig(
+                output_dir=tmp_dir,
+                per_device_train_batch_size=2,
+                per_device_eval_batch_size=2,
+                total_episodes=1,
+                num_train_epochs=1,
+                max_steps=2,
+                report_to="none",
+            )
+
+            # Create a simple dataset
+            dummy_text = [{"content": "Hello World!", "role": "user"}]
+            dummy_data = self.tokenizer.apply_chat_template(dummy_text)
+            dummy_dataset = Dataset.from_dict({"input_ids": [dummy_data, dummy_data]})
+
+            trainer = RLOOTrainer(
+                config=training_args,
+                policy=self.policy_model,
+                reward_model=reward_function,
+                ref_policy=self.policy_ref_model,
+                processing_class=self.tokenizer,
+                train_dataset=dummy_dataset,
+                eval_dataset=dummy_dataset,
+            )
+
+            # Test that training completes without errors
+            trainer.train()
+
+            # Check if objective/rlhf_reward is available
+            self.assertIn("objective/rlhf_reward", trainer.state.log_history[-1])

trl/trainer/rloo_trainer.py

@@ -18,7 +18,7 @@
 import textwrap
 import time
 from collections import defaultdict
-from typing import Optional, Union
+from typing import Callable, Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -79,7 +79,7 @@ def __init__(
         ],
         policy: nn.Module,
         ref_policy: nn.Module,
-        reward_model: nn.Module,
+        reward_model: Union[nn.Module, Callable[[list[str]], list[float]]],
         train_dataset: Dataset,
         data_collator: Optional[DataCollatorWithPadding] = None,
         eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None,
@@ -152,7 +152,8 @@ def __init__(
         # setup model, optimizer, and others
         #########
         for module in [policy, ref_policy, reward_model]:
-            disable_dropout_in_model(module)
+            if isinstance(module, nn.Module):
+                disable_dropout_in_model(module)
         if args.stop_token and args.stop_token == "eos":
             args.stop_token_id = self.processing_class.eos_token_id
         self.model = policy
@@ -219,16 +220,18 @@ def __init__(
         self.eval_dataloader = accelerator.prepare(self.eval_dataloader)
 
         if self.is_deepspeed_enabled:
-            self.reward_model = prepare_deepspeed(
-                self.reward_model, args.per_device_train_batch_size, args.fp16, args.bf16
-            )
+            if isinstance(self.reward_model, nn.Module):
+                self.reward_model = prepare_deepspeed(
+                    self.reward_model, args.per_device_train_batch_size, args.fp16, args.bf16
+                )
             self.ref_policy = prepare_deepspeed(
                 self.ref_policy, args.per_device_train_batch_size, args.fp16, args.bf16
             )
             self.deepspeed = self.model
         else:
             self.ref_policy = self.ref_policy.to(self.accelerator.device)
-            self.reward_model = self.reward_model.to(self.accelerator.device)
+            if isinstance(self.reward_model, nn.Module):
+                self.reward_model = self.reward_model.to(self.accelerator.device)
 
     def get_train_dataloader(self) -> DataLoader:
         return self.dataloader
@@ -350,9 +353,18 @@ def repeat_generator():
                     # Response Processing 2. run reward model on the truncated responses
                     postprocessed_query_response = torch.cat((query, postprocessed_response), 1)
                     sequence_length = first_true_indices(postprocessed_response == processing_class.pad_token_id) - 1
-                    _, score, _ = get_reward(
-                        reward_model, postprocessed_query_response, processing_class.pad_token_id, context_length
-                    )
+
+                    if isinstance(reward_model, nn.Module):
+                        _, score, _ = get_reward(
+                            reward_model, postprocessed_query_response, processing_class.pad_token_id, context_length
+                        )
+                    else:
+                        score = torch.tensor(
+                            reward_model(
+                                processing_class.batch_decode(postprocessed_query_response, skip_special_tokens=True)
+                            ),
+                            dtype=torch.float,
+                        ).to(device)
 
                     # Store batch results
                     responses.append(response)
@@ -595,9 +607,21 @@ def generate_completions(self, sampling: bool = False):
                     )
 
                     postprocessed_query_response = torch.cat((query, postprocessed_response), 1)
-                    _, score, _ = get_reward(
-                        self.reward_model, postprocessed_query_response, processing_class.pad_token_id, context_length
-                    )
+
+                    if isinstance(self.reward_model, nn.Module):
+                        _, score, _ = get_reward(
+                            self.reward_model,
+                            postprocessed_query_response,
+                            processing_class.pad_token_id,
+                            context_length,
+                        )
+                    else:
+                        score = torch.tensor(
+                            self.reward_model(
+                                processing_class.batch_decode(postprocessed_query_response, skip_special_tokens=True)
+                            ),
+                            dtype=torch.float,
+                        ).to(postprocessed_query_response.device)
                     table["score"].extend(self.accelerator.gather_for_metrics(score).float().cpu().numpy())
 
                 if sampling: