RLHF+PPO实战:从奖励模型到策略优化的完整方案

📅 2026/7/24 16:49:28
RLHF+PPO实战:从奖励模型到策略优化的完整方案
RLHFPPO实战从奖励模型到策略优化的完整方案一、引言大语言模型的训练分为三个阶段预训练Pretraining→ 监督微调SFT→ 人类反馈强化学习RLHF。RLHF 是让模型对齐人类偏好的关键也是 ChatGPT 成功的核心秘密。本文将完整实现 RLHF 全流程奖励模型训练、PPO 策略优化、训练稳定性技巧使用 TRL 框架在单卡上完成。二、RLHF 原理2.1 三阶段流程阶段1: SFT → 学习回答格式和内容 阶段2: 奖励模型(RM) → 学习判断回答好坏 阶段3: PPO → 用RM信号优化策略2.2 PPO 核心公式L(θ) E[min(r_t(θ)Â_t, clip(r_t(θ), 1-ε, 1ε)Â_t)] r_t(θ) π_θ(a_t|s_t) / π_old(a_t|s_t)2.3 奖励函数设计总奖励 RM_score KL_penalty KL_penalty β * KL(π_θ || π_ref)β 控制新策略偏离参考模型的程度通常在 0.02-0.2 之间。三、奖励模型训练3.1 偏好数据集// 数据格式每条包含 chosen优选和 rejected劣选{prompt:解释什么是量子计算,chosen:量子计算利用量子比特的叠加和纠缠特性实现经典计算机难以完成的计算任务。核心原理包括量子叠加qubit同时处于0和1和量子纠缠远距离qubit关联。应用包括Shor算法破译RSA加密、Grover搜索加速、量子化学模拟。,rejected:量子计算就是用量子计算机做计算比普通计算机快很多。}3.2 奖励模型训练代码fromtransformersimport(AutoModelForSequenceClassification,AutoTokenizer,TrainingArguments,Trainer)fromdatasetsimportload_datasetimporttorchimporttorch.nnasnnclassRewardModel(nn.Module):基于预训练 LLM 的奖励模型def__init__(self,base_modelQwen/Qwen2.5-0.5B):super().__init__()self.modelAutoModelForSequenceClassification.from_pretrained(base_model,num_labels1,# 输出标量奖励torch_dtypetorch.bfloat16)self.tokenizerAutoTokenizer.from_pretrained(base_model)self.tokenizer.pad_tokenself.tokenizer.eos_tokendefforward(self,input_ids,attention_mask):outputsself.model(input_idsinput_ids,attention_maskattention_mask)returnoutputs.logits# [B, 1]defcompute_reward(self,texts:list)-torch.Tensor:计算文本奖励分数inputsself.tokenizer(texts,paddingTrue,truncationTrue,max_length1024,return_tensorspt)returnself.forward(**inputs).squeeze(-1)deftrain_reward_model():训练奖励模型modelRewardModel()optimizertorch.optim.AdamW(model.parameters(),lr1e-5)# 偏好损失RM(chosen) RM(rejected)forbatchindataloader:# 计算 chosen 和 rejected 的奖励rewards_chosenmodel.compute_reward(batch[chosen])rewards_rejectedmodel.compute_reward(batch[rejected])# 损失 -log(σ(r_chosen - r_rejected))loss-torch.nn.functional.logsigmoid(rewards_chosen-rewards_rejected).mean()optimizer.zero_grad()loss.backward()optimizer.step()# 计算准确率chosen 奖励 rejected 奖励的比例accuracy(rewards_chosenrewards_rejected).float().mean()print(fLoss:{loss.item():.4f}, Accuracy:{accuracy.item():.2%})四、PPO 策略优化4.1 TRL 框架 PPO 训练fromtrlimportPPOTrainer,PPOConfig,AutoModelForCausalLMWithValueHeadfromtrl.coreimportLengthSamplerfromtransformersimportAutoTokenizerimporttorch# 1. 配置configPPOConfig(model_nameQwen2.5-0.5B-SFT,learning_rate1.41e-5,batch_size8,mini_batch_size4,ppo_epochs4,kl_penaltykl,kl_penalty_directionboth,# 双向KL惩罚init_kl_coef0.2,# 初始KL系数target6,# KL目标值horizon10000,gamma1,lam0.95,cliprange0.2,# PPO裁剪范围εcliprange_value0.2,vf_coef0.1,# 价值函数损失权重)# 2. 加载模型modelAutoModelForCausalLMWithValueHead.from_pretrained(config.model_name,torch_dtypetorch.bfloat16)ref_modelAutoModelForCausalLMWithValueHead.from_pretrained(config.model_name,torch_dtypetorch.bfloat16)tokenizerAutoTokenizer.from_pretrained(config.model_name)tokenizer.pad_tokentokenizer.eos_token# 3. PPO Trainerppo_trainerPPOTrainer(configconfig,modelmodel,ref_modelref_model,tokenizertokenizer,datasetdataset,data_collatorcollator,)# 4. 训练循环generation_kwargs{min_length:-1,top_k:0.0,top_p:1.0,do_sample:True,pad_token_id:tokenizer.eos_token_id,max_new_tokens:256,}reward_modelRewardModel()forepoch,batchinenumerate(ppo_trainer.dataloader):query_tensorsbatch[input_ids]# 第一步采样响应 response_tensorsppo_trainer.generate(query_tensors,return_promptFalse,**generation_kwargs)batch[response][tokenizer.decode(r)forrinresponse_tensors]# 第二步计算奖励 # 任务奖励RM打分texts[qrforq,rinzip(batch[query],batch[response])]task_rewardsreward_model.compute_reward(texts)# 归一化均值为0方差为1rewards(task_rewards-task_rewards.mean())/(task_rewards.std()1e-8)# 第三步PPO更新 statsppo_trainer.step(query_tensors,response_tensors,rewards)# 监控指标print(fEpoch{epoch}: freward{rewards.mean():.3f}, fkl{stats[objective/kl]:.4f}, fclipfrac{stats[ppo/clipfrac]:.4f})4.2 PPO 训练稳定性技巧classPPOTrainingStabilizer:PPO训练稳定性增强staticmethoddefadaptive_kl_penalty(kl_div,target_kl0.02):自适应KL惩罚系数ifkl_divtarget_kl*2:return2.0# 增大惩罚elifkl_divtarget_kl/2:return0.5# 减小惩罚return1.0staticmethoddefreward_normalization(rewards,methodzscore):奖励归一化防止梯度爆炸ifmethodzscore:return(rewards-rewards.mean())/(rewards.std()1e-8)elifmethodminmax:r_min,r_maxrewards.min(),rewards.max()return(rewards-r_min)/(r_max-r_min1e-8)staticmethoddefreward_clipping(rewards,clip_range5.0):奖励裁剪防止极端值returntorch.clamp(rewards,-clip_range,clip_range)staticmethoddefearly_stop_check(kl_div,threshold0.5):KL散度过大时早停returnkl_divthresholdstaticmethoddefgradient_clipping(model,max_norm1.0):梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(),max_normmax_norm)五、DPO vs KTO vs PPO方法需要奖励模型需要偏好数据训练成本效果PPO✅❌高需在线采样最强DPO❌✅中强KTO❌部分低无需pair中ORPO❌❌低联合训练中上DPO 实现defdpo_loss(model,ref_model,batch,beta0.1):DPO损失函数# 计算当前模型对 chosen/rejected 的log概率logps_chosenmodel(**batch[chosen]).logits logps_rejectedmodel(**batch[rejected]).logits# 计算参考模型的log概率withtorch.no_grad():ref_logps_chosenref_model(**batch[chosen]).logits ref_logps_rejectedref_model(**batch[rejected]).logits# DPO损失log_ratio_chosenlogps_chosen-ref_logps_chosen log_ratio_rejectedlogps_rejected-ref_logps_rejected loss-torch.nn.functional.logsigmoid(beta*(log_ratio_chosen-log_ratio_rejected)).mean()returnloss六、关键调参指南参数推荐值作用learning_rate1e-6 ~ 1e-5PPO学习率应远小于SFTkl_coef0.02 ~ 0.2KL惩罚系数越大越保守cliprange(ε)0.1 ~ 0.2PPO裁剪范围mini_batch_size4 ~ 8每次更新的样本数ppo_epochs4每批数据重复训练的轮数target_kl0.01 ~ 0.06KL目标值超过则停止训练监控仪表盘Epoch | Reward | KL | ClipFrac | ValueLoss | PolicyLoss | LR 10 | 0.42 | 0.023 | 0.05 | 0.12 | 0.0015 | 1.4e-5 20 | 0.65 | 0.031 | 0.08 | 0.09 | 0.0012 | 1.3e-5 30 | 0.78 | 0.045 | 0.06 | 0.08 | 0.0008 | 1.2e-5 50 | 0.85 | 0.058 | 0.04 | 0.07 | 0.0006 | 1.0e-5 100 | 0.89 | 0.052 | 0.03 | 0.06 | 0.0005 | 8.0e-6 ✅ Reward持续上升 ✅ KL在目标范围内(0.06) ✅ ClipFrac下降(策略收敛) ✅ Loss平稳下降七、总结RLHF 的核心要点奖励模型质量决定一切— RM 准确率应 70%KL惩罚不可忽视— 防止策略偏离太远导致胡言乱语奖励归一化— 防止梯度爆炸小批量— PPO 对 batch size 敏感4-8 即可监控指标— reward/kl/clipfrac 三者都要看DPO 是更简单的替代方案但 PPO 通过在线采样能更好地探索策略空间。生产环境中优先选择 DPO简单稳定追求极致效果时使用 PPO。