基于trl复现DeepSeek-R1的GRPO训练过程

1. 引入

huggingface开发了强化学习训练Transformer的库trl（参考3），借助这个trl，可以用来做GRPO的强化学习训练。魔搭ModelScope社区的文章（参考2）给出了基于Qwen基座模型Qwen2.5-0.5B-Instruct，复现DeepSeek-R1训练过程的描述，参考1还给出了源码。
本文就对这个代码进行复现，并记录重点步骤与过程，以期更好的理解GRPO训练过程、数据格式、训练前后的不同效果。

2. 关键步骤

本文仅讲解关键步骤，完整步骤/代码可看参考2与参考1。

2.1 数据处理

（1）原始数据
数据集选用的是OpenAI的gsm8k（参考4），数据是多组question和answer组成的，原始数据格式如下所示：
在这里插入图片描述
可以将answer看过有两部分组成：推理过程，标准答案（数字）。两部分用####隔开。

（2）处理方法：加上prompt

# 系统提示词：规定回复格式
SYSTEM_PROMPT = """
Respond in the following format:
<reasoning>
...
</reasoning>
<answer>
...
</answer>
"""
# 最终每条数据格式
d = { # type: ignore
	'prompt': [
		{'role': 'system', 'content': SYSTEM_PROMPT},
		{'role': 'user', 'content': x['question']}
	],
	'answer': extract_answer(x['answer'])# 标准答案（数字）
}

最终处理后数据格式入上面代码中的d所示，处理完后，一共7473组数据。如下是1条处理完成后的数据示例：

{
    "prompt": [
        {
            "role": "system",
            "content": "\nRespond in the following format:\n<reasoning>\n...\n</reasoning>\n<answer>\n...\n</answer>\n"
        },
        {
            "role": "user",
            "content": "Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?"
        }
    ],
    "answer": "72"
}

可从prompt中看到，最终模型需要回复的内容中有两个字段，一个是推理过程（<reasoning>字段），另一个是最终答案（<answer>字段）。

2.2 自定义Reward函数

强化学习中比较重要的是reward函数，这里定义了5个reward函数：

（1）correctness_reward_func
如果模型输出答案中的<answer>字段（数字）等于标准答案，则得2分，否则0分。

（2）int_reward_func
如果模型输出答案中的<answer>字段是纯数字类型，则得0.5分，否则得0分。

（3）strict_format_reward_func
如果模型输出答案中的，既有<reasoning>字段也有<answer>字段，则得0.5分，否则0分。
strict表示这里做正则匹配要求是更严格的，必须以<reasoning>开头，以<answer>结束。回复内容收尾必须是这两个字段。具体正则为：pattern = r"^<reasoning>\n.*?\n</reasoning>\n<answer>\n.*?\n</answer>\n$"。

（4）soft_format_reward_func
soft表示这里的格式要求匹配正则不是那么严格。具体正则为pattern = r"<reasoning>.*?</reasoning>\s*<answer>.*?</answer>"。

（5）xmlcount_reward_func
根据如下逻辑对<reasoning>字段和<answer>字段做计数后打分，这也是一种对输出格式的打分。

def count_xml(text) -> float:
    count = 0.0
    if text.count("<reasoning>\n") == 1: # 只出现1次<reasoning>
        count += 0.125
    if text.count("\n</reasoning>\n") == 1: # 只出现1次</reasoning>
        count += 0.125
    if text.count("\n<answer>\n") == 1: # 只出现1次<answer>
        count += 0.125
        count -= len(text.split("\n</answer>\n")[-1])*0.001
    if text.count("\n</answer>") == 1: # 只出现1次</answer>
        count += 0.125
        count -= (len(text.split("\n</answer>")[-1]) - 1)*0.001
    return count

def xmlcount_reward_func(completions, **kwargs) -> list[float]:
    contents = [completion[0]["content"] for completion in completions]
    return [count_xml(c) for c in contents]

2.3 GRPO训练

训练时，定义好超参数，按如下配置即可训练

trainer = GRPOTrainer(
    model=model, # 模型基座
    processing_class=tokenizer, # 分词器
    reward_funcs=[ # reward函数
        xmlcount_reward_func,
        soft_format_reward_func,
        strict_format_reward_func,
        int_reward_func,
        correctness_reward_func],
    args=training_args, # 超参数
    train_dataset=dataset, # 数据
)
trainer.train()

如下为训练中间过程的部分输出（能看出用1张A800训练一次耗时1:41:31）：

