PyTRIO快速入门实战篇(二):用 GRPO 提升 GSM8K 数学推理准确率

发布时间:2026/8/6 5:04:24
PyTRIO快速入门实战篇(二):用 GRPO 提升 GSM8K 数学推理准确率 用 PyTRIO 对 Qwen3.5-4B 进行 50 step 的 GRPO 训练。在 GSM8K test split 的前 100 道题上模型的准确率从 79% 提升到 95%训练与评估总成本仅 1.7 元。在上一篇实战中我们用监督微调让模型学习标准答案。这一篇换一种思路不给模型逐 token 的标准答案而是让模型自己生成多种解法再根据答案好坏获得奖励。本节我们将使用 PyTRIO 在 GSM8K 上完成一次 GRPO 训练并通过同一套评估脚本比较 base model 与训练后模型的数学推理准确率。话不多说我们先看结果模型GRPO step评估样本准确率Qwen3.5-4B Base010079%Qwen3.5-4B GRPO5010095%经过50个step的GRPO训练后Qwen3.5-4B在79%的准确率提升到95%的准确率。50step的训练消耗了0.29M的训练Token评估消耗了0.06M Token所用的成本仅为1.7元可以说很有性价比了训练 Token评估 Token总 Token训练与评估总成本0.29M0.06M0.35M1.7 元代码awesome-pytrio-train数据集gsm8k模型qwen3.5-4b框架pytrio任务介绍GSM8K 是一个小学数学应用题数据集。每条数据包含一道自然语言问题以及带有解题过程和最终答案的参考结果。例如题目可能会问Janet 的鸭子每天产 16 个蛋。她每天早餐吃 3 个并用 4 个做松饼剩下的每个卖 2 美元。她每天能赚多少钱一个符合要求的模型模拟回答可以是Janet 每天一共产 16 个蛋。 吃掉 3 个、做松饼用掉 4 个后还剩 16 - 3 - 4 9 个。 每个蛋卖 2 美元所以每天的收入是 9 × 2 18 美元。 因此Janet 每天能赚 \boxed{18} 美元。ps这段回答只是为了展示模型预期的输出结构并不是从某次 rollout 日志中截取的真实结果。可以看到模型需要写出一段可读的计算过程并把最终数值放进\boxed{18}评估脚本不会逐句判断推理过程而是提取回答中最后一个\boxed{...}。如果回答没有 boxed 答案脚本会退而提取最后一个数字再与 GSM8K 标准答案进行数值比较。GRPO 在做什么GRPO 的全称是 Group Relative Policy Optimization由 DeepSeek 在 2024 年发布的论文 《DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models》 中提出。在此之前大模型强化学习常使用 PPO。PPO 属于 actor-critic 方法除了正在训练的策略模型还需要训练一个 value model也就是 critic用它估计每个状态的价值并计算 advantage。原论文指出这个 value model 通常与策略模型规模相当会额外带来明显的显存和计算负担而大模型的 reward 往往只在回答结束时给出这也增加了为每个 token 学准 value function 的难度。GRPO 的解决方式是不再训练额外的 critic model。它让当前策略针对同一道题生成一组回答再根据这组回答的 reward 估计 baseline并判断每个回答相对组内水平是更好还是更差。这样既保留了 advantage 带来的相对训练信号又减少了 PPO 中 value model 所需的训练资源。回到这次实战它的核心流程可以概括为四步对同一道题采样一组答案。用可验证的 reward 函数为每个答案打分。用“当前答案 reward - 同组平均 reward”得到 advantage。提高组内高分答案的概率降低低分答案的概率。这里的“相对”很重要。为了方便说明我们先用 4 个回答举例假设它们的得分为[1.0, 1.0, 0.2, 0.0]组内平均分是0.55那么对应 advantage 就是[0.45, 0.45, -0.35, -0.55]。模型不需要额外的 critic而是直接从同组答案的比较中得到训练信号。本次实际训练的 group size 是 16也就是每道题会同时采样 16 个回答进行组内比较。如果一组答案的 reward 完全相同它们的 advantage 都是 0无法提供相对优劣信息。训练脚本会跳过这样的 group。为什么使用 PyTRIOPyTRIO 是 TRIO 远程大模型后训练和推理服务的 Python SDK。简单来说我们在本地用 Python 编写数据处理、reward 和训练循环而模型采样、前向反向传播、优化器更新及权重保存等计算密集型任务都由 PyTRIO 的远程服务完成。这种分工很适合想学习大模型后训练、但手边没有 GPU 集群的开发者。我们不需要先配置显卡环境、部署推理服务或搭建分布式训练系统只需要安装pytrio、运行trio login完成登录就可以从普通的本地 Python 环境发起模型训练。本地电脑主要负责控制实验流程因此这次 Qwen3.5-4B 的 GRPO 训练也不要求本地 GPU。PyTRIO 不只是把训练放到远程它还把训练和采样放进了同一套 SDK。在这篇文章的代码里我们先通过ServiceClient创建 LoRATrainingClient每个 step 再从当前训练权重得到SamplingClient模型完成一组 rollout 后本地代码计算 reward 和 advantage并构造Datum交回远程 trainer 更新参数。训练结束后我们可以直接保存 sampler 权重并用同一模型路径启动评估。这条链路尤其适合 GRPO。因为 GRPO 既需要反复采样多个回答又需要根据回答结果立即更新模型如果训练和推理分别使用两套基础设施实验代码和环境管理都会更复杂。使用 PyTRIO 后我们可以把注意力集中在真正影响效果的部分——prompt、reward、group-relative advantage 和训练参数——并用一份 Python 脚本完成从 rollout 到权重保存的闭环。从本次实验的结果看这套方式也足够轻量50 step 训练加评估共使用 0.35M Token实际成本为 1.7 元。对于第一次尝试 RL 后训练的读者它提供了一条不需要先购买硬件、同时又能完整理解 GRPO 数据流的实践路径。准备工作由于PyTRIO不挑设备所以不需准备带有GPU的机器我是在我的Macbook上完成的。PyTRIOpytrio.com基模Qwen/Qwen3.5-4B数据集openai/gsm8k训练代码train.py评估代码eval.py先进入示例目录并安装依赖cdpytrio-quick-start python-mpipinstall-Upytrio transformers datasets numpy addict然后登录 PyTRIOtrio login核心文件结构如下pytrio-quick-start/ ├── train.py # GSM8K 数据加载、GRPO rollout、reward 与训练 └── eval.py # base model 与 GRPO checkpoint 的异步评估数据不需要手动下载。脚本第一次运行时会通过 Hugging Facedatasets自动加载openai/gsm8k的main配置。在开始训练前可以先跑出 base model 的基线python eval.py--limit100这条命令使用 GSM8K test split 的前 100 道题默认设置temperature0.0和max_tokens512。本次实测得到的 base model exact accuracy 为79%。开始训练使用下面的命令启动 50 step GRPO 训练python train.py--steps50--batch-size2--group-size16--max-tokens512--eval-limit0--eval-limit 0表示训练结束后暂时跳过脚本内置的小规模评估。稍后我们会用独立的eval.py在完整的 100 条样本口径上比较结果。看到下面的打印时代表训练已经跑起来了本次训练采用的主要配置如下配置本次运行值作用base modelQwen/Qwen3.5-4B初始策略模型LoRA rank16训练低秩适配器数据加载范围train split 前 200 条候选训练题目实际训练题数前 100 条50 step × 每步 2 道题steps50参数更新次数batch size2每个 step 使用的题目数group size16每道题采样的回答数rollout temperature1.0保持采样多样性max tokens512单条回答的最大生成长度learning rate4e-5Adam 学习率seed42rollout 采样种子脚本默认加载 train split 的前 200 条数据但本次设置为 50 step、每步 2 道题因此实际依次使用其中的前 100 道题。每道题最多产生 16 条 rollout也就是每个 step 最多产生 32 条、本次训练最多产生 1600 条模型自生成轨迹无有效采样或 reward 完全相同的 group 会被跳过。这样的配置减少了每次更新覆盖的题目数量同时增加了同一道题下的候选回答数量。更大的 group 能为组内平均 reward 和相对 advantage 提供更丰富的比较样本这正是本次 GRPO 训练信号的来源。1. 让当前策略对同一道题采样多个答案每个训练 step 都先从当前 LoRA 权重创建 sampler再为每道题一次采样 16 个回答samplertrainer.save_weights_and_get_sampling_client()resultsampler.sample(prompttrio.ModelInput.from_ints(prompt_tokens),num_samplesgroup_size,sampling_paramsparams,return_textTrue,).result()这里必须使用“当前策略”的权重因为后续importance_sampling需要 rollout 生成时的 old logprobs。同步采样调用返回 future.result()表示等待远程采样完成并取得结果。2. 用数值正确性构造 reward训练代码不是只给 0 或 1。回答正确且使用\boxed{}时 reward 为 1.0回答正确但没有 boxed 格式时为 0.85答案错误但数值接近标准答案时会按相对误差得到较低的 shaping reward正确格式还能获得少量加分。ifexact:reward_value1.0ifboxedelse0.85elifpred_valueisnotNoneandgold_valueisnotNone:scalemax(abs(float(gold_value)),1.0)rel_errorabs(float(pred_value-gold_value))/scale reward_valuemax(0.0,0.45*(1.0-min(rel_error,1.0)))ifboxed:reward_value0.10else:reward_value0.10ifboxedelse0.0这种设计同时提供“答案是否正确”“数值是否接近”和“格式是否合规”三个层次的反馈。不过最终 79% 与 95% 的准确率只看答案是否与标准值精确相等不使用 shaping reward 作为准确率。3. 计算 group-relative advantage每个 completion 的 advantage 是它的 reward 减去同一道题所有有效 completion 的平均 rewardmean_rewardsum(rewards)/len(rewards)forsampleinsamples:sample[advantage]sample[reward]-mean_reward同组高于平均分的回答得到正 advantage低于平均分的回答得到负 advantage。如果一组 reward 的标准差接近 0代码会跳过该组避免提交一批全为 0 的训练信号。4. 对齐 token、old logprobs 与 advantageGRPO 在 PyTRIO 中使用importance_samplingloss。prompt token 只提供上下文不参与训练因此对应的 target、logprob 和 advantage 都用 0 占位completion 区间才填入真实值obs_lenlen(prompt_tokens)-1input_tokensprompt_tokenssample[tokens][:-1]target_tokens[0]*obs_lensample[tokens]old_logprobs[0.0]*obs_lensample[logprobs]advantages[0.0]*obs_len[sample[advantage]]*len(sample[tokens])datumtrio.Datum(model_inputtrio.ModelInput.from_ints(input_tokens),loss_fn_inputs{target_tokens:np.asarray(target_tokens,dtypenp.int64),logprobs:np.asarray(old_logprobs,dtypenp.float32),advantages:np.asarray(advantages,dtypenp.float32),},)input_tokens、target_tokens、old_logprobs和advantages的长度必须完全一致。这里obs_len len(prompt_tokens) - 1正是为了配合自回归预测时的一位右移。最后把有训练信号的Datum提交给远程 trainer并完成一次 Adam 更新fwdtrainer.forward_backward(datums,loss_fnimportance_sampling)opttrainer.optim_step(trio.AdamParams(learning_rate4e-5))metricsfwd.result().metrics opt.result()训练日志会逐 step 输出平均 reward、精确答对率、组内 reward 标准差、有效 group 数、跳过的同分 group 数、Datum 数量和 loss 指标。训练结束后脚本会打印可用于推理的 LoRA 权重路径Saved LoRA sampler weights: trio://...评估结果复制训练结束时打印的权重路径然后运行python eval.py --checkpoint-pathtrio://你的权重路径--limit100eval.py对 base model 和 checkpoint 使用相同的 test split 前 100 条数据、prompt 模板、答案解析逻辑、temperature0.0与max_tokens512。它会并发执行采样但并发只影响评估速度不改变计分方式。本次结果如下模型数据范围采样方式正确数Exact AccuracyQwen3.5-4B BaseGSM8K test 前 100 条temperature 079/10079%Qwen3.5-4B GRPO50 stepGSM8K test 前 100 条temperature 095/10095%经过 50 step GRPO 训练准确率从 79% 提升到 95%绝对提升 16 个百分点。这说明在本次小规模实验中基于可验证数学答案的 group-relative reward 已经能提供有效的强化学习信号。从资源消耗来看GRPO 训练使用了0.29M Token评估阶段使用了0.06M Token训练与评估合计 0.35M Token总成本为1.7 元。同时也要注意这里只评估了 test split 排序后的前 100 条数据结果来自一次训练运行并非完整测试集或多随机种子的平均值。因此它适合用于快速验证 GRPO 流程和训练方向不应直接当作模型在完整 GSM8K 上的最终成绩。这次实验最直观的感受是没想到用不到 2 元就能通过 RL 让一个 LLM 的准确率提高这么多。从 79% 到 95% 的结果也让我更直接地感受到当任务的答案可以被可靠验证时即使只进行 50 step 的小规模 GRPO 训练强化学习也可能带来很明显的收益。常用命令先评估 base modelpython eval.py--limit100运行与本文一致的 50 step 训练python train.py--steps50--batch-size2--group-size16--max-tokens512--eval-limit0评估训练后的 checkpointpython eval.py --checkpoint-pathtrio://你的权重路径--limit100如果只想先验证代码链路可以缩小数据、batch、group 和生成长度python train.py--limit8--steps2--batch-size2--group-size2--max-tokens128--eval-limit0