DeepSeek关键RL算法GRPO，有人从头跑通了，贡献完整代码

LH3ErcffX · 7 天前

GRPO（Group Relative Policy Optimization）是 DeepSeek-R1 成功的基础技术之一，我们之前也多次报道过该技术，比如《DeepSeek 用的 GRPO 占用大量内存？有人给出了些破解方法》。
简单来说，GRPO 算法丢弃了 critic model，放弃了价值函数近似，转而通过组内样本的相对比较来计算策略梯度，从而有效降低了训练的不稳定性，同时提高了学习效率。
既然 GRPO 如此有效，那么，你知道如何从头开始实现 GRPO 吗？
近日，AI 工程师和技术作家 Andriy Burkov 发布了一份「从头开始写 GRPO 代码」的教程，其中介绍了如何基于 Qwen2.5-1.5B-Instruct 模型构建一个使用 GRPO 的分布式强化学习流程。
不过，在我们深入这份教程之前，先简单介绍一下它的作者。Andriy Burkov 算得上是 AI 领域的一位著名科普作家，在加拿大拉瓦尔大学取得了计算机科学博士学位，还曾发表过两本颇受欢迎的 AI 主题著作：《100 页语言模型书》和《100 页机器学习书》；书中一步步详实地介绍了相关概念，并附带了简明的实现代码。

接下来我们就来看看这份 GRPO 从头实现教程吧。

教程地址：https://github.com/aburkov/theLMbook/blob/main/GRPO_From_Scratch_Multi_GPU_DataParallel_Qwen_2_5_1_5B_Instruct.ipynb
从头编写 GRPO 代码

使用 Qwen2.5-1.5B-Instruct 的分布式实现

本教程将展示如何使用 GRPO 方法构建分布式强化学习（RL）流程，从而可以针对数学、逻辑和编程任务对语言模型进行微调。
首先需要明确，这些任务都存在一个唯一且正确的 ground truth 答案，可通过简单的字符串比较轻松加以验证。
GRPO 的发明者是 DeepSeek，最早是被用于微调 DeepSeek 的 R1 和 R1-Zero 模型 —— 它们可通过学习生成思维链（CoT）来更好地解决数学和逻辑问题。
本教程的目标是将通用语言模型 Qwen2.5-1.5B-Instruct 转换为数学问题求解器。我们将从头开始编写 GRPO 代码，然后将其与几个流行的库和工具集成起来，以实现分布式训练管道流程，包括：

PyTorch：用于张量运算和分布式训练。
Hugging Face Transformers：用于加载预训练的语言模型和 tokenizer。
FlashAttention2：优化的注意力机制，有助于减少内存使用量并提高训练速度。
Weights & Biases (wandb)：用于实验跟踪、可视化和模型版本控制。

本教程分为几个部分。首先是基本设置和导入，然后是数据格式化和答案提取、数据集准备、评估函数、奖励函数、训练设置和执行，最后加载和测试模型。此过程中，我们将从头实现 GRPO 算法。
Part 1：基础设置和导入

首先是安装并导入所有必要的模块。下面是导入库的一段代码截图。

部分代码截图。完整代码块参见 GitHub。
运行上述代码（参考项目完整代码），可以执行以下任务：

设置随机种子：set_random_seed 函数通过为 Python 的随机模块、NumPy 和 PyTorch 设置种子，确保可复现性；
环境变量配置：设置 WANDB_API_KEY 和 WANDB_PROJECT 环境变量，以启用与 Weights & Biases 的实验跟踪；
导入必要的库，包括 random、copy、re、torch 等等。

Part 2：数据格式以及答案提取

接下来，项目定义了数据格式，以及模型如何从输出和数据集中提取答案段落。
为了确保模型输出格式一致，项目还定义了一个系统提示。该提示指示模型生成包含 < reasoning > 和 < answer > 标签的输出。这一步通过两个函数完成：

extract_answer_from_model_output：此函数获取模型的输出文本，并提取 < answer > 标签内的内容；
extract_answer_from_dataset：此函数从 GSM8K 数据集中提取预期答案，该数据集使用 “####” 分隔符来分隔答案：

