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

选自GitHub

作者： Andriy Burkov

机器之心编译

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_iterations=1：从当前策略模型创建新参考模型的外部迭代次数。一次迭代是指在整个数据集上执行一次通过。
• num_steps=500：训练循环将执行最多 500 个步骤，每个步骤处理一批样本。
• batch_size=7：在 8 台 GPU 的情况下，每个步骤每批处理 7 个样本，每台 GPU 上放置 1 个样本。使用一个 GPU (0) 被 DataParallel 用作主节点来聚合梯度并收集输出。
• num_generations=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 数据执行的策略更新次数。在这里，我们每批只执行一次更新。
• epsilon=0.1：GRPO 的 PPO 组件的 clipping 参数。这可以防止策略在单次更新中发生太大的变化。

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