Mini-R1：用 GRPO 和 Countdown Game 复现 DeepSeek R1 的 aha moment

本文由 Philipp Schmid 撰写，最初发布于 philschmid.de，代码可在此处找到。

Deepseek R1 的发布震惊了整个行业。为什么？因为 DeepSeek-R1 是一个开放模型，在复杂推理任务上可与 OpenAI 的 o1 相媲美，它采用了组相对策略优化（GRPO）和以强化学习为核心的多阶段训练方法。他们不仅发布了模型，还发布了一篇关于如何实现的研究论文。

在论文中，他们描述了在使用纯强化学习训练模型时出现的“顿悟时刻”。在这个阶段，DeepSeek-R1-Zero（DeepSeek-R1 的首次测试）学会了在没有任何人类反馈或描述如何操作的数据的情况下，通过重新评估其初始方法来为问题分配更多的思考时间。他们将此描述为“顿悟时刻”：

这种行为不仅证明了模型不断增长的推理能力，也是一个引人入胜的例子，说明强化学习如何能够带来意想不到的复杂成果。

在这篇博客文章中，我们希望使用组相对策略优化（GRPO）和倒计时游戏来重现 DeepSeek-R1 的小小“顿悟时刻”。我们将使用强化学习训练一个开放模型，尝试教会它自我验证和搜索能力，完全靠自己去解决倒计时游戏。 倒计时游戏是一种数字谜题，玩家使用一组随机抽取的数字和基本算术运算（+、-、×、÷）来达到或尽可能接近目标数字。

Target Number: 952
Available Numbers: 25, 50, 75, 100, 3, 6

(100 × (3 × 3)) + (50 + 6 / 3) = 952

这篇博客文章重点介绍使用 Deepspeed 和 vLLM 进行分布式训练。它在 4 块 NVIDIA H100 GPU 上运行。

1. 搭建开发环境
2. 使用 Deepspeed 和 vLLM 的 GRPO 分布式训练示例
3. 结果与训练观察

注意：这篇博客的灵感来自 Jiayi Pan，他最初探索了这个想法并用一个小模型进行了验证。

但在开始之前，让我们先了解一下组相对策略优化（GRPO），弄清楚它是如何工作的。

组相对策略优化（GRPO）

组相对策略优化（GRPO）是一种强化学习算法，用于提升 LLM 的推理能力。它是在DeepSeekMath 论文中针对数学推理场景提出的。GRPO 对传统的近端策略优化（PPO）进行了改进，去除了对价值函数模型的需求。取而代之的是，它从组分数中估计基线，从而减少内存占用和计算开销。GRPO 现在也被 Qwen 团队使用，既可以配合基于规则/二元的奖励，也可以配合通用奖励模型来提升模型的帮助性。

1. 采样：使用当前策略为每个提示生成多个输出
2. 奖励评分：每次生成都使用奖励函数进行评分，可以是（基于规则或基于结果的）
3. 优势计算：将生成输出的平均奖励用作基线。然后相对于该基线计算组内每个解的优势。奖励在组内进行归一化。
4. 策略优化：策略试图最大化 GRPO 目标，其中包括计算出的优势和 KL 散度项。这与 PPO 在奖励中实现 KL 项的方式不同。

1. 搭建开发环境

我们的第一步是安装 Hugging Face Libraries 和 Pytorch、vllm、trl、transformers 以及 datasets。如果你还没听说过 trl，别担心。它是一个构建在 transformers 和 datasets 之上的新库，让微调、RLHF 和对齐开放 LLM 变得更加容易。

# Install Pytorch & other libraries, make sure to match your GPU driver version
%pip install "torch==2.5.1" tensorboard "setuptools<71.0.0"  --index-url https://download.pytorch.org/whl/cu121

# Install flash-attn
%pip install flash-attn 

# Install Hugging Face libraries
%pip install  --upgrade \
  "transformers==4.48.1" \
  "datasets==3.1.0" \
  "accelerate==1.3.0" \
  "hf-transfer==0.1.9" \
  "deepspeed==0.15.4" \
  "trl==0.14.0"

# install vLLM 
%pip install "vllm==0.7.0"

注意：你可能需要重启内核以使用更新后的包。

我们将使用 Hugging Face Hub 作为远程模型版本管理服务。这意味着我们会在训练过程中自动将模型、日志和信息推送到 Hub。为此你必须在 Hugging Face 上注册。拥有账号后，我们将使用 huggingface_hub 包中的 login 工具登录账号，并将我们的 token（访问密钥）存储在磁盘上。

from huggingface_hub import login

login(token="", add_to_git_credential=True) # ADD YOUR TOKEN HERE

2. 使用 Deepspeed 和 vLLM 的 GRPO 分布式训练示例

