DeepSpeed-RewardModel-Qwen3 实战：从零构建奖励模型

一、前言：为什么需要Reward Model？

在大语言模型的对齐训练中，**Reward Model（奖励模型）**是RLHF（Reinforcement 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的两个不同Response（chosen 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_phase==2:# 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_phase=2):eliftrain_phase==2:# 关键：同时处理chosen和reject，形成对比对fortmp_dataincurrent_dataset:chosen_sentence=raw_dataset.get_prompt_and_chosen(tmp_data)# 偏好回答reject_sentence=raw_dataset.get_prompt_and_rejected(tmp_data)# 非偏好回答ifchosen_sentenceandreject_sentence:# 都加上EOS token，让模型学会何时结束chosen_token=tokenizer(chosen_sentence+eos_token,...)reject_token=tokenizer(reject_sentence+eos_token,...)chosen_dataset.append(chosen_token)reject_dataset.append(reject_token)

关键细节：

• 成对处理：每个样本必须同时包含chosen和rejected，且基于同一个prompt
• EOS标记：添加<|endoftext|>教会模型何时停止生成
• 长度对齐：通过padding="max_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],dim=0)batch["attention_mask"]=torch.cat([f[1]forfindata]+[f[3]forfindata],dim=0)returnbatch

设计意图：[chosen_batch; rejected_batch]的拼接方式，让模型在一个前向传播中同时处理两个response，高效计算pairwise loss。

四、模型架构：从Causal LM到Reward Model

4.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.rwtransformer=base_model# 主干网络（无LM Head）# 关键：添加v_head，将hidden_size映射到1维分数ifhasattr(self.config,"word_embed_proj_dim"):# OPT模型特殊处理self.v_head=nn.Linear(self.config.word_embed_proj_dim,1,bias=False)else:# GPT/LLaMA/Qwen等模型self.config.n_embd=self.config.hidden_size self.v_head=nn.Linear(self.config.n_embd,1,bias=False)

架构特点：

• 轻量级：只增加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_outputs=self.rwtransformer(input_ids,...)hidden_states=transformer_outputs[0]# Step 2: v_head为每个位置打分# shape: (batch_size * 2, seq_len)rewards=self.v_head(hidden_states).squeeze(-1)# Step 3: 拆分为chosen和rejectedbs=input_ids.shape[0]//2# 实际batch大小chosen_rewards=rewards[:bs]# 前一半是chosenrejected_rewards=rewards[bs:]# 后一半是rejected

5.3 关键：对齐片段的确定

# 对于每个样本对，找到"有效对比区域"foriinrange(bs):chosen_id=chosen_ids[i]rejected_id=rejected_ids[i]# 找到chosen和rejected第一个不同的位置（response开始分叉）check_divergence=(chosen_id!=rejected_id).nonzero()divergence_ind=check_divergence[0].item()# 找到各自的有效结束位置（第一个padding token）c_ind=(chosen_id==PAD_ID).nonzero()[0].item()r_ind=(rejected_id==PAD_ID).nonzero()[0].item()end_ind=max(c_ind,r_ind)# 取较长的# 截取对齐片段：从分叉点到结束点c_truncated=chosen_reward[divergence_ind:end_ind]r_truncated=rejected_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_truncated=c_truncated.float()r_truncated=r_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_{i=1}^N \log\sigma(r_\theta(x, y_c) - r_\theta(x, y_r))L=−N1​i=1∑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_only=False,prompt_length=0):# 用于Stage 3 PPO阶段，给单个response打分values=self.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_ind=c_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_split"6,2,2"\--num_train_epochs1\--gradient_accumulation_steps4\--per_device_train_batch_size1\# 实际处理2个response（chosen+rejected）--per_device_eval_batch_size1\--dropout0.0\# 关键：保持排序一致性--max_seq_len512\--learning_rate5e-5\--zero_stage2\--dtypebf16\--output_dir./output

关键参数说明：

6.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倍提升

关键观察：

1. 绝对分数为负是正常的：模型学习的是相对排序，绝对值无意义
2. 差距扩大才是核心：从0.04扩大到49.69，说明模型学会了强区分
3. 准确率99%：模型已充分学习训练集偏好

6.3 显存与耗时

七、模型对比：Reward Model vs 原模型

7.1 文件大小对比

7.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,936
• Reward Model词表：151,672（少了264个token）
• 缺失的主要是：连续空格组合、编码错误字符等无意义token

7.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 size（4卡 × 1 × 4梯度累积 = 16）
• 更快的训练速度
• 支持更大模型（4B参数需多卡才能训练）

九、关键经验与最佳实践

9.1 训练稳定性

9.2 评估指标解读

defevaluation_reward(model,dataloader):# 核心指标：accuracy = (chosen_score > rejected_score)的比例correct=(chosen>rejected).sum()acc=correct/total# 辅助指标：分数差距diff=chosen_mean-rejected_mean# 越大说明区分度越强

9.3 保存格式优化

针对Qwen3模型，建议保存为safetensors格式：

# 替换原有的save_hf_formatfromdschat.utils.utilsimportsave_hf_format_safetensorsifargs.global_rank==0:save_hf_format_safetensors(rm_model,tokenizer,args)

十、总结

本文详细解析了DeepSpeed-Chat框架下Reward Model的完整训练流程：

1. 数据处理：成对构建chosen-rejected对比数据，通过DataCollatorReward拼接batch
2. 模型架构：AutoModel + v_head，轻量级扩展实现评分功能
3. 训练目标：Pairwise Ranking Loss学习相对排序，而非绝对分数
4. 工程实践：ZeRO优化、bf16混合精度、safetensors格式保存

Reward Model作为RLHF的"裁判"，其训练质量直接决定后续PPO的效果。理解其内部机制，有助于更好地调试和优化对齐训练流程。

参考资源：

• DeepSpeed-Chat官方仓库
• InstructGPT论文 - RLHF理论基础
• Qwen3技术报告 - 模型架构细节

本文基于实际训练经验撰写，训练环境：RTX 4090 24GB/48GB，DeepSpeed 0.16.0，PyTorch 2.1