机器之心编译
GRPO（GroupRelativePolicyOptimization）是DeepSeek-R1成功的基础技术之一，我们之前也多次报道过该技术，比如《DeepSeek用的GRPO占用大量内存？有人给出了些破解方法》。
简单来说，GRPO算法丢弃了criticmodel，放弃了价值函数近似，转而通过组内样本的相对比较来计算策略梯度，从而有效降低了训练的不稳定性，同时提高了学习效率。
既然GRPO如此有效，那么，你知道如何从头开始实现GRPO吗？
近日，AI工程师和技术作家AndriyBurkov发布了一份「从头开始写GRPO代码」的教程，其中介绍了如何基于Qwen2.5-1.5B-Instruct模型构建一个使用GRPO的分布式强化学习流程。
不过，在我们深入这份教程之前，先简单介绍一下它的作者。AndriyBurkov算得上是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）流程，从而可以针对数学、逻辑和编程任务对语言模型进行微调。
首先需要明确，这些任务都存在一个唯一且正确的groundtruth答案，可通过简单的字符串比较轻松加以验证。
GRPO的发明者是DeepSeek，最早是被用于微调DeepSeek的R1和R1-Zero模型——它们可通过学习生成思维链（CoT）来更好地解决数学和逻辑问题。
本教程的目标是将通用语言模型Qwen2.5-1.5B-Instruct转换为数学问题求解器。我们将从头开始编写GRPO代码，然后将其与几个流行的库和工具集成起来，以实现分布式训练管道流程，包括：
PyTorch：用于张量运算和分布式训练。HuggingFaceTransformers：用于加载预训练的语言模型和tokenizer。FlashAttention2：优化的注意力机制，有助于减少内存使用量并提高训练速度。Weights&Biases(wandb)：用于实验跟踪、可视化和模型版本控制。
本教程分为几个部分。首先是基本设置和导入，然后是数据格式化和答案提取、数据集准备、评估函数、奖励函数、训练设置和执行，最后加载和测试模型。此过程中，我们将从头实现GRPO算法。
Part1：基础设置和导入
首先是安装并导入所有必要的模块。下面是导入库的一段代码截图。
部分代码截图。完整代码块参见GitHub。
运行上述代码（参考项目完整代码），可以执行以下任务：
设置随机种子：set_random_seed函数通过为Python的随机模块、NumPy和PyTorch设置种子，确保可复现性；环境变量配置：设置WANDB_API_KEY和WANDB_PROJECT环境变量，以启用与Weights&Biases的实验跟踪；导入必要的库，包括random、copy、re、torch等等。
Part2：数据格式以及答案提取
接下来，项目定义了数据格式，以及模型如何从输出和数据集中提取答案段落。
为了确保模型输出格式一致，项目还定义了一个系统提示。该提示指示模型生成包含和标签的输出。这一步通过两个函数完成：
extract_answer_from_model_output：此函数获取模型的输出文本，并提取标签内的内容；extract_answer_from_dataset：此函数从GSM8K数据集中提取预期答案，该数据集使用“####”分隔符来分隔答案：
部分代码截图。完整代码块参见GitHub。
Part3：数据准备
该项目使用GSM8K数据集进行训练。项目使用了该数据集中的示例来训练模型，基于强化学习（RL）训练范式，让模型生成多个问题解答样本，之后作者将这些解答与GSM8K示例中的标准答案进行对比，如果匹配，就为RL算法（GRPO）提供高奖励，然后更新模型权重，以增加模型下次获得高奖励的可能性。
实验过程是这样的。首先从HuggingFace加载数据集，然后格式化每个示例，包括系统提示和用户提示。这段实现代码中还定义了两个辅助函数：prepare_dataset以及build_prompt。
部分代码截图。完整代码块参见GitHub。
Part4：评估函数
评估对于跟踪模型的进展至关重要。因此作者定义了一些函数，从而可以在一组示例上对模型进行评估。该项目的评估函数执行以下任务：
token化提示并生成响应：模型的输出是在token化提示的基础上生成的。提取预测答案：从生成的响应中提取答案。将预测答案与预期答案进行比较：这种比较是通过精确匹配以及数值等价检查来完成的。
在这段代码中，两个辅助函数_extract_last_number和_extract_single_number被用来从文本中提取数字。评估函数evaluate_model使用这些辅助函数来确定预测答案是否正确：
部分代码截图。完整代码块参见GitHub。
Part5：奖励函数
在强化学习中，奖励函数是必不可缺的，作者定义了两个奖励函数：
correctness_reward：这个函数根据生成的答案是否正确来分配奖励。采用两种方式：精确的字符串匹配和数值等价检查，将模型输出的答案与预期答案进行比较。完全匹配会获得更高的奖励（2.0），而基于数值等价的匹配会获得较小的奖励（1.5）。
format_reward：这个函数鼓励模型遵循所需的类似XML的输出格式。它为生成文本中存在、、
和
标签提供小额奖励。
部分代码截图。完整代码块参见GitHub。
Part6：从头开始实现DataParallelGRPO
这一节，我们将从头实现GRPO算法的所有构建模块。首先，这里假设运行代码的机器至少有2台GPU。为此，这里要使用PyTorch的DataParallelAPI来将策略模型放在多个GPU核心上，每个GPU核心都有该模型的一个副本。然后将批量数据分散在这些GPU核心上完成处理。
部分代码截图。完整代码块参见GitHub。
Part7：训练设置和执行
这一节，我们将所有组件组合在一起，完成设置并开始训练。
首先，加载预训练的模型和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的padtoken设置为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目录中。
部分代码截图。完整代码块参见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的想法？