部分代码截图。完整代码块参见 GitHub。
Part 3：数据准备

该项目使用 GSM8K 数据集进行训练。项目使用了该数据集中的示例来训练模型，基于强化学习（RL）训练范式，让模型生成多个问题解答样本，之后作者将这些解答与 GSM8K 示例中的标准答案进行对比，如果匹配，就为 RL 算法（GRPO）提供高奖励，然后更新模型权重，以增加模型下次获得高奖励的可能性。
实验过程是这样的。首先从 Hugging Face 加载数据集，然后格式化每个示例，包括系统提示和用户提示。这段实现代码中还定义了两个辅助函数：prepare_dataset 以及 build_prompt。

部分代码截图。完整代码块参见 GitHub。
Part 4：评估函数

评估对于跟踪模型的进展至关重要。因此作者定义了一些函数，从而可以在一组示例上对模型进行评估。该项目的评估函数执行以下任务：

token 化提示并生成响应：模型的输出是在 token 化提示的基础上生成的。
提取预测答案：从生成的响应中提取答案。
将预测答案与预期答案进行比较：这种比较是通过精确匹配以及数值等价检查来完成的。

在这段代码中，两个辅助函数 _extract_last_number 和 _extract_single_number 被用来从文本中提取数字。评估函数 evaluate_model 使用这些辅助函数来确定预测答案是否正确：

部分代码截图。完整代码块参见 GitHub。
Part 5：奖励函数

在强化学习中，奖励函数是必不可缺的，作者定义了两个奖励函数：
correctness_reward：这个函数根据生成的答案是否正确来分配奖励。采用两种方式：精确的字符串匹配和数值等价检查，将模型输出的答案与预期答案进行比较。完全匹配会获得更高的奖励（2.0），而基于数值等价的匹配会获得较小的奖励（1.5）。
format_reward：这个函数鼓励模型遵循所需的类似 XML 的输出格式。它为生成文本中存在 < reasoning>、、和标签提供小额奖励。

部分代码截图。完整代码块参见 GitHub。
Part 6：从头开始实现 DataParallel GRPO

这一节，我们将从头实现 GRPO 算法的所有构建模块。首先，这里假设运行代码的机器至少有 2 台 GPU。为此，这里要使用 PyTorch 的 DataParallel API 来将策略模型放在多个 GPU 核心上，每个 GPU 核心都有该模型的一个副本。然后将批量数据分散在这些 GPU 核心上完成处理。

部分代码截图。完整代码块参见 GitHub。
Part 7：训练设置和执行

这一节，我们将所有组件组合在一起，完成设置并开始训练。
首先，加载预训练的模型和 tokenizer，准备评估数据，然后使用上面从头实现的 train_with_grpo 进行强化学习微调。
关键步骤包括：

模型和 tokenizer 初始化：使用优化设置（使用 torch.bfloat16 和 FlashAttention2）加载模型 Qwen/Qwen2.5-1.5B-Instruct。tokenizer 也要加载，其填充 token 设置为序列末尾 token。使用 torch.bfloat16 加载模型会将其参数转换为每个数值使用 16 位而不是 32 位的形式，这可将模型的内存使用量减少一半，并且可加快在现代 GPU 上的训练速度。
初步评估：在微调之前，根据几个示例对模型进行评估，以确定基准性能。
强化学习微调：为从头开始实现 GRPO 的训练函数 train_with_grpo 配置适当的训练参数和奖励函数。然后，在剩余的训练数据上执行强化学习训练。
最终评估和模型保存：强化学习微调后，再次评估模型，并保存最终模型。

下面的代码会执行以下功能：

确定设备（如果有 GPU 就用 GPU，否则就用 CPU）。
加载预训练版 Qwen2.5-1.5B-Instruct 模型和 tokenizer。tokenizer 的 pad token 设置为 eos_token。
保留一小部分数据集用于评估，以提供基线。
通过启用梯度检查点和禁用 KV 缓存，优化模型的内存效率。
步骤 1：在微调之前评估模型，以建立基线准确性。
步骤 2：使用 train_with_grpo 函数和我们定义的奖励函数（format_reward 和 correctness_reward，合并为 combined_reward）执行强化学习微调。这里使用了多台 GPU 训练模型。
步骤 3：将最终的微调模型和 tokenizer 保存到磁盘。

