☰
LLM强化学习实战:GRPO训练主循环、KL控制与踩坑清单
2026/10/21 5:07:19 网站建设 项目流程

这个标题我盯了一阵子了。最近半年 LLM 强化学习的热度一直没下来,社区里流传的 reinforcement-learning-in-llm 实现库也从一个“看起来很高级的玩具”变成了不少团队做推理模型微调的首选参考。我花了三个晚上把这类代码从头到尾读了一遍,又用一个小规模模型完整跑通了一个基于 GRPO 的端到端训练流程,期间翻车了好几次,最后总算把训练主循环、参考模型、奖励归一化、KL 控制这一整套逻辑都理顺了。

这篇就是把我的解读和复现过程完整写下来。目标是给已经会做 SFT 微调、但还没真正碰过 RL 训练的人一条捷径:我会从代码反推开原理,讲清楚为什么现在的主流实现普遍放弃 PPO 而选 GRPO,然后把一组可以抄作业的配置放在最后。这是系列的第一篇,重点放在 GRPO 主循环、loss 计算和踩坑实录上,后面再聊性能优化和奖励模型怎么训。

1. 先搞清楚这套代码到底在训练什么

1.1 从 SFT 到 RL:训练目标发生了本质变化

做过 SFT 的人都知道,有监督微调本质上是在做 teacher forcing 的最大似然估计:给定 prompt,让模型去拟合标注好的 response 序列,loss 只关注“标答”长什么样。它的问题很明显:模型学到的是“答案的样子”,而不是“答案对错的标准”。你给它几万条格式规整的样本,它确实能模仿得像模像样,但遇到训练分布之外的题,照样会一本正经地胡说八道。

强化学习在这套代码里要做的事完全不同。它不是让模型去拟合一个固定答案,而是让模型自己生成多个候选回答,然后根据一个外部信号——比如答案是否正确、格式是否符合要求——来调整自己的生成策略。换句话说,SFT 是照抄标准答案,RL 是开卷考试写作文,写完交给老师打分,分数低的写法以后少用,分数高的写法以后多用。

我需要特别强调,这套 reinforcement-learning-in-llm 代码里的“训练对象”是一个 7B 量级的开放权重基座模型,而不是某个魔改过的专用架构。很多时候大家误以为 RL 训练需要特殊的模型结构,其实不需要。模型照样是标准的 Decoder-only Transformer,区别只存在于训练循环:多了采样生成、奖励计算、优势估计这几步。

1.2 为什么经典 PPO 在 LLM 场景里显得笨重

很多第一次接触 LLM 强化学习的人会先去看 PPO(Proximal Policy Optimization),然后被一堆概念劝退:critic 网络、GAE、广义优势估计、价值函数归一化……这些在传统机器人控制、游戏 AI 里都很好用,但搬到语言模型上立刻变得非常别扭。

问题的根源在于,LLM 的“状态”是当前已经生成的 token 序列,而动作空间是整个词表,动辄几万甚至十几万个候选。想让一个价值网络去准确预测“当前生成到某个 token 时,未来的预期总回报是多少”,这个回归任务本身就比生成任务还难。我见过不少 PPO 微调实验,训了几个 epoch 之后 critic 的 loss 还在那里震荡,策略网络反而被带偏了。

所以现在社区里传播最广的这类实现,普遍换成了 GRPO(Group Relative Policy Optimization)。它的核心思路很“暴力”:我不需要价值网络去估计基线,我只需要对同一个 prompt 采样出 G 条回答,然后在这组样本内部做奖励归一化,就能得到每个样本的优势值。你比同组平均水平高,就是正样本,就把它往高推;比平均水平低,就是负样本,就让它往低退。这个“组内相对比较”的做法,把 PPO 里最复杂的优势估计环节直接干掉了,训练稳定性和实现难度都改善了一大截。

1.3 一个典型的代码库文件结构与整体流程

我这次解读的实现库,文件结构大概是这样的(为了叙述方便,我管它叫 rl-lab):

rl-lab/ ├── train_rl.py # 训练入口:超参、数据加载、主循环 ├── data_utils.py # prompt 数据读取与打包 ├── model_utils.py # 策略模型、参考模型加载与保存 ├── reward_utils.py # 规则奖励与格式化检查 ├── rollout_utils.py # 采样生成与 logprob 计算 ├── grpo_trainer.py # GRPO 训练主循环与 loss 计算 └── eval_utils.py # 验证集评估

