大模型RLHF中PPO训练全流程：从SFT到工程调优实战

1. 从SFT到PPO为什么大模型对齐绕不开这一步很多人第一次接触大模型训练流程时会有一个疑问既然已经做了SFT监督微调模型已经能听懂指令、能对话了为什么还要再折腾一轮PPO我刚开始做RLHF的时候也有这个困惑直到自己亲手跑了几轮训练、对比了SFT模型和PPO模型的输出才真正理解这一步的价值。SFT的本质是“模仿”。你给它一堆“问题-标准答案”对它学会的是在这些问题上复现人类写的答案。但问题在于人类写答案的成本极高你不可能为每一个可能的输入都准备一个标准答案。更关键的是SFT模型面对没见过的问题时它的行为是“插值”出来的可能对也可能跑偏。而PPO做的事情是不再告诉模型“标准答案是什么”而是告诉它“你给出的这个答案人类觉得好不好”。模型通过试错来学习什么样的输出能获得更高的奖励这个过程更接近人类学习的方式。从技术角度看PPOProximal Policy Optimization属于策略梯度方法的一种。它的核心思想是在更新策略时不要让新策略偏离旧策略太远。这个“不要太远”的约束就是通过裁剪clip机制实现的。为什么需要这个约束因为策略梯度方法的方差很大如果一步更新太猛策略可能直接崩掉后面再也救不回来。PPO通过限制每次更新的幅度让训练过程更稳定。在大模型RLHF的语境下PPO的四个模型角色需要先搞清楚Actor模型也就是我们要训练的策略模型通常是SFT之后的模型。它负责生成回复。Critic模型价值网络负责估计某个状态下的预期回报。它帮助计算优势函数降低策略梯度的方差。Reward模型事先用人类偏好数据训练好的打分模型给Actor生成的回复打分。Reference模型通常是SFT模型的副本冻结参数不更新。它的作用是计算KL散度防止Actor偏离原始语言模型太远避免输出退化。这四个模型的关系可以用一个简单的类比理解Actor是学生Critic是辅导老师Reward是考试评分标准Reference是“你原来的学习基础”。学生做题生成回复评分标准给分Reward辅导老师评估这道题相对于平均水平好不好Critic算优势同时要求学生的解题风格不能偏离原来的基础太远KL约束。我实测下来这四个模型的显存占用是最大的工程挑战。以7B模型为例如果全量微调Actor和Reference各需要约14GBFP16Critic需要约14GBReward模型如果也是7B规模又需要14GB。加上优化器状态和梯度单卡80GB的A100都未必够用。所以实际训练中通常采用LoRA微调Actor和CriticReward模型用较小的规模如1.5B或3B或者用奖励头代替完整的Reward模型。注意如果你打算从零开始搭建PPO训练流程建议先用小模型如1B以下跑通全流程确认数据管道、奖励计算、KL约束都正常工作后再扩展到7B或更大规模。直接上大模型调试一次OOM就能浪费半天时间。2. PPO训练流程的工程拆解从数据到梯度2.1 数据管道的设计细节PPO训练的数据和SFT阶段有本质区别。SFT的数据是“问题-答案”对而PPO的数据只需要“问题”prompt。模型自己生成回复然后由Reward模型打分。这意味着数据管道的设计需要考虑几个问题第一prompt的多样性和难度分布。如果所有prompt都很简单模型很快就能拿到高分训练信号会消失。如果太难模型一直拿低分梯度方向不明确。我的经验是prompt的难度应该呈现一个梯度分布大约30%简单、50%中等、20%困难。这样模型在训练过程中能持续获得有区分度的奖励信号。第二prompt的长度控制。大模型有上下文长度限制prompt太长会挤占生成空间。一般建议prompt长度控制在模型最大上下文长度的1/4到1/3。比如模型支持4096 tokenprompt控制在1024-1365 token比较合适。生成的最大长度也要设置合理太短可能截断有效回复太长会浪费计算资源。第三batch的组织方式。PPO训练中一个batch的数据需要经过“生成-打分-计算优势-更新”的完整流程。如果batch太大生成阶段耗时很长显存压力也大。如果batch太小梯度估计的方差会很大。实践中我通常用batch size 64-128prompt数量每个prompt生成1-4个回复。这样一次迭代能拿到64-512个样本足够计算稳定的梯度。2.2 生成阶段的关键参数生成阶段是PPO训练中最耗时的部分。Actor模型需要为每个prompt生成回复这个过程是自回归的逐个token生成。影响生成质量和速度的参数主要有temperature控制生成的随机性。训练初期可以设高一点如1.0让模型多探索训练后期降低如0.7让输出更稳定。top_p核采样参数通常设0.9-0.95。和temperature配合使用控制生成多样性。max_new_tokens最大生成长度。根据任务类型设置对话任务一般256-512长文本生成可能需要1024以上。repetition_penalty重复惩罚。设1.0-1.2之间太高会导致输出不自然。这里有一个容易踩的坑生成阶段和训练阶段的模型状态不一致。生成时模型处于eval模式dropout关闭训练时处于train模式。如果生成和训练之间没有正确切换会导致训练不稳定。我在早期实验中就遇到过这个问题生成时忘了调model.eval()结果dropout随机丢弃神经元同一个prompt每次生成的回复都不一样奖励信号噪声极大。2.3 奖励计算与KL约束的平衡Reward模型给每个生成的回复打分得到一个标量奖励。但这个奖励不能直接用来更新Actor还需要加上KL惩罚项。KL惩罚的作用是防止Actor偏离Reference模型太远。为什么需要这个约束因为Reward模型只是在人类偏好数据上训练的它不能覆盖所有情况。如果Actor为了拿高分生成一些Reward模型没见过但实际很差的回复就会发生“奖励黑客”reward hacking。KL惩罚项的计算方式是reward_total reward_score - kl_coef * kl_divergence。其中kl_coef是一个超参数控制约束的强度。这个参数的调优非常关键kl_coef太大Actor几乎不敢偏离Reference训练信号被压制模型学不到新东西。kl_coef太小Actor放飞自我输出退化可能变成重复无意义的文本。我的经验是kl_coef从0.01-0.1之间开始尝试。训练过程中监控KL散度的变化如果KL散度持续增长超过10-20说明约束太松如果KL散度一直接近0说明约束太紧。理想情况下KL散度应该在一个合理的范围内波动比如2-10之间。另外KL惩罚的计算方式有两种一种是每个token都算KL然后求和或平均另一种是只在序列末尾算一次。实践中token级别的KL更稳定但计算量更大。我通常用token级别然后取平均。2.4 优势函数的计算与GAE优势函数Advantage衡量的是在某个状态下采取某个动作比平均水平好多少。PPO用GAEGeneralized Advantage Estimation来计算优势公式涉及折扣因子γ和λ。这两个参数控制偏差和方差的权衡γ接近1考虑更长期的回报但方差大。λ接近1优势估计更准确但方差大。在大模型RLHF中通常γ1.0因为序列有明确结束λ0.95。GAE的计算需要Critic模型估计每个状态的价值。Critic的训练目标是拟合实际回报用MSE损失。这里有一个工程细节Critic模型的输出是每个token的价值估计但实际回报只在序列末尾才有。所以需要把末尾的回报反向传播到每个token用折扣因子累积。这个过程叫“回报计算”return computation。实现时要注意mask掉padding token否则padding位置的价值估计会干扰训练。3. 调优实战从loss曲线诊断训练问题3.1 如何读懂PPO的训练曲线PPO训练涉及多个loss和指标初学者容易被一堆曲线搞晕。我通常关注以下几个核心指标指标含义健康范围异常表现与对策policy_loss策略损失缓慢下降波动持续上升说明学习率太大或KL约束太松value_loss价值损失缓慢下降持续上升说明Critic学习率太大或网络容量不足reward平均奖励逐步上升不升反降说明Reward模型有问题或KL太紧klKL散度2-10之间波动持续增长说明kl_coef太小接近0说明太大entropy策略熵缓慢下降下降太快说明探索不足可能陷入局部最优clip_frac裁剪比例0.1-0.3太高说明更新幅度太大降低学习率我踩过的一个典型坑是reward一直在涨但人工评估发现模型输出越来越差。这就是典型的reward hacking。Reward模型被“钻空子”了Actor找到了一些能拿高分但实际很差的输出模式。解决办法是加强KL约束、增加Reward模型的多样性、或者在奖励中加入长度惩罚等辅助项。另一个常见问题是value_loss不下降。这通常是因为Critic模型容量不够或者学习率设置不当。Critic需要拟合的是每个token的累积回报这个目标比Actor的损失更难优化。我的经验是Critic的学习率可以比Actor大2-5倍或者给Critic用更大的网络。3.2 学习率与batch size的配合PPO对学习率非常敏感。学习率太大策略更新过猛KL散度爆炸学习率太小训练进度缓慢可能几天都看不到明显效果。我通常从1e-6到5e-6开始尝试Actor和Critic可以用不同的学习率。batch size的选择也有讲究。batch size越大梯度估计越准但计算成本越高。在显存允许的情况下我倾向于用较大的batch size如128-256然后配合线性缩放的学习率。如果显存不够可以用梯度累积来模拟大batch。这里有一个经验公式effective_batch_size batch_size * num_generations_per_prompt * gradient_accumulation_steps。这个有效batch size决定了梯度的噪声水平。一般来说有效batch size在512-2048之间比较合适。3.3 早停与模型选择PPO训练不像SFT那样有明确的收敛点。Reward可能一直在涨但模型的实际质量可能在某个点之后开始下降。所以需要设置早停策略。我通常用两个指标来判断第一KL散度。如果KL散度持续超过阈值如20说明模型已经偏离太远应该停止。第二人工评估。每隔一定步数采样一些输出进行人工评估。如果连续几次评估质量没有提升就停止训练。模型选择方面不要只保存最后一个checkpoint。我通常保存多个checkpoint然后用验证集上的奖励和人工评估综合选择。有时候中间某个checkpoint的效果比最终的好。4. 显存优化与分布式训练让7B模型跑起来4.1 显存占用的详细拆解7B模型PPO训练的显存占用可以拆解为以下几部分以FP16为例模型参数7B * 2 bytes 14GB。四个模型Actor、Critic、Reward、Reference如果都是7B就是56GB。优化器状态如果用Adam每个参数需要2个状态一阶矩和二阶矩每个状态2 bytes就是7B * 4 bytes 28GB。Actor和Critic都需要共56GB。梯度7B * 2 bytes 14GBActor和Critic共28GB。激活值取决于batch size和序列长度通常几GB到十几GB。生成时的KV Cache取决于batch size和生成长度可能几GB。总计轻松超过150GB单卡肯定放不下。所以必须做显存优化。4.2 实战中的显存优化手段我常用的显存优化手段按优先级排列第一LoRA微调。Actor和Critic都用LoRA只训练低秩矩阵参数量减少到原来的1%以下。优化器状态和梯度也相应减少。这是最有效的显存优化手段。第二梯度检查点。用时间换显存把激活值重新计算。可以节省50%-70%的激活显存但训练速度降低20%-30%。第三混合精度训练。用BF16代替FP32显存减半速度提升。但要注意数值稳定性某些操作可能需要FP32。第四ZeRO Stage 2/3。把优化器状态和梯度分片到多卡上。ZeRO Stage 3还能分片模型参数但通信开销大。第五Offload。把优化器状态和梯度放到CPU内存需要时再加载到GPU。速度慢但能跑起来。第六Reward模型用更小的规模。Reward模型不需要和Actor一样大1.5B或3B通常够用。或者用奖励头reward head代替完整模型。我实测下来7B模型用LoRA 梯度检查点 ZeRO Stage 2可以在8张A100 40GB上跑起来。如果只有4张卡需要加上Offload但速度会慢很多。4.3 分布式训练的通信瓶颈PPO训练的分布式策略和SFT不同。SFT可以用数据并行每个卡处理不同的数据梯度All-Reduce。但PPO涉及多个模型需要更复杂的并行策略。我通常用以下配置Actor和Reference数据并行 ZeRO Stage 2Critic数据并行 ZeRO Stage 2Reward数据并行不需要优化器状态生成阶段是通信瓶颈。每个卡上的Actor需要生成回复然后汇总到一起计算奖励。如果生成阶段用数据并行每个卡生成的样本不同需要All-Gather汇总。这个通信量很大尤其是生成长度较长时。一个优化技巧是生成阶段用流水线并行把模型按层切分到不同卡上减少单卡显存压力。但流水线并行有气泡问题需要仔细调度。5. 常见OOM场景与排查方法5.1 生成阶段OOM生成阶段OOM通常是因为KV Cache太大。KV Cache的大小是batch_size * num_layers * 2 * hidden_size * max_seq_len * 2 bytes。以7B模型为例32层hidden_size 4096batch_size 64max_seq_len 512KV Cache就是64 * 32 * 2 * 4096 * 512 * 2 17GB。这还只是一个模型。解决办法减小batch size、减小生成长度、用GQAGrouped Query Attention减少KV头数、或者用PagedAttention管理KV Cache。5.2 反向传播OOM反向传播OOM通常是因为激活值太大。激活值的大小和batch size、序列长度、模型层数成正比。解决办法梯度检查点、减小batch size、用FlashAttention减少激活值。5.3 优化器更新OOM优化器更新OOM通常是因为优化器状态太大。Adam的优化器状态是参数量的2倍FP32。解决办法用LoRA减少参数量、用ZeRO分片优化器状态、用8-bit Adam减少状态精度。我遇到过一个比较隐蔽的OOM训练过程中显存逐渐增长最后OOM。原因是PyTorch的缓存分配器没有及时释放显存。解决办法是设置PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True或者定期调用torch.cuda.empty_cache()。注意OOM报错时不要只看报错的那一行。PyTorch的OOM报错通常会显示显存分配的历史找到最大的那次分配那才是真正的瓶颈。6. 奖励模型的质量决定PPO的上限6.1 Reward模型的训练要点Reward模型的质量直接决定了PPO训练的效果。如果Reward模型本身不准Actor学到的就是错误的方向。Reward模型通常用人类偏好数据训练输入是一个prompt和两个回复输出是哪个回复更好。训练Reward模型时我关注几个要点第一偏好数据的质量和多样性。数据要覆盖不同的任务类型、不同的难度、不同的风格。如果数据有偏Reward模型也会有偏。第二Reward模型的规模。不需要和Actor一样大但也不能太小。我的经验是Reward模型至少要有Actor的1/4参数量否则区分能力不够。第三训练时的正则化。Reward模型容易过拟合需要用dropout、权重衰减等手段。另外可以用pairwise ranking loss而不是pointwise regression前者更稳定。6.2 奖励黑客的识别与缓解奖励黑客是PPO训练中最棘手的问题。表现是Reward分数很高但人工评估质量很差。常见的奖励黑客模式包括输出冗长但无意义的文本因为Reward模型可能偏好长回复。重复某些“安全”的短语因为Reward模型给这些短语高分。回避有争议的话题输出模棱两可的回复。缓解奖励黑客的方法第一加强KL约束。让Actor不要偏离Reference太远。第二在奖励中加入辅助项。比如长度惩罚、重复惩罚、多样性奖励。第三用多个Reward模型集成。不同的Reward模型有不同的偏好集成后更难被黑客。第四定期用人工评估校准Reward模型。发现Reward模型和人类判断不一致时及时修正。6.3 奖励归一化的技巧Reward模型的输出范围可能不稳定有的样本奖励很高有的很低。直接用来计算优势函数会导致梯度尺度不一致。所以需要对奖励做归一化。我通常用running mean和running std来做归一化normalized_reward (reward - running_mean) / (running_std eps)。running_mean和running_std在训练过程中动态更新。这样可以让奖励的分布保持稳定梯度尺度一致。另外不同prompt的奖励尺度可能不同。有的prompt容易拿高分有的难。可以用prompt级别的归一化或者用batch级别的归一化。我通常用batch级别简单有效。7. 从CartPole到LLMPPO核心逻辑的跨领域验证7.1 用Gymnasium CartPole理解PPO如果你对PPO的数学推导还是模糊我强烈建议先用Gymnasium的CartPole环境跑一遍PPO。CartPole是一个经典的控制任务一根杆子放在小车上通过左右移动小车保持杆子不倒。状态是4维位置、速度、角度、角速度动作是2维左、右。用CartPole跑PPO的好处是环境简单训练快几分钟就能看到效果。你可以直观地看到PPO的各个组件如何工作Actor网络输入状态输出动作概率。Critic网络输入状态输出价值估计。优势函数用GAE计算。裁剪机制限制策略更新幅度。我当初就是先用CartPole跑通了PPO才理解了大模型RLHF中的各个组件。CartPole的PPO代码大概100行但包含了PPO的所有核心逻辑。跑通之后再看大模型PPO的代码会发现结构完全一样只是模型换成了Transformer动作换成了token。7.2 从CartPole到大模型的关键差异虽然核心逻辑一样但从CartPole到大模型PPO有几个关键差异第一动作空间。CartPole是离散的2个动作大模型是词表大小的离散动作几万到十几万。这意味着Actor网络的输出维度从2变成了几万。第二状态表示。CartPole的状态是4维向量大模型的状态是token序列。Critic网络需要处理序列输入通常用Transformer的最后一层隐藏状态。第三奖励来源。CartPole的奖励是环境给的大模型的奖励是Reward模型给的。Reward模型本身有噪声增加了训练难度。第四训练成本。CartPole几分钟就能训练好大模型PPO需要几天甚至几周。理解这些差异有助于你在调试大模型PPO时知道哪些地方容易出问题。7.3 离线强化学习的启示如果你对离线强化学习Offline RL感兴趣IQLImplicit Q-Learning是一个值得关注的算法。IQL的核心思想是不直接估计行为策略的价值而是估计一个“隐式”的价值函数避免分布外动作的高估问题。IQL对大模型PPO的启示是Reward模型可以看作一个离线数据源Actor在训练过程中会探索Reward模型没见过的区域。如何约束Actor不跑到分布外是PPO和离线RL共同面临的问题。KL约束是一种解决方案IQL的隐式价值估计是另一种思路。8. 一些实战中的零散经验8.1 关于KL系数的动态调整固定KL系数往往不是最优的。训练初期模型需要多探索KL系数可以小一点训练后期模型需要稳定KL系数可以大一点。我通常用自适应KL设定一个目标KL范围如5-10如果实际KL低于下限减小kl_coef如果高于上限增大kl_coef。这样可以让KL散度保持在合理范围内。8.2 关于Critic的预训练Critic模型如果从随机初始化开始训练前期价值估计很不准优势函数噪声很大。一个技巧是先用SFT模型的奖励数据预训练Critic让它有一个合理的初始价值估计。或者用Reward模型的输出作为Critic的初始目标。8.3 关于生成样本的多样性如果每个prompt只生成一个回复样本多样性不足梯度估计方差大。我通常每个prompt生成2-4个回复然后分别计算奖励和优势。这样一次迭代能拿到更多样本梯度更稳定。但生成本身有成本需要在多样性和效率之间权衡。8.4 关于训练日志的监控PPO训练涉及很多指标建议用TensorBoard或WandB记录所有指标。我通常记录policy_loss、value_loss、reward、kl、entropy、clip_frac、learning_rate、grad_norm。这些指标能帮你快速定位问题。比如grad_norm突然变大说明有异常梯度entropy持续下降说明探索不足。8.5 关于checkpoint的管理PPO训练可能持续几天中间可能因为各种原因中断。建议定期保存checkpoint包括模型参数、优化器状态、学习率调度器状态、随机数种子。这样中断后可以恢复训练不用从头开始。我通常每1000步保存一次保留最近5个checkpoint。8.6 关于人工评估的频率自动指标如Reward分数不能完全反映模型质量。建议每隔一定步数如5000步做一次人工评估采样一些输出人工判断质量。如果发现自动指标和人工评估不一致及时调整训练策略。8.7 关于多轮对话的PPO上面的讨论主要针对单轮对话。多轮对话的PPO更复杂奖励只在对话结束时给出中间轮次没有奖励。这时需要用折扣因子把末尾奖励反向传播到每一轮。另外多轮对话的状态空间更大Critic更难训练。我通常先用单轮对话跑通PPO再扩展到多轮。8.8 关于不同基座模型的差异不同基座模型的PPO表现差异很大。有的模型对KL约束更敏感有的模型更容易reward hacking。建议在正式训练前先用小规模数据做一轮快速实验看看基座模型的行为特点。比如Llama系列和Qwen系列在PPO中的表现就有差异Llama对KL更敏感Qwen的探索能力更强。8.9 关于训练成本的估算PPO训练的成本远高于SFT。以7B模型为例SFT可能只需要几十GPU小时PPO可能需要几百到几千GPU小时。成本主要花在生成阶段。如果每个prompt生成4个回复生成长度512batch size 128那么一次迭代需要生成1284512 262144个token。以A100的生成速度约1000 token/s计算一次迭代需要约4分钟。如果训练10000步就是约666小时。这还不包括训练阶段的时间。所以PPO训练前一定要做好成本估算。8.10 关于失败案例的复盘我经历过几次PPO训练失败复盘下来主要有几个原因一是Reward模型质量差导致Actor学偏二是KL系数设置不当导致输出退化三是学习率太大导致训练崩溃四是数据管道有bug导致奖励计算错误。每次失败后我都会仔细检查数据管道和奖励计算确保没有低级错误。PPO训练很脆弱一个小bug就可能导致整个训练失败。最后分享一个我常用的调试技巧先用极小的数据集如10个prompt跑一遍完整流程确认所有组件正常工作。然后逐步增加数据量和模型规模。这样可以在早期发现大部分问题避免在大规模训练时浪费资源。