-------------------- Question:
Billy wants to watch something fun on YouTube but doesn't know what to watch.  He has the website generate 15 suggestions but, after watching each in one, he doesn't like any of them.  Billy's very picky so he does this a total of 5 times before he finally finds a video he thinks is worth watching.  He then picks the 5th show suggested on the final suggestion list.  What number of videos does Billy watch?
Answer:
65
Response:
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
Extracted:
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
-------------------- Question:
Xander read 20% of his 500-page book in one hour.  The next night he read another 20% of the book.  On the third night, he read 30% of his book.  How many pages does he have left to read?
Answer:
150
Response:
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
Extracted:
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
{'loss': 0.0, 'grad_norm': nan, 'learning_rate': 4.918502993675464e-06, 'rewards/xmlcount_reward_func': 0.0, 'rewards/soft_format_reward_func': 0.0, 'rewards/strict_format_reward_func': 0.0, 'rewards/int_reward_func': 0.0, 'rewards/correctness_reward_func': 0.0, 'reward': 0.0, 'reward_std': 0.0, 'completion_length': 200.0, 'kl': nan, 'epoch': 0.17}
 17%|███████████████████                                                                                           | 324/1868 [17:38<1:22:29,  3.21s/it]-------------------- Question:
Jordan gave Danielle two dozen roses and a box of chocolates as a birthday day gift.  Later that day, after Jordan left, Danielle traded the box of chocolates for another dozen roses.  Overnight, half of the roses wilted, and  Danielle decided to throw the wilted flowers away.  On the second day, another half of the remaining flowers wilted, and she threw the wilted ones away.  How many unwilted flowers remained?
Answer:
9
Response:
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
Extracted:
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
...
...
-------------------- Question:
It's payday but Jebb has to pay 10% for the tax. If his pay is $650, how much is his take-home pay?
Answer:
585
Response:
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
Extracted:
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
-------------------- Question:
In a car racing competition, Skye drove a 6-kilometer track. For the first 3 kilometers, his speed was 150 kilometers per hour. For the next 2 kilometers, his speed was 50 kilometers per hour more. For the remaining 1 kilometer, his speed was twice as fast as his speed on the first 3 kilometers. What is Skye's average speed for the entire race?
Answer:
180
Response:
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
Extracted:
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
{'loss': 0.0, 'grad_norm': nan, 'learning_rate': 0.0, 'rewards/xmlcount_reward_func': 0.0, 'rewards/soft_format_reward_func': 0.0, 'rewards/strict_format_reward_func': 0.0, 'rewards/int_reward_func': 0.0, 'rewards/correctness_reward_func': 0.0, 'reward': 0.0, 'reward_std': 0.0, 'completion_length': 200.0, 'kl': nan, 'epoch': 1.0}
{'train_runtime': 6091.2875, 'train_samples_per_second': 1.227, 'train_steps_per_second': 0.307, 'train_loss': 0.0, 'epoch': 1.0}
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1868/1868 [1:41:31<00:00,  3.26s/it]

2.4 基本验证

对GRPO训练前后的大模型输入一个问题：

prompt = "Xiao Ming bought 4 apples, ate 1, and gave 1 to his sister. How many apples were left?"

（1）GRPO训练前模型的输出：

To determine how many kilometers Natalia rode in total, we need to calculate her distances for each day and then sum them up.

1. **Monday**: Natalia rode 40 kilometers.
2. **Tuesday**: Natalia rode 50 kilometers.
3. **Wednesday**: She rode 50% fewer kilometers than on Tuesday. This means she rode \(50 - (0.5 \times 50) = 50 - 25 = 25\) kilometers.
4. **Thursday**: She rode as many kilometers as the sum of the kilometers from Monday and Wednesday. So, she rode \(40 + 25 = 65\) kilometers.

Now, let's add up the distances from all four days:

- Monday: 40 kilometers
- Tuesday: 50 kilometers
- Wednesday: 25 kilometers
- Thursday: 65 kilometers

Total distance = 40 + 50 + 25 + 65 = 180 kilometers

Therefore, the total number of kilometers Natalia rode is \(\boxed{180}\).

（2）GRPO训练后模型的输出：

<reasoning> 
Initially, Xiao Ming had 4 apples. After eating 1 apple, he was left with 4 - 1 = 3 apples. Then, he gave 1 apple to his sister, leaving him with 3 - 1 = 2 apples. 
</reasoning>
<answer> 
2 
</answer>

从输出内容，输出格式，最终结果可以看到GRPO训练前后模型的不同表现。