整个训练流程是一条单向管道:

  1. 从数据集中采样一个 batch 的 prompt;
  2. 用当前策略模型对每个 prompt 生成 G 条回答,同时记录每条回答的 token logprob;
  3. 用参考模型(通常是 SFT 阶段的模型)给这些回答打一遍分,计算出 ref logprob,用于后续 KL 散度控制;
  4. 根据规则奖励或奖励模型给每条回答打分;
  5. 在组内对奖励做归一化,得到每个样本的优势值;
  6. 用这些优势值去更新策略模型,更新的时候做 ratio 计算和 clip 约束;
  7. 周期性保存 checkpoint,并跑一遍验证集看真实效果。

这个顺序是固定的,每一步都不能乱。很多人第一次复现失败,就是因为在第 4 步之前就把模型参数更新了,导致后续所有 logprob 对不上。

2. 核心概念与数学符号速览

2.1 奖励从哪里来:规则奖励与奖励模型

训练信号是整个 RL 循环里最容易被低估的部分。我阅读的这个代码库里,奖励来源分两类。

一类是规则奖励,也就是写一个纯 Python 函数去判断模型输出是否符合预期。比如做数学题,代码里会预设一个标准答案或者表达式校验逻辑,模型输出的结果能被解析出来,并且数值或符号表达式一致,就给 1 分,否则给 0 分。对于格式要求,比如“必须包含解析过程后给出最终答案”,则用字符串匹配或正则来判断。

另一类是奖励模型,用一个专门的打分模型给生成结果打一个连续分值。奖励模型一般是从人类偏好数据里训练出来的,它可以捕捉规则写不出来的细微偏好,但缺点是本身的容量和偏差会直接影响训练效果。在这个代码库里,作者的选择是优先用规则奖励,只有规则覆盖不了的任务才切换成奖励模型。我给这个设计点个赞,因为规则奖励是白盒、可解释、不会漂移,奖励模型却可能被模型生成分布的变化骗过去。

我抄一段这个库里 reward_utils 的示意逻辑,思路大致是这样的:

def reward_function(prompt, response): # 示意代码:规则奖励,不依赖外部库 answer = extract_answer_text(response) if answer is None: return 0.0 if ans_str_to_float(answer) == expected_answer(prompt): return 1.0 return 0.0

这类逻辑看起来简单,但它有个隐蔽的坑:如果解析函数写的太宽松,模型很快会学会“怎么让解析器通过”,而不是“怎么把题做对”。我后面会在排障清单里专门讲这一点。

2.2 参考模型与 KL 散度:为什么要拉住模型别飘

这是整个 LLM 强化学习里头最容易劝退新人的概念,但也是这套代码里实现得最优雅的部分。

强化学习的目标是最大化奖励,但如果你让模型在这个目标上自由狂奔,它很快就会走向奖励黑客:生成一些在字符串层面能骗过奖励函数、但人类一眼就看出来是废话的内容。为了防止这种漂移,代码里会保持一个“参考模型”,它就是 SFT 训练完成之后、进入 RL 阶段之前的那份模型权重。参考模型的参数在整个 RL 过程中是冻结的。

每一轮训练时,策略模型生成一个 response,代码会同时计算策略模型对这个 response 的 token logprob,以及参考模型对同一 response 的 token logprob。两者的差异就是 KL 散度的估计值。如果策略模型开始朝着偏离参考模型的方向走,KL 散度就会变大,我们就在 loss 里减去这个值,相当于给漂移踩刹车。

代码里比较常见的是 k3 估计器:

# 示意代码:token 级 KL 估计(k3) kl = (policy_logprob - ref_logprob).exp() - 1.0 - (policy_logprob - ref_logprob)

这个式子表面看着有点绕,其实就是用指数形式的散度近似来避免方差过大。实际操作里不需要逐 token 手写,框架内部都会封装好,但理解它可以帮助你读懂日志里的 KL 曲线:如果 KL 在一路暴涨,说明你的 β 系数太小或学习率太大,模型正在脱离参考分布。

