Repository: huggingface/trl
Pull request: 🏆 Custom reward function for GRPO and shiny doc
# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly. They may suggest changes to make the code even better.
-->

<!-- Remove if not applicable -->

Fixes #2599


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/trl/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a GitHub issue? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/trl/tree/main/docs).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

Problem statement:
Add the training method for DeepSeek-R1
### Feature request

The pure RL method that has been used in training DeepSeek-R1 seems very cool and i think the community will apprechiate it if we can have our hands on it.

### Motivation

DeepSeek-R1 is better in a lot of areas

### Your contribution

If there is something i can help with i will do it

Hints from issue discussion:
The training method for DeepSeek-R1 is GRPO, right?
> The training method for DeepSeek-R1 is GRPO, right?

Yeah, but they don't use a reward model.

Commit message: it's probably the best of both worlds [ci skip]

Patch to review:
diff --git a/docs/source/grpo_trainer.md b/docs/source/grpo_trainer.md
index 59abe04356..a304f52610 100644
--- a/docs/source/grpo_trainer.md
+++ b/docs/source/grpo_trainer.md
@@ -114,6 +114,78 @@ The GRPO Trainer logs the following metrics:
 - `reward_std` : The average standard deviation within reward groups.
 - `kl` : The average KL divergence between the model and the reference model calculated on completions.
 
+## Customization
+
+### Using a custom reward function
+
+The [`GRPOTrainer`] supports using custom reward function. To ensure compatibility, your reward function must satisfy the following requirements:
+
+1. **Input arguments**:
+   - The function must accept two arguments: `prompts` and `completions`.
+   - Depending on the dataset format, the input will vary:
+     - For [standard format](dataset_formats#standard), `prompts` and `completions` will be lists of strings.
+     - For [conversational format](dataset_formats#conversational), `prompts` and `completions` will be lists of message dictionaries.
+
+2. **Return value**: The function must return a list of floats. Each float represents the reward corresponding to a single completion.
+
+#### Example 1: Standard format
+
+Below is an example of a reward function for a standard format that rewards longer completions:
+
+```python
+def reward_func(prompts, completions):
+    """Reward function that gives higher scores to longer completions."""
+    return [float(len(completion)) for completion in completions]
+```
+
+You can test it as follows:
+
+```python
+>>> prompts = ["The sky is", "The sun is"]
+>>> completions = [" blue.", " in the sky."]
+>>> print(reward_func(prompts, completions))
+[6.0, 12.0]
+```
+
+#### Example 2: Conversational format
+
+For conversational format, completions consist of structured messages. Here’s an example that rewards longer completion content:
+
+```python
+def reward_func(prompts, completions):
+    """Reward function that gives higher scores to longer completion content."""
+    completion_contents = [completion[0]["content"] for completion in completions]
+    return [float(len(content)) for content in completion_contents]
+```
+
+You can test this function as follows:
+
+```python
+>>> prompts = [
+...     [{"role": "user", "content": "What color is the sky?"}],
+...     [{"role": "user", "content": "Where is the sun?"}],
+... ]
+>>> completions = [
+...     [{"role": "assistant", "content": "It is blue."}],
+...     [{"role": "assistant", "content": "In the sky."}],
+... ]
+>>> print(reward_func(prompts, completions))
+[11.0, 11.0]
+```
+
+#### Passing the Reward Function to the Trainer
+
+To use your custom reward function, pass it to the `GRPOTrainer` as follows:
+
+```python
+from trl import GRPOTrainer
+
+trainer = GRPOTrainer(
+    reward_func=reward_func,
+    ...,
+)
+```
+
 ## GRPOTrainer
 
 [[autodoc]] GRPOTrainer
diff --git a/tests/test_grpo_trainer.py b/tests/test_grpo_trainer.py
index 1afd490bd1..d1d6a36248 100644
--- a/tests/test_grpo_trainer.py
+++ b/tests/test_grpo_trainer.py
@@ -146,3 +146,72 @@ def test_training_different_reward_model(self):
             for n, param in previous_trainable_params.items():
                 new_param = trainer.model.get_parameter(n)
                 self.assertFalse(torch.equal(param, new_param), f"Parameter {n} has not changed.")
+
+    def test_training_reward_func_standard(self):
+        dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train")
+
+        def reward_func(prompts, completions):
+            """Reward function that rewards longer completions."""
+            return [float(len(completion)) for completion in completions]
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            training_args = GRPOConfig(
+                output_dir=tmp_dir,
+                learning_rate=0.1,  # increase the learning rate to speed up the test
+                per_device_train_batch_size=2,  # reduce the batch size to reduce memory usage
+                num_generations=3,  # reduce the number of generations to reduce memory usage
+                max_completion_length=32,  # reduce the completion length to reduce memory usage
+                report_to="none",
+            )
+            trainer = GRPOTrainer(
+                model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",
+                reward_model=reward_func,
+                args=training_args,
+                train_dataset=dataset,
+            )
+
+            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}
+
+            trainer.train()
+
+            self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
+
+            # Check the params have changed
+            for n, param in previous_trainable_params.items():
+                new_param = trainer.model.get_parameter(n)
+                self.assertFalse(torch.equal(param, new_param), f"Parameter {n} has not changed.")
+
+    def test_training_reward_func_conversational(self):
+        dataset = load_dataset("trl-internal-testing/zen", "conversational_prompt_only", split="train")
+
+        def reward_func(prompts, completions):
+            """Reward function that gives higher scores to longer completion content."""
+            completion_contents = [completion[0]["content"] for completion in completions]
+            return [float(len(content)) for content in completion_contents]
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            training_args = GRPOConfig(
+                output_dir=tmp_dir,
+                learning_rate=0.1,  # increase the learning rate to speed up the test
+                per_device_train_batch_size=2,  # reduce the batch size to reduce memory usage
+                num_generations=3,  # reduce the number of generations to reduce memory usage
+                max_completion_length=32,  # reduce the completion length to reduce memory usage
+                report_to="none",
+            )
+            trainer = GRPOTrainer(
+                model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",
+                reward_model=reward_func,
+                args=training_args,
+                train_dataset=dataset,
+            )
+
+            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}
+
+            trainer.train()
+
+            self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
+
+            # Check the params have changed
+            for n, param in previous_trainable_params.items():
+                new_param = trainer.model.get_parameter(n)
+                self.assertFalse(torch.equal(param, new_param), f"Parameter {n} has not changed.")
diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index 579f55ec38..664fd62463 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -17,7 +17,6 @@
 from typing import Any, Callable, Optional, Union
 
 import torch
