AB
AiBoss
Tutorials

DeepSeek

Tutorials

用 GRPO 复现 DeepSeek-R1 的「顿悟时刻」:Countdown 游戏强化学习实战

DeepSeek-R1 论文里描述过一个「顿悟时刻」:模型在纯强化学习下自发学会重新审视自己的解法、分配更多思考时间。本教程用 GRPO 算法和 Countdown 数字游戏,在一个 3B 参数的开源模型上复现这一现象,涵盖环境搭建、奖励函数设计、DeepSpeed + vLLM 分布式训练配置,以及分阶段的训练观察。

DeepSeek-R1 的发布之所以引起关注,不只是因为它是一个在复杂推理任务上表现接近顶尖闭源模型的开源模型,更因为它公开了训练方法:以 Group Relative Policy Optimization(GRPO)为核心、围绕强化学习设计的多阶段训练流程。论文里提到一个被称为「顿悟时刻」(aha moment)的现象——在纯强化学习阶段,模型在没有任何人类反馈、也没有任何示范数据的情况下,自己学会了重新评估最初的解题思路,从而把更多思考时间分配给真正困难的问题。

本教程要做的,就是在小规模上复现这个「顿悟时刻」。我们使用 DeepSeek 论文中提出的 GRPO 算法,配合 Countdown 数字游戏作为任务,训练一个开源模型,观察它能否自己长出自我验证与搜索的能力。整个流程基于 TRL 的 GRPOTrainer,用 DeepSpeed 做分布式训练、用 vLLM 加速生成,参考配置在 4 张 NVIDIA H100 80GB 上验证过。

适合阅读的人群:已经了解 Transformer 微调基本流程、想动手跑一次强化学习对齐实验的工程师或研究者。如果你只是想调用现成的推理模型,这篇内容偏底层,可能不是最省力的路径。

准备工作

任务与算法背景

Countdown 游戏是一个数字谜题:给出一组随机抽取的数字和一个目标数字,玩家用加、减、乘、除四种基本运算,把给定数字各用一次,凑出目标数字(或尽量接近)。例如目标数字是 952,可用数字是 25、50、75、100、3、6,一个可行解是 (100 × (3 × 3)) + (50 + 6 / 3) = 952。这个任务的好处是:答案对错可以自动判定,不需要人工标注,非常适合用规则型奖励来做强化学习。

GRPO 是一种用于提升大模型推理能力的强化学习算法,最早在数学推理场景中被提出。它相对传统 PPO 的关键改动是去掉了价值函数模型:不再单独训练一个 critic 来估计基线,而是用同一组采样结果的得分均值作为基线。这样做显著降低了显存占用和计算开销。它的工作流程大致是:

  • 采样:对每个 prompt,用当前策略生成多条输出。
  • 打分:用奖励函数给每条输出打分,奖励可以是规则型、结果型,也可以是通用奖励模型。
  • 优势计算:把组内输出的平均奖励当作基线,每条解法的优势值相对这个基线计算,奖励在组内做归一化。
  • 策略优化:最大化 GRPO 目标函数,其中包含优势项和一个 KL 散度项。与 PPO 不同的是,KL 项的处理方式不一样,PPO 是把 KL 塞进奖励里。

硬件与账号

训练脚本在 4 张 H100 80GB 的节点上做过验证。显存需求与模型规模、序列长度、每样本生成条数直接相关,换更小的卡需要相应下调这些参数。此外需要一个 Hugging Face 账号:训练过程中模型、日志和相关信息会自动推送到 Hub 上做版本管理,因此要先注册并准备好访问令牌。

安装依赖

第一步是装 PyTorch、TensorBoard 以及 Hugging Face 生态的几个库。注意 PyTorch 的安装源要匹配你自己的 GPU 驱动版本,下面给出的版本号只是参考,实际以各库官方当前说明为准。

# 安装 PyTorch 及其他基础库,注意匹配你的 GPU 驱动版本
%pip install "torch==2.5.1" tensorboard "setuptools<71.0.0" --index-url

# 安装 flash-attn
%pip install flash-attn

# 安装 Hugging Face 相关库
%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"

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

安装完成后可能需要重启内核才能让新装的包生效。这里的 trl 是构建在 transformers 和 datasets 之上的库,专门用来简化开源大模型的微调、RLHF 和对齐流程,本教程用到的 GRPOTrainer 就来自它。

接着登录 Hugging Face,把访问令牌写到本地凭据里:

from huggingface_hub import login

login(token="<你的令牌>", add_to_git_credential=True)

操作步骤

第一步:选定数据集与基座模型

数据集使用 Jiayi-Pan/Countdown-Tasks-3to4,它包含使用 3 到 4 个数字的题目及其解答。模型选择 Qwen/Qwen2.5-3B-Instruct,一个 30 亿参数的指令微调模型。选指令微调版本的原因是它已经能听懂指令,更容易把「顿悟」过程展示出来。

这里有一个经验性的门槛:模型本身需要具备一定的基础能力,才能学会推理过程。相关探索表明,参数量低于 15 亿的模型很难跑出效果,所以不要用太小的模型做这个实验。

第二步:设计奖励函数

GRPOTrainer 支持通用的结果奖励模型(ORM)以及自定义奖励函数,后者可以用来实现规则型奖励。DeepSeek-R1 论文中就是用规则型奖励来验证生成解的正确性,本教程采用类似思路,定义两个奖励函数:

  • 格式奖励:检查生成内容是否符合 <think> [思考过程] </think><answer> [答案] </answer> 的结构。
  • 准确率奖励:从 <answer> 标签中提取算式,代入目标数字求值,同时检查每个给定数字是否恰好被使用了一次。