2.3 优势函数:一组奖励如何变成可学习的信号

GRPO 最核心的数学动作就是组内归一化。假设同一个 prompt 采样了 G 条回答,各自的奖励是 r_1 到 r_g,那么第 i 条回答的优势值就是:

A_i = (r_i - mean(r)) / (std(r) + eps)

这其实是在表达一个非常朴素的思想:这条回答跟同组比,到底是好还是坏。好处第一是不用训价值网络,省掉一整块显存和训练负担;好处第二是它天然做出正负对比,模型能同时看到“什么该学”和“什么该避开”,训练信号量比单样本绝对奖励丰富得多。

但这里有一个非常关键的注意点:如果 G 设置得太小,比如只采样 2 到 4 条,那么组内的均值和方差会很不稳定,优势值的波动会很大;如果 G 太大,比如 32 以上,训练很稳,但生成阶段的时间开销会线性上涨。这个库的默认值一般落在 8 到 16,我在复现时使用 G=16,效果比 G=4 明显稳定。

3. 核心实现细节:GRPO 训练循环逐行拆解

3.1 主循环:采样生成与旧策略 logprob 的保存

我把 grpo_trainer.py 里最核心的训练循环拆出来,逻辑大致长这样:

for step, batch in enumerate(dataloader): prompts = batch["prompt"] # rollout 阶段:不更新梯度,只生成样本并记录旧 logprob with torch.no_grad(): seqs, old_logprobs = policy_model.generate( prompts, G=G, # 每个 prompt 生成 G 条 temperature=0.8, top_p=0.95, max_new_tokens=max_len, return_logprobs=True # 关键:返回每个 token 的 logprob ) # 参考模型计算 ref logprob,同样不更新梯度 with torch.no_grad(): ref_logprobs = ref_model(seqs) # 计算所有样本的奖励 rewards = [reward_fn(p, s) for p, s in zip(prompts, seqs)] # 按组 reshape 并做归一化,得到 advantage rewards = torch.tensor(rewards).view(-1, G) advantages = (rewards - rewards.mean(dim=-1, keepdim=True)) / (rewards.std(dim=-1, keepdim=True) + 1e-4) # 训练阶段:用 advantage 更新策略模型 for _ in range(update_epochs): # 这里的 logprobs 需要重新前向计算,并且要 detach old_logprobs new_logprobs = policy_model(seqs) # 带梯度 loss = grpo_loss(new_logprobs, old_logprobs, ref_logprobs, advantages) loss.backward() optimizer.step()

这里面有一个看起来不起眼、但直接决定成败的细节:模型生成时需要返回每个 token 的 logprob,而且必须是采样完成时的那个 logprob,不是训练阶段前向再算一遍的 logprob。原因很简单,训练阶段我们要算“新旧比例”:

ratio = exp(new_logprob - old_logprob)

如果 old logprob 来源不对,ratio 就全错了,整个 RL 更新就失去了意义。所以代码里 rollout 阶段和训练阶段是严格分开的,前者没有任何梯度,后者才走反向传播。

3.2 策略损失与 KL 惩罚项的代码实现

接下来看 loss 函数。GRPO 的 loss 在代码里一般写成一个小组件,逻辑很清晰:

def grpo_loss(new_logprobs, old_logprobs, ref_logprobs, advantages): # 1. 新旧策略比例 ratio = torch.exp(new_logprobs - old_logprobs) # 2. clip 约束:ratio 超过 [1-eps, 1+eps] 就截断 eps = 0.2 clipped_ratio = torch.clamp(ratio, 1.0 - eps, 1.0 + eps) # 3. 策略梯度损失:取 min 保证不要更新过猛 policy_loss = -torch.min(ratio * advantages, clipped_ratio * advantages) # 4. KL 惩罚:约束策略模型不要偏离参考模型太远 kl_loss = torch.exp(new_logprobs - ref_logprobs) - 1.0 - (new_logprobs - ref_logprobs) loss = policy_loss + beta * kl_loss return loss.mean()

这个实现里有几个“为什么”值得讲清楚。