-import torch.nn as nn
 import torch.utils.data
 import transformers
 from datasets import Dataset, IterableDataset
@@ -26,8 +25,6 @@
     AutoModelForCausalLM,
     AutoModelForSequenceClassification,
     AutoTokenizer,
-    DataCollator,
-    EvalPrediction,
     GenerationConfig,
     PreTrainedModel,
     PreTrainedTokenizerBase,
@@ -51,22 +48,89 @@
 
 
 class GRPOTrainer(Trainer):
+    """
+    Trainer for the Group Relative Policy Optimization (GRPO) method. This algorithm was initially proposed in the
+    paper [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://huggingface.co/papers/2402.03300).
+
+    Example:
+
+    ```python
+    from datasets import load_dataset
+    from trl import GRPOTrainer
+
+    dataset = load_dataset("trl-lib/tldr", split="train")
+
+    trainer = GRPOTrainer(
+        model="Qwen/Qwen2-0.5B-Instruct",
+        reward_model="weqweasdas/RM-Gemma-2B",
+        train_dataset=dataset,
+    )
+
+    trainer.train()
+    ```
+
+    Args:
+        model (`Union[str, PreTrainedModel]`):
+            Model to be trained. Can be either:
+
+            - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or
+              a path to a *directory* containing model weights saved using
+              [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is
+              loaded using [`~transformers.AutoModelForCausalLM.from_pretrained`] with the keywork arguments
+              in `args.model_init_kwargs`.
+            - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported.
+        reward_model (`Union[str, PreTrainedModel, Callable[[list, list], list[float]]`):
+            Reward model to be used for computing the rewards. Can be either:
+
+            - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or
+              a path to a *directory* containing model weights saved using
+              [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is
+              loaded using [`~transformers.AutoModelForSequenceClassification.from_pretrained`] with `num_labels=1` and
+              the keywork arguments in `args.model_init_kwargs`.
+            - A [`~transformers.PreTrainedModel`] object. Only sequence classification models are supported.
+            - A custom reward function that takes a list of prompts and completions and returns a list of rewards. For
+              more details, see [Using a custom reward function](#using-a-custom-reward-function).
+        args ([`GRPOConfig`], *optional*, defaults to `None`):
+            Configuration for this trainer. If `None`, a default configuration is used.
+        train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]):
+            Dataset to use for training. It must include a column `"prompt"`. Any additional columns in the dataset is
+            ignored. The format of the samples can be either:
+
+            - [Standard](dataset_formats#standard): Each sample contains plain text.
+            - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role
+              and content).
+        eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Union[Dataset, IterableDataset]]`):
+            Dataset to use for evaluation. It must meet the same requirements as `train_dataset`.
+        processing_class ([`~transformers.PreTrainedTokenizerBase`], *optional*, defaults to `None`):
+            Processing class used to process the data. The padding side must be set to "left". If `None`, the
+            processing class is loaded from the model's name with [`~transformers.AutoTokenizer.from_pretrained`].
+        reward_processing_class ([`~transformers.PreTrainedTokenizerBase`], *optional*, defaults to `None`):
+            Processing class used to process the data for the reward model. If `None`, the processing class is loaded
+            from the reward model's name with [`~transformers.AutoTokenizer.from_pretrained`].
+        callbacks (list of [`~transformers.TrainerCallback`], *optional*, defaults to `None`):
+            List of callbacks to customize the training loop. Will add those to the list of default callbacks
+            detailed in [here](https://huggingface.co/docs/transformers/main_classes/callback).
+
+            If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`]
+            method.
+        optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`, *optional*, defaults to `(None, None)`):
+            A tuple containing the optimizer and the scheduler to use. Will default to an instance of [`AdamW`] on your
+            model and a scheduler given by [`get_linear_schedule_with_warmup`] controlled by `args`.
+        peft_config ([`~peft.PeftConfig`], *optional*, defaults to `None`):
+            PEFT configuration used to wrap the model. If `None`, the model is not wrapped.
+    """
+
     def __init__(
         self,
-        model: Union[str, PreTrainedModel, nn.Module] = None,
-        reward_model: Optional[Union[PreTrainedModel, nn.Module]] = None,
+        model: Union[str, PreTrainedModel],
+        reward_model: Union[str, PreTrainedModel, Callable[[list, list], list[float]]],
         args: GRPOConfig = None,
-        data_collator: Optional[DataCollator] = None,
         train_dataset: Optional[Union[Dataset, IterableDataset]] = None,
         eval_dataset: Optional[Union[Dataset, IterableDataset, dict[str, Union[Dataset, IterableDataset]]]] = None,
         processing_class: Optional[PreTrainedTokenizerBase] = None,
         reward_processing_class: Optional[PreTrainedTokenizerBase] = None,
-        model_init: Optional[Callable[[], PreTrainedModel]] = None,
-        compute_loss_func: Optional[Callable] = None,
-        compute_metrics: Optional[Callable[[EvalPrediction], dict]] = None,
         callbacks: Optional[list[TrainerCallback]] = None,
         optimizers: tuple[Optional[torch.optim.Optimizer], Optional[torch.optim.lr_scheduler.LambdaLR]] = (None, None),
-        preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,
         peft_config: Optional["PeftConfig"] = None,
     ):
         # Args
