GRPO Trainer

Overview

TRL supports the GRPO Trainer for training language models, as described in the paper DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models by Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Mingchuan Zhang, Y. K. Li, Y. Wu, Daya Guo.

The abstract from the paper is the following:

Mathematical reasoning poses a significant challenge for language models due to its complex and structured nature. In this paper, we introduce DeepSeekMath 7B, which continues pre-training DeepSeek-Coder-Base-v1.5 7B with 120B math-related tokens sourced from Common Crawl, together with natural language and code data. DeepSeekMath 7B has achieved an impressive score of 51.7% on the competition-level MATH benchmark without relying on external toolkits and voting techniques, approaching the performance level of Gemini-Ultra and GPT-4. Self-consistency over 64 samples from DeepSeekMath 7B achieves 60.9% on MATH. The mathematical reasoning capability of DeepSeekMath is attributed to two key factors: First, we harness the significant potential of publicly available web data through a meticulously engineered data selection pipeline. Second, we introduce Group Relative Policy Optimization (GRPO), a variant of Proximal Policy Optimization (PPO), that enhances mathematical reasoning abilities while concurrently optimizing the memory usage of PPO.

This post-training method was contributed by Quentin Gallouédec.

Quick start

This example demonstrates how to train a model using the GRPO method. We train a Qwen 0.5B Instruct model with the prompts from the TLDR dataset (completion column is ignored!). You can view the data in the dataset here:

Below is the script to train the model.

# train_grpo.py
from datasets import load_dataset
from trl import GRPOConfig, GRPOTrainer
dataset = load_dataset("trl-lib/tldr", split="train")
# Define the reward function, which rewards completions that are close to 20 characters
def reward_len(completions, **kwargs):
return [-abs(20 - len(completion)) for completion in completions]
training_args = GRPOConfig(output_dir="Qwen2-0.5B-GRPO", logging_steps=10)
trainer = GRPOTrainer(
model="Qwen/Qwen2-0.5B-Instruct",
reward_funcs=reward_len,
args=training_args,
train_dataset=dataset,
)
trainer.train()

Execute the script using the following command:

accelerate launch train_grpo.py

Distributed across 8 GPUs, the training takes approximately 1 day.

Looking deeper into the GRPO method

GRPO is an online learning algorithm, meaning it improves iteratively by using the data generated by the trained model itself during training. The intuition behind GRPO objective is to maximize the advantage of the generated completions, while ensuring that the model remains close to the reference policy. To understand how GRPO works, it can be broken down into four main steps: Generating completions, computing the advantage, estimating the KL divergence, and computing the loss.

Generating completions

At each training step, we sample a batch of prompts and generate a set of \( G \) completions for each prompt (denoted as \( o_i \)).

Computing the advantage

For each of the \( G \) sequences, we compute the reward using a reward model. To align with the comparative nature of reward models—typically trained on datasets of comparisons between outputs for the same question—the advantage is calculated to reflect these relative comparisons. It is normalized as follows:

This approach gives the method its name: Group Relative Policy Optimization (GRPO).

Estimating the KL divergence

KL divergence is estimated using the approximator introduced by Schulman et al. (2020). The approximator is defined as follows:

$$\mathbb{D}{\text{KL}}\left[\pi\theta |\pi_{\text{ref}}\right] = \frac{\pi_{\text{ref}}(o_{i,t} \mid q, o_{i,<t})}{\pi_\theta(o_{i,t} \mid q, o_{i,<t})} - \log \frac{\pi_{\text{ref}}(o_{i,t} \mid q, o_{i,<t})}{\pi_\theta(o_{i,t} \mid q, o_{i,<t})} - 1, $$

Computing the loss

The objective is to maximize the advantage while ensuring that the model remains close to the reference policy. Consequently, the loss is defined as follows:

$$ \mathcal{L}{\text{GRPO}}(\theta) = -\frac{1}{G} \sum{i=1}^G \frac{1}{|o_i|} \sum_{t=1}^{|o_i|} \left[ \frac{\pi_\theta(o_{i,t} \mid q, o_{i,< t})}{\left[\pi_\theta(o_{i,t} \mid q, o_{i,< t})\right]{\text{no grad}}} \hat{A}{i,t} - \beta \mathbb{D}{\text{KL}}\left[\pi\theta | \pi_{\text{ref}}\right] \right], $$

where the first term represents the scaled advantage and the second term penalizes deviations from the reference policy through KL divergence.

In the original paper, this formulation is generalized to account for multiple updates after each generation by leveraging the clipped surrogate objective:

$$ \mathcal{L}{\text{GRPO}}(\theta) = - \frac{1}{G} \sum{i=1}^G \frac{1}{|o_i|} \sum_{t=1}^{|o_i|} \left[ \min \left( \frac{\pi_\theta(o_{i,t} \mid q, o_{i,< t})}{\pi_{\theta_{\text{old}}}(o_{i,t} \mid q, o_{i,< t})} \hat{A}{i,t}, , \text{clip}\left( \frac{\pi\theta(o_{i,t} \mid q, o_{i,< t})}{\pi_{\theta_{\text{old}}}(o_{i,t} \mid q, o_{i,< t})}, 1 - \epsilon, 1 + \epsilon \right) \hat{A}{i,t} \right) - \beta \mathbb{D}{\text{KL}}\left[\pi_\theta | \pi_{\text{ref}}\right] \right], $$