GRPO 训练流程使用的超参数如下。
训练配置
以下参数设定了使用上面的 GRPO 算法实现强化学习微调运行的配置：

num_iteratinotallow=1：从当前策略模型创建新参考模型的外部迭代次数。一次迭代是指在整个数据集上执行一次通过。
num_steps=500：训练循环将执行最多 500 个步骤，每个步骤处理一批样本。
batch_size=7：在 8 台 GPU 的情况下，每个步骤每批处理 7 个样本，每台 GPU 上放置 1 个样本。使用一个 GPU (0) 被 DataParallel 用作主节点来聚合梯度并收集输出。
num_generatinotallow=14：对于训练数据中的每个提示词，训练器将生成 14 个不同的完成结果。这些生成结果将被用于计算指导强化学习更新的相对优势（或奖励信号）。如果你的 GPU 的 VRAM 较少，请减少此数字。
max_completion_length=400：在生成完成结果（序列的 response 部分）时，生成上限为 400 个 token。这限制了模型在 RL 阶段生成的输出的长度。如果你的 GPU 的 VRAM 较少，请减少此数字。
beta=0.04：GRPO 损失函数中 KL 散度惩罚的系数。这控制的是模型与参考模型的偏差程度。
learning_rate=5e-6：RL 微调的学习率。为了实现稳定的策略更新，这里使用了相对较低的学习率。
mu=1：对每批 rollout 数据执行的策略更新次数。在这里，我们每批只执行一次更新。
epsilnotallow=0.1：GRPO 的 PPO 组件的 clipping 参数。这可以防止策略在单次更新中发生太大的变化。

在微调之前和之后都会对模型进行评估，以衡量准确率的提高情况。最后，将微调后的模型保存到 grpo_finetuned_model 目录中。

部分代码截图。完整代码块参见 GitHub。
教程中还给出了详细的执行情况，可作参考。

下面我们也简单看看其训练过程。
首先，初始配置后，我们可以看到运行 GRPO 之前的准确度为 23.33%。

然后经过 500 步的 1 轮 GRPO 迭代，下图展示了相关的训练动态：

训练完成后，自然还需要对模型进行新一轮的评估。这里采用了 30 个评估样本来进行评估，以下展示了其中一个模型回答正确的示例：

整体表现如何呢？可以看到，经过一轮 GRPO 之后，Qwen-2.5-1.5B-Instruct 模型答对了 30 问题中的 27 题，实现了 90% 的准确度。相较于 GRPO 之前的 23.33%，可说是实现了性能飞跃。

上面两张图展示了模型的学习过程动态，可以看到：平均奖励在 2.25 左右就趋于稳定了（理论最大值为 0.8 + 2.0 = 2.8）。相比于另一处微调的 Qwen-2.5-0.5B-Instruct（获得的平均奖励为 1.4），这个数字相当高了，参阅：https://github.com/aburkov/theLMbook/blob/main/GRPO_Qwen_0_5_Instruct.ipynb
如果使用更大的模型并允许更长的生成时间，模型正确解答问题的能力还将进一步提升。但是，如果要训练更大的模型，不仅需要将数据分布在多台 GPU 上，还需要将模型分开放在多台 GPU 上，这需要用到 DeepSpeed 或 FSDP（完全分片数据并行）等模型并行工具。
下面加载和测试已经微调的模型：

完整代码见原笔记本
加载完成后测试一下，首先问问 1+1 等于几：

可以看到，模型反复思考了很多次，终于认定确实等于 2。
多次测试后还可以发现，该模型没有学会生成序列结束（EOS）token，因此即使在 token 之后，输出序列仍会继续。这是预期的行为，因为我们使用的奖励函数中没有包含一个用于停止生成的奖励。我们也没有执行监督微调步骤 —— 该步骤可以让模型学会在之后立即生成 EOS。
你对这篇代码密集的教程怎么看？有没有让你产生在自己的电脑上实现 GRPO 的想法？