首先,clip 操作本质上是在控制单步更新的幅度。如果没有 clip,某条高优势样本的 ratio 可能瞬间变得很大,导致模型一次更新就把某个 token 的概率推得极高,后面想拉回来就难了。有了 clip,无论优势多大,ratio 都被限制在 0.8 到 1.2 范围内,每一步最多只能小幅调整策略。这是 PPO 家族能稳定训练的核心机制。

其次,KL 惩罚项不应该加在 ratio 的分子上,而是作为独立的一项加到最终 loss 里。有些新手实现会把 KL 直接揉到奖励里,然后跟着优势一起归一化,这其实改变了优势的量纲,会让平衡奖励和 KL 的 β 系数变得非常难调。独立加项之后,β 的语义很单纯:偏离参考模型一个 KL 单位,等价于多少奖励损失。

最后,训练阶段需要把 old_logprobs 和 ref_logprobs 都做 detach。旧策略值是一个既定事实,不能让它参与梯度回传,否则梯度路径会乱掉。

3.3 参考模型推理与缓存技巧

参考模型在每次训练迭代都要跑一次前向,计算量跟策略模型差不多。如果模型是 7B 量级,并且组大小是 16,那参考模型的推理开销很容易成为整个训练流程的瓶颈。

rl-lab 代码里做了一个很实用的优化:在 rollout 阶段一次性把参考模型对这条生成的 ref_logprob 计算完,然后缓存在内存/磁盘里,训练阶段直接读取,不再重复前向。这个优化我的理解非常强烈,因为策略模型在 update_epochs 内会反复前向计算,如果每次前向都临时拉参考模型再算一次 KL,显存和耗时都会翻倍。

更进一步的优化是让参考模型和策略模型共享同一个底层模型权重。具体做法是:策略模型加载一份参数,参考模型不单独 allocate 一份完整权重,而是复用策略模型的参数,但在计算 ref logprob 时用 stop-gradient 隔离。不过这个共享方案在代码可读性上差一些,它要求你对模型的参数管理非常熟悉,否则很容易出现无意间把参考模型的权重也更新掉的情况。我建议第一次复现时老实一点,加载两份独立模型,用显存换正确性,跑通之后再优化。

4. 实操中踩过的坑,挑严重的说

4.1 把组大小 G 当成普通 batch size,显存直接爆掉

我第一轮复现时犯了这个错。习惯性地以为 batch size 8 就是每次前向 8 条样本,结果 rollout 阶段其实是每个 prompt 乘上 G,8 个 prompt 乘 16 就是 128 条完整序列,每条约 1024 个 token。这还没算 logprob 的存储开销。

解决方案是把 rollout 和训练彻底拆成两个阶段:生成阶段可以一段一段来,比如每次只生成 2 个 prompt 的 32 条回答,全部存到缓存里,然后再统一进训练阶段。训练阶段再按显存大小拆 micro-batch 做梯度累积。这样做的代价是代码结构会稍微复杂一点,但省显存效果立竿见影。

另外一个很实用的显存优化是使用 LoRA。这类 RL 代码默认可能支持全参数微调,但对于 7B 模型,我建议第一轮复现就先冻结基座权重,只训练一个低秩适配器。RL 阶段模型飘移的问题在 LoRA 下会更可控,而且显存占用直线下降。我最终跑通实验就是 LoRA 方案,训练稳定性反而更好。

4.2 KL 系数 β 不是越大越好

β 太大,模型被死死摁在参考模型附近,奖励几乎不涨;β 太小,KL 散度迅速抬头,生成内容开始出现奇怪的重复和格式违规。我踩过的直观表现是:日志里 reward 在涨,但人工翻生成的样本,发现模型的思考过程变成了同一句话循环十几次。

后来我把 β 从 0.5 一路下调,观察 KL 曲线的斜率,最终落在 0.05 附近才找到平衡点。建议你在自己的任务里也这样调:先固定其他参数,把 β 从 0.1 往上往下各扫两个量级,然后选一个 KL 缓慢上升但奖励保持增长的值。不要指望一个 β 通吃所有任务,规则奖励和奖励模型的分布不一样,最优 β 也会差不少。

4.3 奖励恒为 0:先查格式解析,别急着调算法

这个坑特别蠢,但特别容易遇到。我一开始用的数据集里,标准答案带了一些特殊符号,而我的答案解析函数写得很严格,导致模型明明答对了,却因为多一个空格或标点被判定为 0 分。结果就是整个 batch 的奖励全部是 0,优势全部是 0,模型根本学不到东西。