注意答案里要包含完整算式,例如 <answer> 55 + 36 - 7 - 19 </answer>,而不是只写一个结果数字,否则准确率奖励无法校验数字使用情况。

第三步:配置分布式训练

训练入口是一个 run_r1_grpo.py 脚本,配合 receipes/grpo-qwen-2.5-3b-deepseek-r1-countdown.yaml 配置文件。这套配置在 4 张 H100 80GB 上测试通过,单步耗时约 45 到 60 秒,因为生成阶段交给 vLLM、分布式训练交给 DeepSpeed,两者各司其职。

关键点在于进程数与 GPU 的分配关系:最后一张 GPU 要留给 vLLM 做生成,所以 num_processes 必须设成「GPU 总数减一」。4 卡机器就是 3 个训练进程。如果卡更多,还要同步修改配置文件里的 vllm_device,把它指向最后一张卡的索引。例如 8 卡机器,vllm_device=7num_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

按上述配置,每个样本生成 8 条输出,单步 45 到 60 秒,跑满 450 步大约需要 6 小时。

第四步:观察训练过程

脚本会把随机采样的生成结果保存到 completion_samples 目录,用来检查模型进展。其中有两个文件:

  • completion_samples.txt:包含全部生成结果。
  • success_completion_samples.txt:只保留正确解出算式的样本。

训练过程中每 25 步保存一次检查点,方便回看不同阶段的模型状态。TensorBoard 日志可以同时观察奖励曲线和成功率的变化。

一个完整示例

下面把整个流程串一遍,从零到能跑起来。

1. 环境准备。按前面的命令安装 PyTorch、flash-attn、transformers、datasets、accelerate、hf-transfer、deepspeed、trl 和 vllm。装完重启内核。

2. 登录 Hub。huggingface_hub.login() 写入令牌,训练过程中的模型和日志会自动推送上去。

3. 准备数据与模型。数据集用 Jiayi-Pan/Countdown-Tasks-3to4,基座模型用 Qwen/Qwen2.5-3B-Instruct

4. 写奖励函数。格式奖励校验 <think><answer> 标签结构;准确率奖励解析 <answer> 中的算式,验证求值结果等于目标数字,且每个给定数字恰好使用一次。

5. 准备配置文件。在 YAML 里设置学习率、KL 系数 beta、每样本生成条数等超参数,并确认 vllm_device 指向最后一张 GPU。

6. 启动训练。accelerate launch 拉起脚本,num_processes 等于 GPU 数减一。4 卡机器上单步约 45 到 60 秒,450 步约 6 小时。

7. 检查结果。训练结束后查看 completion_samples 目录下的两个文本文件,以及 TensorBoard 中的成功率曲线,判断模型是否学会了正确格式、是否开始出现自我验证式的推理。

训练观察与超参数调整

超参数

实验最初沿用 DeepSeekMath 论文的超参数:学习率 1e-6,KL 系数 beta 为 0.04。结果训练在 150 步左右开始变得不稳定。经过几组小规模消融,把学习率降到 5e-7、beta 降到 0.001 之后训练才稳定下来。这个调整参考了 OpenRLHF 的相关测试。

另一个没能验证的变量是每样本生成条数:实验用的是 8,而 DeepSeekMath 论文用的是 64。把 8 提到 64 会带来什么影响,这里没有测试数据,不做推测。其余参数都在配置文件中。

分阶段表现

  • 约 50 步:模型学会了正确的输出格式 <think>...</think>\n<answer>...</answer>
  • 约 100 步:解出算式的成功率约 25%,模型开始用自然语言「推理」,出现类似「先看看手上有哪些数字」「这个组合似乎有希望」这样的表述。
  • 约 200 步:性能提升明显放缓,成功率约 40%。同时模型开始形成一种新的「格式」——它不再用文字描述思路,而是像程序一样枚举不同组合、逐个检查结果。
  • 约 450 步:成功率达到 50%,仍在缓慢提升,且保持了 200 步时形成的程序化风格。

为什么从「文字推理」转向「程序化执行」

对于这个风格转变,有三种可能的解释,都还没有定论:

  • Qwen2.5-3B 本身能力不够强或规模偏小。DeepSeek 方面提到过,这类训练需要非常强的基座模型。
  • 奖励函数定义得不够好,模型找到了某种「奖励 hack」的捷径来解算式。
  • 只训练 Countdown 这一种任务,模型自然会收敛到最高效的解法,因为没有其他格式的要求。

还有一种可能是训练步数不够——R1 论文里展示的训练曲线超过了 8000 步。如果想让模型维持文字推理风格,可以尝试加入约束,比如对数字与文字的出现频率设置条件。

注意事项

  • 进程数必须留一张卡给 vLLM。这是最容易踩的坑:num_processes 设成 GPU 总数减一,同时把 vllm_device 指向最后一张卡的索引,两者要一致。
  • 学习率和 beta 不要照搬论文默认值。1e-6 配 0.04 的组合在本实验里 150 步左右就发散了,需要下调。
  • 模型规模有下限。参数量低于 15 亿的模型很难学会推理过程,不要指望用小模型复现。
  • 答案必须包含完整算式。只写结果数字的话,准确率奖励无法校验数字是否各用一次。
  • 训练时长与硬件强相关。450 步约 6 小时是 4 张 H100 80GB 上的数字,换硬件后单步耗时和总时长都会变。
  • 版本号、显存需求、配额等信息会随时间变化,动手前请以各依赖库和平台的官网当前信息为准。