我们将使用 Jiayi-Pan/Countdown-Tasks-3to4 数据集，其中包含带有 3 到 4 个数字的样本及其解答。作为模型，我们将使用 Qwen/Qwen2.5-3B-Instruct，这是一个 3B 参数的指令微调模型。这使得展示“顿悟时刻”更加容易，因为它已经能够遵循指令。Jiayi-Pan 探索发现，模型需要具备一定的质量才能学习推理过程，从 > 1.5B 参数开始。

TRL 通过专用的 GRPOTrainer 支持 Group Relative Policy Optimization (GRPO)，用于从偏好数据中对齐 LLM，如 DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models 中所述。GRPOTrainer 是来自 transformers 库的 Trainer 的子类，支持所有相同的功能，包括日志记录、检查点保存、分布式训练和参数高效微调 (PEFT)。

GRPOTrainer 支持通用结果奖励模型 (ORM) 和自定义奖励函数，可用于实现基于规则的奖励模型。在 Deepseek R1 论文中，他们实现了基于规则的奖励模型来验证生成解答的正确性。在我们的示例中，我们将采用类似的方法，创建 2 个奖励函数：

1. 格式奖励：检查生成的格式是否正确 <think> [thinking] </think><answer> [answer] </answer>
2. 准确率奖励：从 <answer> 标签中提取等式，并与目标进行比对评估，且检查每个数字是否只使用一次。

注意：在我们的示例中，正确的 <answer> 包含等式，例如 <answer> 55 + 36 - 7 - 19 </answer>

Hugging Face TRL 添加了对使用 Deepspeed 的分布式训练以及使用 vLLM 进行更快生成的支持。我准备了一个 run_r1_grpo.py 脚本和一个 receipes/grpo-qwen-2.5-3b-deepseek-r1-countdown.yaml 配置文件来运行训练。

此配置已在 4x H100 80GB 的节点上测试和验证，单步大约需要 45-60 秒，因为我们可以利用 vLLM 进行生成，并利用 DeepSpeed 进行分布式训练。因此，我们需要确保将 num_processes 正确设置为 GPU 数量减 1，因为最后一个 GPU 将用于 vLLM 生成。如果你使用更多 GPU，你需要将配置文件中的 vllm_device 改为最后一个 GPU 的索引，例如，如果你有 8 个 GPU，你需要设置 vllm_device=7 并将你的 num_processes 设为 7。

运行训练的命令：

accelerate launch --num_processes 3 --config_file configs/accelerate_configs/deepspeed_zero3.yaml scripts/run_r1_grpo.py --config receipes/grpo-qwen-2.5-3b-deepseek-r1-countdown.yaml

在优化后的分布式训练下，在 4x H100 80GB 上每个样本生成 8 次的单步大约需要 45-60 秒。450 步的完整训练大约需要 6 小时。

3. 结果与训练观察

该脚本将随机补全结果保存到 completion_samples 文件夹中，你可以用它来检查模型的进展。它包括 completion_samples.txt 和 success_completion_samples.txt。completion_samples.txt 包含所有补全结果，而 success_completion_samples.txt 则正确解出了方程。下面你可以找到关于性能随时间变化的有趣训练观察，以及 Tensornoard 日志和成功推理样本。

每 25 步带有检查点的模型可以在 philschmid/qwen-2.5-3b-r1-countdown 找到。

超参数

我使用 DeepSeekMath 论文中的超参数开始了实验，学习率为 1e-6，beta（KL 系数）为 0.04，这导致训练在约 150 步后变得不稳定。我进行了一些小型消融实验，并根据 OpenRLHF 的测试将学习率降低到 5e-7，将 beta 降低到 0.001。我无法测试将 num_generations 从 8 增加到 64 会如何影响训练。64 是生成值，DeepSeekMath 论文中使用了该值。所有其他参数可以在 grpo-qwen-2.5-3b-deepseek-r1-countdown.yaml 配置文件中找到。

训练观察：

• 在约 50 步时，模型已经学会了正确的格式 <think>...</think>\n<answer>...</answer>。
• 在 100 步时，解方程的成功率约为 25%。模型开始用词语进行“推理”，参见下面的示例。
• 在 200 步时，性能似乎收敛得慢得多，我们达到了约 40% 的成功率。模型开始学习一种新的“格式”，它像你以编程方式那样解方程，通过尝试不同的组合并检查结果，参见“第 200 步到第 450 步之间的成功推理样本”。
• 在 450 步时，我们解方程的成功率为 50%。性能仍在缓慢提升，模型保持了从第 200 步开始的新格式形式。

关于模型为何从“词语推理”转向“程序化执行”，我有 3 个可能的假设：

