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 参数区分：

class PromptDataset(Dataset):
"""
三阶段数据加载器
train_phase:
1 – SFT阶段：只返回chosen response
2 – RM阶段：返回chosen + rejected成对数据
3 – RL阶段：只返回prompt用于生成
"""
def __getitem__(self, idx):
if self.train_phase == 2:
# RM关键：同时返回chosen和reject的pair
return (
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阶段数据处理流程

def create_dataset_split(..., train_phase=2):
elif train_phase == 2:
# 关键：同时处理chosen和reject，形成对比对
for tmp_data in current_dataset:
chosen_sentence = raw_dataset.get_prompt_and_chosen(tmp_data) # 偏好回答
reject_sentence = raw_dataset.get_prompt_and_rejected(tmp_data) # 非偏好回答

if chosen_sentence and reject_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组装

class DataCollatorReward:
def __call__(self, data):
batch = {}
# 关键：将chosen和rejected在batch维度拼接
# 前一半是chosen，后一半是rejected
batch["input_ids"] = torch.cat(
[f[0] for f in data] + [f[2] for f in data], dim=0
)
batch["attention_mask"] = torch.cat(
[f[1] for f in data] + [f[3] for f in data], dim=0
)
return batch

设计意图：[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核心实现

class RewardModel(nn.Module):
def __init__(self, base_model, tokenizer, ...):
super().__init__()
self.rwtransformer = base_model # 主干网络（无LM Head）

# 关键：添加v_head，将hidden_size映射到1维分数
if hasattr(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函数）

def forward(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和rejected
bs = input_ids.shape[0] // 2 # 实际batch大小
chosen_rewards = rewards[:bs] # 前一半是chosen
rejected_rewards = rewards[bs:] # 后一半是rejected

5.3 关键：对齐片段的确定

# 对于每个样本对，找到"有效对比区域"
for i in range(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))

if self.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使用）

def forward_value(self, ..., return_value_only=False, prompt_length=0):
# 用于Stage 3 PPO阶段，给单个response打分
values = self.v_head(hidden_states).squeeze(–1)

if return_value_only:
return values # 所有位置分数，用于GAE计算

# 找到response最后一个有效token的分数作为奖励
for i in range(bs):
# 从prompt之后找第一个PAD
c_inds = (input_id[prompt_length:] == self.PAD_ID).nonzero()
c_ind = c_inds[0].item() + prompt_length if len(c_inds) > 0 else seq_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/bash
deepspeed –num_gpus 1 main.py \\
–model_name_or_path /root/vscode/Tuling/Qwen3-0.6B \\
–data_path /path/to/train.jsonl \\
–data_split "6,2,2" \\
–num_train_epochs 1 \\
–gradient_accumulation_steps 4 \\
–per_device_train_batch_size 1 \\ # 实际处理2个response（chosen+rejected）
–per_device_eval_batch_size 1 \\
–dropout 0.0 \\ # 关键：保持排序一致性
–max_seq_len 512 \\
–learning_rate 5e-5 \\
–zero_stage 2 \\
–dtype bf16 \\
–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倍提升

关键观察：

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_gpus 4 main.py \\
–model_name_or_path /root/vscode/Tuling/Qwen3-4B \\
... # 其他参数同单卡

多卡优势：

• 更大的有效batch size（4卡 × 1 × 4梯度累积 = 16）
• 更快的训练速度
• 支持更大模型（4B参数需多卡才能训练）

九、关键经验与最佳实践

9.1 训练稳定性

9.2 评估指标解读

def evaluation_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_format
from dschat.utils.utils import save_hf_format_safetensors

if args.global_rank == 0:
save_hf_format_safetensors(rm_model, tokenizer, args)

十、总结

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

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

参考资源：

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

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