排查方法是把 reward 为 0 的样本原始输出打印出来,直接肉眼检查。如果发现答对了但解析失败,那就放宽解析逻辑,或者对回答做格式化预处理。我后来还专门给解析函数写了一组单元测试,把各种边界情况(空回答、带单位、带换行、大小写不一致)都覆盖了一遍,之后再没出现过这类问题。

4.4 优势全部为 0:任务难度和奖励区分度出了问题

当同一个 prompt 生成的 16 条回答全部做对或者全部做错时,组内归一化就会失效,因为 std 变成 0,优势全部是 0。我一轮实验中任务太简单,基座模型本身就能答对 95%,于是组内几乎全是满分样本,模型推不动。

这种情况有两个解决方向:一是提高任务难度,让同组样本产生区分度;二是把奖励从离散的 0/1 改成部分奖励,比如过程分加结果分,让同组样本即使全对也有高低差异。还有一个小技巧是把旧策略 checkpoint 东拉西扯混合进训练数据,故意拉低一些样本的正确率,给优势计算提供差异化空间。这个方法在一些实现里叫“数据混合热身”,我用了之后确实有效。

4.5 NaN 和 Inf:混合精度翻车

训练过程中 loss 突然变 NaN,第一反应不要怀疑算法,要查数值精度。GRPO 里有一个指数运算 exp(new_logprob - old_logprob),如果为 fp16,中间结果很容易下溢或溢出。我测试时因为用了不合适的精度设置,比率一下子就变成了 Inf。

解决方案是把模型的计算精度切到 bf16,或者至少在 loss 计算这一块保持 fp32。如果基础设施不支持 bf16,那就在计算 ratio 前对 old_logprobs 和 new_logprobs 做一次浮点升精度。还有一个小经验:检查一下 max_new_tokens 设置,如果响应被截断,某些 logprob 会被 padding token 污染,也可能带来异常数值。

4.6 复现不稳定的另一大坑:随机种子未统一

RL 训练相比 SFT 对随机性敏感得多。生成阶段有采样温度,如果你的采样器没有固定随机种子,同一份代码跑两次,结果可能差很远。

我在复现时发现,模型加载阶段和数据加载阶段的 seed 都设了还不够,因为 rollout 生成和训练阶段之间如果发生随机重排,就会导致某一轮的优势完全乱掉。最终我在 train_rl.py 的开头统一设置了全局随机种子,并显式传入给采样器,之后的对比实验才变得可复现。这里要特别提醒:如果你用到了批量推理加速工具,它本身可能带有异步线程,seed 控制要单独处理。

4.7 日志只看 loss 是不够的

我一开始只看 train loss 和 reward 均值,结果过了 200 步才发现模型生成的回答长度越来越离谱,从 512 个 token 膨胀到 2000 多个 token,这就是经典的奖励黑客信号——模型学会了通过凑字数让解析器“碰巧”提取到答案。

后来我在训练循环里额外记录了四个指标:KL 散度均值、生成回答的平均长度、格式正确率、以及奖励分布的标准差。这四个指标和策略 loss 放在一起,才能判断训练是否健康。尤其是“奖励分布的标准差”,它直接反映优势信号是否还存在,一旦标准差持续萎缩,说明模型已经收敛到同质化了。

我把这几个高频问题的排查要点合并成一个速查表:

问题现象排查方向常见解法
显存溢出rollout 阶段还是训练阶段拆分生成批次、开启梯度检查点、改 LoRA
reward 全是 0解析函数 / 数据格式打印原始输出,放宽解析,补单测
优势全是 0组内奖励无区分度提高难度、改连续奖励、混合旧 checkpoint
KL 爆炸beta 太小 / lr 太大调大 beta,降低学习率
NaN / Inffp16 溢出 / padding 污染换 bf16、loss 内部 fp32
结果不可复现seed 未统一全局固定 seed,并传给采样器
回答越来越长奖励黑客加长度惩罚、检查生成日志

5. 复现路线图与一份可抄的超参配置

5.1 三天跑通:从玩具数据集开始