1. Qwen 2.5 3B 不够强或太小，Deepseek 提到你需要一个非常强大的基础模型。
2. 奖励函数定义得不够好，模型通过奖励黑客来解方程。我们可以尝试强制它使用词语，例如加入数字到词频的条件。（我们对 Deepseek 的奖励函数了解不多）
3. 只在 Countdown Game 任务上训练可能会自然地迫使模型学习解方程的最有效方式，因为不需要其他格式。
4. 我们没有训练模型足够长的时间，在 R1 论文中他们展示了训练超过 8000 步的可视化结果。

Tensorboard 结果

截至第 200 步的成功推理样本

<think> First, I need to think about the numbers we have: 14, 70, and 80. To get 85, I think 80 might be a good starting point because it's right in between 85. If I subtract or add 80, it won't lead to 85 easily. Next, 70 seems to be a good intermediate number, especially to use in conjunction with 14. 
Now, I can think about different combinations using addition, subtraction and multiplication/division. If I take 80 - 70 = 10, which is not needed. So, I need to find a way to use 14. If I subtract 14 from 80, I get 66, which still isn't 85. If I subtract 70 from 80, I get 10 again, this isn't constructive yet. 
Multiplication could create numbers larger than the target 85. Let's consider division: If I divide 70 by 14, I get 5. Then if I add this result to 80, I get 85. This combination seems promising: 80 - (70 / 14) = 80 - 5 = 75, this still isn't quite right. Let me correct my approach: Let's use division and then addition. If I take 80 / (70 / 14) → 80 / 5 = 16, and then 80 - 16 = 64, still incorrect. Let me simplify: 80 / (70 / 14) → 80 / 5 = 16, and then 80 - 16 gives 64 again, let me correct: 80 - (70 / 14) → 80 - 5 = 75, but I need 85. Thus I will use the combinations in a slightly different order to insert numbers: I get 80 + (70 / 14) = 80 + 5 = 85</think>
<answer> 80 + (70 / 14) </answer>

<think> First, I need to consider how I can use the numbers 1, 21, and 90 to reach 68. I can try different combinations of operations. Let's try to use subtraction and addition. If I subtract 90 from 21, I get -69, which is too far from 68. Now, if I add 1 to 90, I get 91, which is also not right. Another way is to see what happens if I subtract 21 from 90, which gives me 69. This is closer but still not 68. Let me use multiplication instead of subtraction. 90 * 1 = 90. Now, if I try to reduce this slightly by subtracting 21, I am not getting close to 68.
Now, let's try another multiplication. I'll think of 1 * (larger value) + (smaller value). Let's try 1 * 90 + 21. This evaluates to 90 + 21 = 111, which is too much. Instead, I will try: 90 / 1 + 21 - 7. But this gives a division by 7, which seems complex. Let's re-evaluate: I'll use 1 to divide 21: 21 / 1 = 21. Now, subtract 21 from 90, and add 1, we have 90 - 21 = 69, using one...
Hey, wait on re-evaluating I will subtract 21 from 90, and include 1 to complete the solution: 90 - (21 + 1) which would be 68.</think>
<answer> 90 - (21 + 1) </answer>

<think> First, I need to consider the numbers we have: 59, 80, and 66. I need to use these numbers to get 45. Let's try to use subtraction first since it can quickly reduce large numbers.
Let's see if we can get close by subtracting or adding these numbers. If I subtract 66 from 80, I get 14. Then I need to make this 14 equal 45 when combined with 59. That gives me 14 + 59 = 73, which is too high.
Now let's try adding numbers. If I add 59 and 66, I get 125. Then I would need to subtract 80 from 125 to get 45, but 80 is too large to fit well with this.
Let me try a different approach. What if I subtract 66 from 80 first and then add 59? That would be 14 + 59 = 73, which is still not 45.
Hmm, I need to think of another combination. What if I start with 59 and then add 80 and subtract 66? That would be (59 + 80) - 66, let me calculate it. That equals 139 - 66, which is 73.
That's still not 45, so let me try rearranging it slightly: maybe 80 - 66 + 59 would be closer to 45. Calculating again: 80 - 66 is 14, then adding 59 gives 73.
This way doesn't work either. Let's flip this: maybe (59 + 66) - 80. Let me check: that's 125 - 80, which equals 45.</think>
<answer>(59 + 66) - 80</answer>