@@ -119,23 +183,27 @@ def __init__(
             reward_model = AutoModelForSequenceClassification.from_pretrained(
                 reward_model, num_labels=1, **model_init_kwargs
             )
+        # The GRPO implementation supports either a pretrained reward model (sequence classifier) or a custom reward
+        # function. A private attribute is used to indicate the type of the reward model.
+        self._reward_model_type = "pretrained" if isinstance(reward_model, PreTrainedModel) else "function"
         self.reward_model = reward_model
 
         # Reward processing class
-        if reward_processing_class is None:
-            reward_processing_class = AutoTokenizer.from_pretrained(reward_model.config._name_or_path)
-        if reward_processing_class.pad_token_id is None:
-            reward_processing_class.pad_token = reward_processing_class.eos_token
-        self.reward_processing_class = reward_processing_class
-        # The reward model computes the reward for the latest non-padded token in the input sequence.
-        # So it's important to set the pad token ID to the padding token ID of the processing class.
-        self.reward_model.config.pad_token_id = reward_processing_class.pad_token_id
-
-        # Data loading and preprocessing
-        if data_collator is None:
-
-            def data_collator(features):  # No data collation is needed in GRPO
-                return features
+        if self._reward_model_type == "pretrained":
+            if reward_processing_class is None:
+                reward_processing_class = AutoTokenizer.from_pretrained(reward_model.config._name_or_path)
+            if reward_processing_class.pad_token_id is None:
+                reward_processing_class.pad_token = reward_processing_class.eos_token
+            self.reward_processing_class = reward_processing_class
+            # The reward model computes the reward for the latest non-padded token in the input sequence.
+            # So it's important to set the pad token ID to the padding token ID of the processing class.
+            self.reward_model.config.pad_token_id = reward_processing_class.pad_token_id
+        elif self._reward_model_type == "function":
+            self.reward_processing_class = None
+
+        # Data collator
+        def data_collator(features):  # No data collation is needed in GRPO
+            return features
 
         # Training arguments
         self.max_prompt_length = args.max_prompt_length
@@ -168,17 +236,15 @@ def data_collator(features):  # No data collation is needed in GRPO
             train_dataset=train_dataset,
             eval_dataset=eval_dataset,
             processing_class=processing_class,
-            model_init=model_init,
-            compute_loss_func=compute_loss_func,
-            compute_metrics=compute_metrics,
             callbacks=callbacks,
             optimizers=optimizers,
-            preprocess_logits_for_metrics=preprocess_logits_for_metrics,
         )
 
         if self.ref_model is not None:
             self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)