where \(\text{clip}(\cdot, 1 - \epsilon, 1 + \epsilon) \) ensures that updates do not deviate excessively from the reference policy by bounding the policy ratio between \( 1 - \epsilon \) and \( 1 + \epsilon \). In TRL though, as in the original paper, we only do one update per generation, so we can simplify the loss to the first form.

Logged metrics

The GRPO Trainer logs the following metrics:

completion_length: The average completion length.
reward/{reward_func_name}: The reward computed by each reward function.
reward: The average reward.
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

Speed up training with vLLM-powered generation

Generation is often the main bottleneck that makes training slow with online methods. To accelerate generation, you can use vLLM, a library that enables fast generation. To enable it, pass use_vllm=True in the training arguments.

from trl import GRPOConfig
training_args = GRPOConfig(..., use_vllm=True)

For more information, see Speeding up training with vLLM.

Using a custom reward function

The [GRPOTrainer] supports using custom reward functions instead of dense reward models. To ensure compatibility, your reward function must satisfy the following requirements:

Input arguments:
- The function must accept the following as keyword arguments:
  - prompts (contains the prompts),
  - completions (contains the generated completions),
  - All columns names (but prompt) that the dataset may have. For example, if the dataset contains a column named ground_truth, the function will be called with ground_truth as a keyword argument.
  The easiest way to comply with this requirement is to use **kwargs in the function signature.
- Depending on the dataset format, the input will vary:
  - For standard format, prompts and completions will be lists of strings.
  - For conversational format, prompts and completions will be lists of message dictionaries.
Return value: The function must return a list of floats. Each float represents the reward corresponding to a single completion.

Example 1: Reward longer completions

Below is an example of a reward function for a standard format that rewards longer completions:

def reward_func(completions, **kwargs):
"""Reward function that gives higher scores to longer completions."""
return [float(len(completion)) for completion in completions]

You can test it as follows:

>>> prompts = ["The sky is", "The sun is"]
>>> completions = [" blue.", " in the sky."]
>>> print(reward_func(prompts=prompts, completions=completions))
[6.0, 12.0]

Example 2: Reward completions with specific format

Below is an example of a reward function that checks if the completion has a specific format. This example is inspired by the format reward function used in the paper DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. It is designed for conversational format, where prompts and completions consist of structured messages.

import re
def format_reward_func(completions, **kwargs):
"""Reward function that checks if the completion has a specific format."""
pattern = r"^<think>.*?</think><answer>.*?</answer>$"
completion_contents = [completion[0]["content"] for completion in completions]
matches = [re.match(pattern, content) for content in completion_contents]
return [1.0 if match else 0.0 for match in matches]

You can test this function as follows:

>>> prompts = [
... [{"role": "assistant", "content": "What is the result of (1 + 2) * 4?"}],
... [{"role": "assistant", "content": "What is the result of (3 + 1) * 2?"}],
... ]
>>> completions = [
... [{"role": "assistant", "content": "<think>The sum of 1 and 2 is 3, which we multiply by 4 to get 12.</think><answer>(1 + 2) * 4 = 12</answer>"}],
... [{"role": "assistant", "content": "The sum of 3 and 1 is 4, which we multiply by 2 to get 8. So (3 + 1) * 2 = 8."}],
... ]
>>> format_reward_func(prompts=prompts, completions=completions)
[1.0, 0.0]

Example 3: Reward completions based on a reference

Below is an example of a reward function that checks if the completion is correct. This example is inspired by the accuracy reward function used in the paper DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. This example is designed for standard format, where the dataset contains a column named ground_truth.

import re
def reward_func(completions, ground_truth, **kwargs):
# Regular expression to capture content inside \boxed{}
matches = [re.search(r"\\boxed\{(.*?)\}", completion) for completion in completions]
contents = [match.group(1) if match else "" for match in matches]
# Reward 1 if the content is the same as the ground truth, 0 otherwise
return [1.0 if c == gt else 0.0 for c, gt in zip(contents, ground_truth)]

You can test this function as follows:

>>> prompts = ["Problem: Solve the equation $2x + 3 = 7$. Solution:", "Problem: Solve the equation $3x - 5 = 10$."]
>>> completions = [r" The solution is \boxed{2}.", r" The solution is \boxed{6}."]
>>> ground_truth = ["2", "5"]
>>> reward_func(prompts=prompts, completions=completions, ground_truth=ground_truth)
[1.0, 0.0]

Passing the reward function to the trainer

To use your custom reward function, pass it to the [GRPOTrainer] as follows:

from trl import GRPOTrainer
trainer = GRPOTrainer(
reward_funcs=reward_func,
...,
)

If you have multiple reward functions, you can pass them as a list:

from trl import GRPOTrainer
trainer = GRPOTrainer(
reward_funcs=[reward_func1, reward_func2],
...,
)

and the reward will be computed as the sum of the rewards from each function, or the weighted sum if reward_weights is provided in the config.

Note that [GRPOTrainer] supports multiple reward functions of different types. See the parameters documentation for more details.

GRPOTrainer

[[autodoc]] GRPOTrainer

GRPOConfig

[[autodoc]] GRPOConfig

GitHub: GRPO implementation for code