<think> First, I need to find a combination of the numbers 61, 4, 94, and 21 that can be manipulated using basic arithmetic operations to equal 50. One way to approach this is to look for numbers that can be used to reach a target number close to 50 and then adjust with the remaining numbers.
Let's try combining 94 and 21 first:
94 + 21 = 115
Now, we need to get from 115 to 50. We can try subtracting 65:
115 - 65 = 50
Now we need to find a way to get 65 using the remaining numbers 61 and 4. We can try:
61 - 4 = 57
Now we need to get from 57 to 65. Adding 8 would give us that:
57 + 8 = 65
So, we need to find a way to get 8 using 61 and 4. We can try:
61 - 94 = -36
This is negative, so let's try dividing:
61 / 94 = 0.648987...
This is not helping, so let's try multiplication and division in a different way:
(61 - 4) * (94 / 21) = 57 * 4.476... = 254.4
This is not working either. Let's try a different approach:
(94 - 61) = 33
Now we have:
33 + 21 = 54
We need to adjust from 54 to 50, so we can:
54 - 4 = 50
So, the equation becomes:
(94 - 61) + (21 - 4) = 50
</think>
<answer>(94 - 61) + (21 - 4)</answer>

<think> First, I'll consider the possibility of directly subtracting two numbers to get close to 31. If I subtract 57 from 63, I get 7. Now, if I subtract this 7 from 25, I get 25 - 7 = 18. Far from 31, so this doesn't work.
Next, let's try adding and subtracting. Dividing 63 by 25 will give a decimal, so that's unlikely to work. Now, let's see if adding and subtracting in another order could work. For example, let's try 63 - 25, and then we see this gets us 38. Now, we need a way to get to 31 from here. Subtracting 7 from 38 would work, but we can't figure out how to create that 7 from the given numbers.
Hmm, let me try to incorporate all numbers by adding and subtracting. If I try 63 + 25, I get 88. This is quite far from 31, but if 88 - 57, I get 31 exactly!
This seems to be the equation: 63 + 25 - 57 equals 31.
</think>
<answer> 63 + 25 - 57 </answer>

第 200 步到第 450 步之间的成功推理样本

<think> We need to use the numbers 39, 16, and 55 with basic arithmetic operations to get 78. Let's try different combinations:
- 55 + 39 - 16 = 84 - 16 = 68 (too low)
- 55 + 16 - 39 = 71 - 39 = 32 (too low)
- 55 - 16 + 39 = 39 + 39 = 78 (this works)
So, the equation is 55 - 16 + 39 = 78. </think>
<answer> 55 - 16 + 39 </answer>

<think> We need to use the numbers 21, 5, and 40 with basic arithmetic operations to get 24. Let's try different combinations:
- 21 - 5 - 40 = -15
- 21 - 5 + 40 = 56
- 21 + 5 - 40 = -14
- 21 * 5 / 40 = 2.1 (not an integer, let's try another)
- 40 * 5 / 21 = 9.52 (not 24)
- 40 / 5 - 21 = 9 - 21 = -12
- 5 * 40 / 21 = 100 / 21 ~= 4.76 (not 24)
- 5 + 40 - 21 = 24 ! This works.
</think>
<answer> 5 + 40 - 21 </answer>

<think> We need to use the numbers 75, 99, 62, and 34 with basic arithmetic operations to get 72. Let's try different combinations:
- 99 - 75 + 62 - 34 = 92 - 34 = 58 (too high)
- 99 - 75 - 62 + 34 = 24 - 62 + 34 = -38 + 34 = -4 (too low)
- 99 + 75 - 62 - 34 = 174 - 96 = 78 (too high)
- 99 + 75 - 62 + 34 = 174 + 34 - 62 = 208 - 62 = 146 (too high)
- 99 - 75 + 62 + 34 = 24 + 96 = 120 (too high)
- 75 + 99 - 62 - 34 = 174 - 96 = 78 (too high)
- 75 + 99 - 62 + 34 = 174 + 34 - 62 = 208 - 62 = 146 (too high)
- 75 + 62 - 99 + 34 = 137 - 99 + 34 = 38 + 34 = 72
So, 75 + 62 - 99 + 34 equals 72.
</think>
<answer> 75 + 62 - 99 + 34 </answer>

结论

DeepSeek R1 及其研究论文的发布可能是开放科学和开源开发的一个突破点。在 DeepSeek 发布仅一周后，我们就能够使用 GRPO 和 Countdown Game 复现 R1 所学“推理”的简单版本。虽然我们的实现专注于特定任务而非通用推理，并收敛到一种非常特定的“推理”格式，但它表明该方法正在发挥作用。

展望 2025 年，很明显我们正处于更重大进展的边缘。RL 将变得更加易用和用户友好，更多研究者和开发者将探索其潜力，但与监督微调相比，也将像以前一样需要更多的计算资源。

我对 2025 年感到兴奋。如果你有任何问题或想法，欢迎联系我。

如果这听起来有趣，我们非常希望得到你的帮助！无论是贡献代码，还是加入 Hugging Face 上的讨论。