一、前言：为什么需要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])

数学推导：
$$\mathcal{L}=-\frac{1}{N}\sum _{i=1}^N\log \sigma (r_{\theta }(x,y_c)-r_{\theta }(x,y_r))$$
其中 $\sigma$ 是sigmoid函数，目标是让 $r_{\theta }(x,y_c)>r_{\theta }(x,y_r)$。

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

关键参数说明：

参数	值	说明
per_device_train_batch_size	1	实际处理1对（chosen+rejected）= 2个序列
dropout	0.0	必须设为0，避免训练/推理排序不一致
zero_stage	2	ZeRO-2优化，显存占用从5.55GB降至3.33GB

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 显存与耗时

指标	数值
峰值显存	17.4 GB / 24 GB (4090)
训练时间	23分44秒
处理速度	~7条/秒
训练集	10,000对
验证集	1,000对

---

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

7.1 文件大小对比

模型	大小	差异分析
Qwen3-0.6B	1.50 GB	含LM Head (151936×1024≈1.55亿参数)
Reward Model	1.19 GB	减小309MB (20.6%)，移除LM Head

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 训练稳定性

问题	解决方案
排序不一致	设置 dropout=0.0
分数发散	使用 compute_fp32_loss=True
显存不足	启用ZeRO-2/3 + 梯度检查点

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的完整训练流程：

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