-        self.reward_model = self.accelerator.prepare_model(self.reward_model, evaluation_mode=True)
+
+        if self._reward_model_type == "pretrained":
+            self.reward_model = self.accelerator.prepare_model(self.reward_model, evaluation_mode=True)
 
     def _set_signature_columns_if_needed(self):
         # If `self.args.remove_unused_columns` is True, non-signature columns are removed.
@@ -252,24 +318,28 @@ def get_per_token_logps(model, input_ids):
 
         # Decode the generated completions
         completions = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True)
+        if is_conversational(inputs[0]):
+            completions = [[{"role": "assistant", "content": completion}] for completion in completions]
 
         # Compute the rewards
         prompts = [prompt for prompt in prompts for _ in range(self.num_generations)]
-        if is_conversational(inputs[0]):
-            completions = [[{"role": "assistant", "content": completion}] for completion in completions]
-            messages = [{"messages": p + c} for p, c in zip(prompts, completions)]
-            texts = [apply_chat_template(x, self.reward_processing_class)["text"] for x in messages]
-            reward_inputs = self.reward_processing_class(
-                texts, return_tensors="pt", padding=True, padding_side="right", add_special_tokens=False
-            )
-        else:
-            texts = [p + c for p, c in zip(prompts, completions)]
+
+        if self._reward_model_type == "pretrained":
+            if is_conversational(inputs[0]):
+                messages = [{"messages": p + c} for p, c in zip(prompts, completions)]
+                texts = [apply_chat_template(x, self.reward_processing_class)["text"] for x in messages]
+            else:
+                texts = [p + c for p, c in zip(prompts, completions)]
             reward_inputs = self.reward_processing_class(
                 texts, return_tensors="pt", padding=True, padding_side="right", add_special_tokens=False
             )
-        reward_inputs = super()._prepare_inputs(reward_inputs)
-        with torch.inference_mode():
-            rewards = self.reward_model(**reward_inputs).logits[:, 0]  # Shape (B*G,)
+            reward_inputs = super()._prepare_inputs(reward_inputs)
+            with torch.inference_mode():
+                rewards = self.reward_model(**reward_inputs).logits[:, 0]  # Shape (B*G,)
+
+        elif self._reward_model_type == "function":
+            rewards = self.reward_model(prompts, completions)
+            rewards = torch.tensor(rewards, dtype=torch.float, device=device)
 
         # Compute grouped-wise rewards
         mean_grouped_rewards = rewards.view(-1, self.num_generations).mean(dim=1)


Code review comments:
diff_hunk: @@ -51,22 +48,89 @@
 
 
 class GRPOTrainer(Trainer):
+    """
+    Trainer for the Group Relative Policy Optimization (GRPO) method. This algorithm was initially proposed in the
+    paper [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://huggingface.co/papers/2402.03300).
+
+    Example:
+
+    ```python
+    from datasets import load_dataset
+    from trl import GRPOTrainer
+
+    dataset = load_dataset("trl-lib/tldr", split="train")
+
+    trainer = GRPOTrainer(
+        model="Qwen/Qwen2-0.5B-Instruct",
+        reward_model="weqweasdas/RM-Gemma-2B",
+        train_dataset=dataset,
+    )
+
+    trainer.train()
+    ```
+
+    Args:
+        model (`Union[str, PreTrainedModel]`):
+            Model to be trained. Can be either:
+
+            - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or
+              a path to a *directory* containing model weights saved using
+              [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is
+              loaded using [`~transformers.AutoModelForCausalLM.from_pretrained`] with the keywork arguments
+              in `args.model_init_kwargs`.
+            - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported.
+        reward_model (`Union[str, PreTrainedModel, Callable[[list, list], list[float]]`):
path: trl/trainer/grpo_trainer.py
text: @user1:
nit: I feel like the `reward_model` naming might be a bit confusing for users, i.e. if they didn't read the docstring, they might assume this has to be a LLM.

Can we rename it to sth like `reward_provider` or `reward_accessor`?

@author:
> be a bit confusing for users, i.e. if they didn't read the docstring

Who the hell doesn't read the doc???

@author:
Just kidding, what about `reward_func`? A torch module is a callable, so somehow it's a function. Maybe less confusing

@user1:
sgtm 🤣