我不建议任何人第一次复现就直奔 70B 全参数微调。比较稳的路线是先用一个小规模模型和一个可验证答案的数据集把整个流程跑通。我的做法是准备了 200 条数学运算题,每题都有确定整数答案,基座模型使用 7B 量级的开源权重,LoRA 方式接入训练。

第一天做数据准备和奖励函数,第二天跑 rollout 和 loss 验证,第三天调参并完成一次 20 步的小实验。小实验的目标不是效果惊艳,而是让下面三个数字动起来:reward 均值上升、KL 保持缓慢增长、格式正确率稳定在 80% 以上。只要这三个指标变化符合预期,就说明训练管线本身是通的,后面换大数据集和大模型只是在堆计算资源。

我那次小实验里,格式正确率从初始的 35% 左右爬到了 78%,reward 均值从 0.3 上升到 0.81,但生成长度也从 300 涨到了 600 多,最后靠加长度奖励惩罚才压回来。这说明小实验阶段发现的问题,放大到正式训练前解决,成本是最低的。

5.2 超参默认值速查清单

这里放一份我在复现时最终定下来的配置,应该覆盖了 GRPO 训练的核心选项。不同的任务类型可以在这个基础上调整,但不建议一次改动多个参数:

参数默认值说明与调整建议
组大小 G168 起步,显存允许可到 32
学习率1e-6RL 训练用极小学习率,防止策略突变
更新轮数 update_epochs1增大到 2-4 提升样本利用率,但注意过拟合
KL 系数 beta0.05观察 KL 曲线,0.01 到 0.1 之间搜索
clip epsilon0.2基本不需要动
生成温度0.8太低会减少探索,太高会让优势噪声变大
top_p0.95保持输出多样性
max_new_tokens1024和任务输出长度相关,要防止截断污染
梯度累积8显存不够时有效扩大等效 batch
LoRA rank647B 模型做 RL 时比较平衡

特别提醒学习率。SFT 时我们经常用 1e-5 甚至 5e-5,但 RL 阶段模型是在和环境交互,学习率过大会导致一步更新就崩掉。我见过最离谱的情况是学习率高了 3 倍,模型在 10 步之内就只会反复生成同一个 token。把学习率调低通常比调任何其他参数都管用。

5.3 从单机到多卡:扩展规模时的注意点

当你想把实验从 7B 放大到 70B,或者想把 batch size 显著增大时,代码的瓶颈会从算法变成工程调度。常见做法是让策略模型、参考模型、奖励模型分别落在不同的 GPU 或不同节点上,尽量减少单卡显存压力。

从单卡到多卡,最容易出问题的其实不是模型并行,而是 rollout 阶段和训练阶段的时间配比。如果策略模型生成太慢,训练卡就会一直空等。很多实现里会把生成和训练做成异步管道:一个进程负责采样生成,另一个进程负责梯度更新,中间通过缓存队列通信。虽然这会引入略微过时的策略梯度(不能严格保证用的是当前参数),但在实践中,一点延迟对收敛影响很有限,吞吐量却能提升好几倍。

模型并行方面建议先使用简单的张量并行或流水线并行,而不急着上复杂的显存调度。70B 模型做全参数 RL 时,LoRA 方案依然是显存最友好的选择。另外一个很容易忽略的点是参考模型的精度可以适当降低,甚至量化到 8bit,因为 KL 散度计算对参考模型的精度要求不如策略模型高。我在实验里给参考模型做了量化,训练效果几乎不受影响,显存却省出了不少。

最后分享一个我在复现过程中最有价值的体会:GRPO 这类算法确实大幅度降低了 LLM 强化学习的门槛,但它并没有把真正的难点干掉,难点只是从“调 GAE 和 critic”转移到了“调奖励函数和生成策略”。你花在奖励函数设计上的时间,应该比花在算法实现上的时间多得多。这个库值得你读、值得你复现、值得你改成适合自己业务的形状,但它不是开箱即用的大模型万能增强器,你还是要亲手去理解每一条 reward 规则背后到底在鼓励什么行为。

下一篇我会接着写 rollout 阶段吞吐优化的具体手段,以及如何从零训练一个奖励模型去覆盖规则覆盖不了的任务。到时候我们继续用这份代码做底子,慢慢把一条完整的 RL 训练链路拼齐。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询