一、前言为什么需要Reward Model在大语言模型的对齐训练中**Reward Model奖励模型**是RLHFReinforcement Learning from Human Feedback流程的核心组件。它扮演着裁判的角色学习人类偏好为后续PPO训练提供奖励信号。本文将深入解析基于DeepSpeed-Chat框架的Reward Model训练全流程以Qwen3-0.6B/4B为例涵盖数据处理与偏好对构建Reward Model架构设计v_head原理Pairwise Ranking Loss详解单卡/多卡训练实战二、Reward Model在RLHF中的定位┌─────────────────────────────────────────────────────────┐ │ RLHF 三阶段训练 │ ├─────────────────────────────────────────────────────────┤ │ Stage 1: SFT (监督微调) │ │ └── 用高质量对话数据微调基础模型 → 得到 SFT Model │ │ │ │ Stage 2: Reward Model训练 ← 本文重点 │ │ └── 用偏好对比数据训练奖励模型 → 得到 RM │ │ │ │ Stage 3: PPO (强化学习) │ │ └── SFT Model RM → 通过PPO算法优化策略 │ └─────────────────────────────────────────────────────────┘Reward Model的核心任务给定同一个Prompt的两个不同Responsechosen vs rejected学会打分让chosen的分数显著高于rejected。三、数据处理构建偏好对比数据集3.1 数据格式与PromptDataset类DeepSpeed-Chat使用统一的data_utils.py处理三个阶段的数据通过train_phase参数区分classPromptDataset(Dataset): 三阶段数据加载器 train_phase: 1 - SFT阶段只返回chosen response 2 - RM阶段返回chosen rejected成对数据 3 - RL阶段只返回prompt用于生成 def__getitem__(self,idx):ifself.train_phase2:# RM关键同时返回chosen和reject的pairreturn(self.chosen_dataset[idx][input_ids],self.chosen_dataset[idx][attention_mask],self.reject_dataset[idx][input_ids],self.reject_dataset[idx][attention_mask])3.2 RM阶段数据处理流程defcreate_dataset_split(...,train_phase2):eliftrain_phase2:# 关键同时处理chosen和reject形成对比对fortmp_dataincurrent_dataset:chosen_sentenceraw_dataset.get_prompt_and_chosen(tmp_data)# 偏好回答reject_sentenceraw_dataset.get_prompt_and_rejected(tmp_data)# 非偏好回答ifchosen_sentenceandreject_sentence:# 都加上EOS token让模型学会何时结束chosen_tokentokenizer(chosen_sentenceeos_token,...)reject_tokentokenizer(reject_sentenceeos_token,...)chosen_dataset.append(chosen_token)reject_dataset.append(reject_token)关键细节成对处理每个样本必须同时包含chosen和rejected且基于同一个promptEOS标记添加|endoftext|教会模型何时停止生成长度对齐通过paddingmax_length确保同一对样本长度一致3.3 DataCollatorReward特殊的Batch组装classDataCollatorReward:def__call__(self,data):batch{}# 关键将chosen和rejected在batch维度拼接# 前一半是chosen后一半是rejectedbatch[input_ids]torch.cat([f[0]forfindata][f[2]forfindata],dim0)batch[attention_mask]torch.cat([f[1]forfindata][f[3]forfindata],dim0)returnbatch设计意图[chosen_batch; rejected_batch]的拼接方式让模型在一个前向传播中同时处理两个response高效计算pairwise loss。四、模型架构从Causal LM到Reward Model4.1 架构对比AutoModel vs AutoModelForCausalLM┌────────────────────────────────────────────────────────┐ │ AutoModelForCausalLM (SFT阶段使用) │ ├────────────────────────────────────────────────────────┤ │ Transformer Backbone LM Head │ │ ↓ ↓ │ │ Hidden States Linear(hidden_size, vocab_size)│ │ (batch, seq, hidden) → Logits │ │ (batch, seq, vocab_size) │ │ │ │ 用途预测下一个token学习生成文本 │ └────────────────────────────────────────────────────────┘ ┌────────────────────────────────────────────────────────┐ │ AutoModel (RM阶段使用) │ ├────────────────────────────────────────────────────────┤ │ Transformer Backbone Value Head │ │ ↓ ↓ │ │ Hidden States Linear(hidden_size, 1) │ │ (batch, seq, hidden) → Rewards │ │ (batch, seq) │ │ │ │ 用途为每个位置打分学习评估文本质量 │ └────────────────────────────────────────────────────────┘4.2 RewardModel核心实现classRewardModel(nn.Module):def__init__(self,base_model,tokenizer,...):super().__init__()self.rwtransformerbase_model# 主干网络无LM Head# 关键添加v_head将hidden_size映射到1维分数ifhasattr(self.config,word_embed_proj_dim):# OPT模型特殊处理self.v_headnn.Linear(self.config.word_embed_proj_dim,1,biasFalse)else:# GPT/LLaMA/Qwen等模型self.config.n_embdself.config.hidden_size self.v_headnn.Linear(self.config.n_embd,1,biasFalse)架构特点轻量级只增加1024个参数hidden_size→1相比原模型可忽略位置敏感为每个token位置输出分数而非整个序列一个分数灵活评分可取最后一个token、平均值、或特定位置作为最终奖励五、训练目标Pairwise Ranking Loss详解5.1 核心思想Reward Model不直接拟合绝对分数而是学习相对排序——确保chosen的分数高于rejected。5.2 前向传播流程forward函数defforward(self,input_ids,attention_mask,...):# Step 1: 通过主干网络获取hidden states# shape: (batch_size * 2, seq_len, hidden_size)transformer_outputsself.rwtransformer(input_ids,...)hidden_statestransformer_outputs[0]# Step 2: v_head为每个位置打分# shape: (batch_size * 2, seq_len)rewardsself.v_head(hidden_states).squeeze(-1)# Step 3: 拆分为chosen和rejectedbsinput_ids.shape[0]//2# 实际batch大小chosen_rewardsrewards[:bs]# 前一半是chosenrejected_rewardsrewards[bs:]# 后一半是rejected5.3 关键对齐片段的确定# 对于每个样本对找到有效对比区域foriinrange(bs):chosen_idchosen_ids[i]rejected_idrejected_ids[i]# 找到chosen和rejected第一个不同的位置response开始分叉check_divergence(chosen_id!rejected_id).nonzero()divergence_indcheck_divergence[0].item()# 找到各自的有效结束位置第一个padding tokenc_ind(chosen_idPAD_ID).nonzero()[0].item()r_ind(rejected_idPAD_ID).nonzero()[0].item()end_indmax(c_ind,r_ind)# 取较长的# 截取对齐片段从分叉点到结束点c_truncatedchosen_reward[divergence_ind:end_ind]r_truncatedrejected_reward[divergence_ind:end_ind]图示Prompt: [今天天气怎么样] Chosen: [今天天气怎么样][今天晴天适合出游][PAD][PAD] Rejected:[今天天气怎么样][不知道 maybe下雨][PAD][PAD] ↑divergence_ind ↑c_ind/r_ind 对比区域: [今天晴天适合出游] vs [不知道 maybe下雨] ↑ 高分期望 ↑ 低分期望5.4 损失函数LogSigmoid Ranking Loss# 核心损失最大化 P(chosen rejected) sigmoid(c_score - r_score)# 等价于最小化 -log(sigmoid(c_score - r_score))ifself.compute_fp32_loss:c_truncatedc_truncated.float()r_truncatedr_truncated.float()# 对对齐区域的每个位置计算ranking loss取平均loss-torch.nn.functional.logsigmoid(c_truncated_reward-r_truncated_reward).mean()# 最终得分取最后一个有效token与训练目标一致chosen_mean_scores.append(chosen_reward[c_ind-1])rejected_mean_scores.append(rejected_reward[r_ind-1])数学推导L − 1 N ∑ i 1 N log ⁡ σ ( r θ ( x , y c ) − r θ ( x , y r ) ) \mathcal{L} -\frac{1}{N}\sum_{i1}^N \log\sigma(r_\theta(x, y_c) - r_\theta(x, y_r))L−N1​i1∑N​logσ(rθ​(x,yc​)−rθ​(x,yr​))其中σ \sigmaσ是sigmoid函数目标是让r θ ( x , y c ) r θ ( x , y r ) r_\theta(x, y_c) r_\theta(x, y_r)rθ​(x,yc​)rθ​(x,yr​)。5.5 推理函数forward_value供PPO使用defforward_value(self,...,return_value_onlyFalse,prompt_length0):# 用于Stage 3 PPO阶段给单个response打分valuesself.v_head(hidden_states).squeeze(-1)ifreturn_value_only:returnvalues# 所有位置分数用于GAE计算# 找到response最后一个有效token的分数作为奖励foriinrange(bs):# 从prompt之后找第一个PADc_inds(input_id[prompt_length:]self.PAD_ID).nonzero()c_indc_inds[0].item()prompt_lengthiflen(c_inds)0elseseq_len chosen_end_scores.append(value[c_ind-1])return{values:values,# 全序列价值估计chosen_end_scores:torch.stack(chosen_end_scores)# 最终奖励}六、训练实战Qwen3-0.6B单卡训练6.1 训练脚本#!/bin/bashdeepspeed--num_gpus1main.py\--model_name_or_path/root/vscode/Tuling/Qwen3-0.6B\--data_path/path/to/train.jsonl\--data_split6,2,2\--num_train_epochs1\--gradient_accumulation_steps4\--per_device_train_batch_size1\# 实际处理2个responsechosenrejected--per_device_eval_batch_size1\--dropout0.0\# 关键保持排序一致性--max_seq_len512\--learning_rate5e-5\--zero_stage2\--dtypebf16\--output_dir./output关键参数说明参数值说明per_device_train_batch_size1实际处理1对chosenrejected 2个序列dropout0.0必须设为0避免训练/推理排序不一致zero_stage2ZeRO-2优化显存占用从5.55GB降至3.33GB6.2 训练日志分析# 初始状态Epoch 0 chosen_last_scores: 0.1047 # 随机初始化接近0 rejected_last_scores: 0.0614 # 与chosen接近 acc: 51.0% # 几乎随机猜测 # 训练结束Epoch 1 chosen_last_scores: -17.93 # 分数下降绝对值不重要 rejected_last_scores: -67.62 # 显著低于chosen acc: 99.0% # 完美区分偏好对 score差距: 49.69 # ↑1147倍提升关键观察绝对分数为负是正常的模型学习的是相对排序绝对值无意义差距扩大才是核心从0.04扩大到49.69说明模型学会了强区分准确率99%模型已充分学习训练集偏好6.3 显存与耗时指标数值峰值显存17.4 GB / 24 GB (4090)训练时间23分44秒处理速度~7条/秒训练集10,000对验证集1,000对七、模型对比Reward Model vs 原模型7.1 文件大小对比模型大小差异分析Qwen3-0.6B1.50 GB含LM Head (151936×1024≈1.55亿参数)Reward Model1.19 GB减小309MB (20.6%)移除LM Head7.2 参数量对比Qwen3-0.6B: 596,049,920 参数 (5.96亿) Reward Model: 595,779,584 参数 (5.96亿 - 27万) └─ 差值: 270,336 1024 × 264 (词表差异)词表变化Qwen3-0.6B词表151,936Reward Model词表151,672少了264个token缺失的主要是连续空格组合、编码错误字符等无意义token7.3 架构差异// Reward Model config.json 新增字段{n_embd:1024,// 为v_head统一hidden_size命名dtype:float32,// 训练精度end_token_id:151645,// 明确结束标记pad_token_id:151645// 与eos共用}八、多卡训练Qwen3-4B扩展8.1 多卡脚本调整deepspeed--num_gpus4main.py\--model_name_or_path/root/vscode/Tuling/Qwen3-4B\...# 其他参数同单卡多卡优势更大的有效batch size4卡 × 1 × 4梯度累积 16更快的训练速度支持更大模型4B参数需多卡才能训练九、关键经验与最佳实践9.1 训练稳定性问题解决方案排序不一致设置dropout0.0分数发散使用compute_fp32_lossTrue显存不足启用ZeRO-2/3 梯度检查点9.2 评估指标解读defevaluation_reward(model,dataloader):# 核心指标accuracy (chosen_score rejected_score)的比例correct(chosenrejected).sum()acccorrect/total# 辅助指标分数差距diffchosen_mean-rejected_mean# 越大说明区分度越强9.3 保存格式优化针对Qwen3模型建议保存为safetensors格式# 替换原有的save_hf_formatfromdschat.utils.utilsimportsave_hf_format_safetensorsifargs.global_rank0:save_hf_format_safetensors(rm_model,tokenizer,args)十、总结本文详细解析了DeepSpeed-Chat框架下Reward Model的完整训练流程数据处理成对构建chosen-rejected对比数据通过DataCollatorReward拼接batch模型架构AutoModel v_head轻量级扩展实现评分功能训练目标Pairwise Ranking Loss学习相对排序而非绝对分数工程实践ZeRO优化、bf16混合精度、safetensors格式保存Reward Model作为RLHF的裁判其训练质量直接决定后续PPO的效果。理解其内部机制有助于更好地调试和优化对齐训练流程。参考资源DeepSpeed-Chat官方仓库InstructGPT论文 - RLHF理论基础Qwen3技术报告 - 模型架构细节本文基于实际训练经验撰写训练环境RTX 4090 24GB/48GBDeepSpeed 0.16.0PyTorch 2.1