RewardModel 这块阵地，最近因为 DeepSpeed 和 Qwen3 的组合又热闹起来了。把 Qwen3 这类本身就具备很强语义理解能力的基座模型，通过 DeepSpeed 做分布式训练，微调成奖励模型，是我正在跑的一条路线。这篇文章把我从零开始构建奖励模型的完整过程写出来，覆盖数据处理、损失函数、DeepSpeed 配置、训练监控和踩坑记录。适合正在做 RLHF、偏好排序，或者想给生成模型加打分的同学参考。不需要你有过模型训练经验，但最好能看懂基本的 PyTorch 张量操作。
很多新人上来就去查训练参数，其实最容易被忽略的是“奖励模型到底在学什么”。我会先把设计思路讲清楚，再给可直接复制的实操配置。整个过程跑通之后，你手里会有一个能给任意文本序列输出标量分数的模型，这个分数可以直接喂给强化学习，也可以单独做成偏好排序接口。
1. 项目整体设计与方案选型
1.1 奖励模型到底在解决什么问题
奖励模型本质上是给一个“文本序列”打分的回归模型。输入通常是 user prompt 和 assistant response 拼起来的完整文本，输出一个连续标量分数。这个分数在 RLHF 里充当裁判：同一个 prompt 下，让模型知道哪个回答更符合人类偏好。
它不是分类器，不需要给句子贴“好/坏”标签。大多数时候，我们只需要一个相对分数，而不是绝对分数。也就是说，模型可以给所有回答打 1.2 分、1.5 分，只要好的那个分数更高就行。真实训练中用的是偏好对数据：一条 prompt 对应两个回答，一个被人类更偏好（chosen），一个没那么偏好（rejected）。模型要学的是把 chosen 的打分压过 rejected。
所以损失函数自然就用 pairwise ranking loss，最常用的是 Bradley-Terry 模型导出的 log-sigmoid 损失：
loss = -log(sigmoid(reward_chosen - reward_rejected))
这个公式的意思是，chosen 和 rejected 的分数差越大，loss 越小。训练过程不关心绝对值，只关心相对排序。这一点会在后续监控里反复提到，很多新手只盯着 loss 下降，却忽略了分数分布正在异常漂移。
1.2 为什么选 Qwen3 作为底座，DeepSpeed 作为训练框架
Qwen3 在对话理解、指令跟随和多语言能力上都很稳。奖励模型虽然只需要最终打一个分，但底层模型必须先理解prompt和response的语义关系，否则分数没有依据。如果用太小的模型，比如几亿参数，可能连基础的指令约束都跟不上，偏好排序也就容易退化成长度、标点、格式等表面特征。
选 Qwen3 还有一个现实原因：它与当前主流训练框架的兼容性很好，权重加载、hidden states 获取、tokenizer 处理都很直接，不需要自己写一堆适配层。
DeepSpeed 解决的则是显存和训练效率问题。奖励模型通常从 7B、14B 甚至更大的底座训练，单卡全参微调几乎不可能。DeepSpeed 的 ZeRO 优化能把优化器状态、梯度、模型参数分布到多卡上，极大降低单卡显存压力。而且它直接把多卡同步、梯度裁剪、混合精度训练这些脏活封装好了，我可以更专注在模型逻辑上。
当然也可以不用 DeepSpeed，用原生 DDP 或者 ZeRO Offload 手动实现，但工程量和踩坑成本会高不少。对于“从零构建奖励模型”这个目标，DeepSpeed 是我目前觉得性价比最高的选择。
1.3 数据形态与得分头的接入位置
奖励模型的训练数据长这样：
{ "prompt": "请给我推荐一款适合办公的机械键盘", "chosen": "办公优先建议选茶轴或静音红轴，键程适中，手感不吵。推荐...", "rejected": "随便买一个就行了，能打字就行。" }chosen 和 rejected 会构造为两个独立序列，分别送给模型，得到两个 reward 分数。
关键点在于得分头（reward head）接在哪里。Qwen3 是 decoder-only 结构，最后几层隐藏状态已经编码了整个序列的语义。我们把最后一层 hidden states 拉出来，每个 token 位置都会有一个 hidden_size 维向量，然后只取序列最后一个非 padding token 的向量，经过一个线性层映射到 1 维标量。
为什么不取所有 token 的平均？这样做容易让模型把分数平滑掉，丢失关键位置的信息。大模型在很多任务里都是靠最后 token 作为整个序列的 summary，奖励模型沿用这个习惯也很自然。
2. 环境准备与数据预处理
2.1 依赖版本选择
训练奖励模型的前置环境并不复杂，但版本搭配不对会踩各种隐蔽的坑。我自己实测过的组合如下：
| 组件 | 推荐版本 | 说明 |
|---|---|---|
| Python | 3.10 | 对 DeepSpeed 兼容性最好 |
| PyTorch | 2.1 ~ 2.3 | 支持 bf16、FlashAttention |
| DeepSpeed | 0.14 以上 | ZeRO、CPUAdam、bf16 支持完整 |
| Transformers | 4.43 以上 | 适配 Qwen3 权重 |
| Accelerate | 0.30 以上 | 部分启动流程会用到 |
安装直接用 pip：
pip install deepspeed==0.14.4 transformers accelerate torch注意不要随手把 transformers 升到最新版，有些中间版本和 DeepSpeed 的 ZeRO-3 兼容性有问题，报错信息还不直观。如果已经装了，出问题就按上面这组版本回退。
2.2 偏好数据集的整理与 tokenization
数据质量直接决定奖励模型上限。我的处理流程是：
- 按 prompt 字段 group by，一个 prompt 下可能有多条偏好对。
- 过滤掉 chosen 和 rejected 完全一样的样本。
- 过滤掉长度过长、明显语无伦次的样本。
- 对每个独立的 prompt 记录它对应的所有 pair。
然后做 tokenization。这里有一个容易犯的错误：直接把 prompt 和 response 用普通文本拼接。Qwen3 有自己约定好的 chat template，包含角色标识 token，如果训练和后续推理时的格式不一致，打分会漂移。
我建议先用apply_chat_template把 prompt 和 chosen/rejected 分别构造成完整文本，再 tokenize：
from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained("qwen3-7b", trust_remote_code=True) def build_text(prompt, response): messages = [ {"role": "user", "content": prompt}, {"role": "assistant", "content": response} ] return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)然后对两个完整文本分别做 tokenize，加上 truncation 和 padding：
max_len = 2048 def tokenize_pair(example): chosen_text = build_text(example["prompt"], example["chosen"]) rejected_text = build_text(example["prompt"], example["rejected"]) chosen_enc = tokenizer(chosen_text, truncation=True, max_length=max_len) rejected_enc = tokenizer(rejected_text, truncation=True, max_length=max_len) return { "chosen_input_ids": chosen_enc["input_ids"], "chosen_attention_mask": chosen_enc["attention_mask"], "rejected_input_ids": rejected_enc["input_ids"], "rejected_attention_mask": rejected_enc["attention_mask"] }padding 策略我建议用padding="longest"，让同一个 batch 内对齐到最长样例，而不是固定 max_len。固定长度会浪费大量显存，尤其当多数样本都只有几百 token 时。训练时的动态 padding 交给 DataLoader 里的collate_fn统一做。
2.3 训练集/验证集拆分原则
奖励模型的验证集拆分和普通分类任务不一样。必须按 prompt 维度拆，不能按 sample 随机拆。
举个例子：如果同一条 prompt 的 10 对数据里，8 对进了 train，2 对进了 eval，那么 eval 里出现的 prompt 在 train 里已经被模型见过了。模型可能记住了这个 prompt 对应回答的打分模式，验证集准确率虚高。奖励模型要评估的是“对没见过的问题能正确排序”，而不是“对见过的 prompt 记住排序”。所以我的做法是：
unique_prompts = train_data.unique_prompts train_prompts = set(random.sample(unique_prompts, int(len(unique_prompts) * 0.9))) eval_prompts = set(unique_prompts) - train_prompts train_set = data[data.prompt.isin(train_prompts)] eval_set = data[data.prompt.isin(eval_prompts)]在实际项目中，一个 prompt 下往往有多个 chosen/rejected 对，拆分时要把这些 pair 看成一个整体移动，不要打散。
3. DeepSpeed 训练奖励模型的核心实现
3.1 加载 Qwen3 并替换为回归头
直接用AutoModel加载 Qwen3 的隐藏表示部分，不需要加载 lm_head。这样既省内存，又避免模型误用语言建模头。
核心模型结构如下：
import torch import torch.nn as nn from transformers import AutoModel, AutoConfig class RewardModel(nn.Module): def __init__(self, model_path, hidden_size=3584): super().__init__() config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) self.base_model = AutoModel.from_pretrained( model_path, config=config, torch_dtype=torch.bfloat16, trust_remote_code=True, use_cache=False, device_map=None ) self.reward_head = nn.Linear(hidden_size, 1, bias=False) self._init_head() def _init_head(self): nn.init.zeros_(self.reward_head.weight) def forward(self, input_ids, attention_mask): outputs = self.base_model( input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True ) hidden_states = outputs.hidden_states[-1] # [B, S, H] last_index = attention_mask.sum(dim=1) - 1 # [B] batch_index = torch.arange(input_ids.size(0), device=input_ids.device) last_hidden = hidden_states[batch_index, last_index, :] # [B, H] reward = self.reward_head(last_hidden).squeeze(-1) # [B] return rewardhidden_size 要根据具体 Qwen3 规格调整。7B 级别一般是 3584，14B 级别可能是 5120，14B 一般 5120。最稳的办法是从config.hidden_size读取。
初始化 reward_head 权重为 0 是我习惯的做法。这样初始阶段所有输入的 reward 都是 0，chosen 和 rejected 的差也是 0，不会刚开始训练就给模型一个错误的方向。
另外注意这里没有device_map="auto"。 DeepSpeed 对模型的管理和 transformers 的 device_map 容易冲突，尤其 ZeRO-3 会把参数切分到多卡，如果 transformers 已经把模型分好卡了，后面会报出奇怪的 place 错误。
3.2 损失函数与训练循环
训练循环比较直接。每步需要把 chosen 和 rejected 分别过一遍模型得到两个 reward 向量，再算差值。完整片段：
import torch.nn.functional as F def compute_loss(model, chosen, rejected): chosen_rew = model(chosen["input_ids"], chosen["attention_mask"]) rejected_rew = model(rejected["input_ids"], rejected["attention_mask"]) logits = chosen_rew - rejected_rew loss = -F.logsigmoid(logits).mean() with torch.no_grad(): acc = (logits > 0).float().mean().item() diff = logits.mean().item() return loss, acc, diff这个logits就是 chosen 和 rejected 的分数差。如果差大于 0，说明模型判断正确。
我还会加一个很小的 L2 正则，防止 reward 分数无边界漂移：
if reward_l2_reg > 0: reg = (chosen_rew.pow(2) + rejected_rew.pow(2)).mean() loss = loss + reward_l2_reg * regreward_l2_reg 我通常设为 0.001 ~ 0.01。这个正则不是为了优化精度，而是为了让分数保持在一个合理区间。后期接强化学习的时候，reward 分过大会导致 KL 项和 reward 比例失衡。
优化器和学习率建议直接用 DeepSpeed 配置里的参数，不单独写在训练脚本里。逻辑保持简单：
for step, batch in enumerate(dataloader): chosen = {"input_ids": batch["chosen_input_ids"].cuda(), "attention_mask": batch["chosen_attention_mask"].cuda()} rejected = {"input_ids": batch["rejected_input_ids"].cuda(), "attention_mask": batch["rejected_attention_mask"].cuda()} loss, acc, diff = compute_loss(model, chosen, rejected) loss.backward() model_engine.step() model_engine.zero_grad() if step % 20 == 0: print(f"step {step}, loss {loss.item():.4f}, acc {acc:.4f}, diff {diff:.4f}")这里model_engine是 deepspeed 初始化的 engine，不是原始 model。初始化方式看 3.3。
3.3 DeepSpeed ZeRO 配置与启动命令
DeepSpeed 的配置文件是训练的核心之一。我用得最多的配置如下：
{ "train_batch_size": 32, "gradient_accumulation_steps": 4, "gradient_clipping": 1.0, "zero_optimization": { "stage": 2, "offload_optimizer": { "device": "cpu" }, "contiguous_gradients": true }, "bf16": { "enabled": true }, "steps_per_print": 20 }train_batch_size代表全局 batch size，它必须等于 micro batch size × GPU 数 × gradient_accumulation_steps。比如每个 GPU 上实际跑 4 条样本，4 张卡，累计 4 步，那么全局 batch size 就是 4 × 4 × 4 = 64。如果这里配置错了，DeepSpeed 会直接报错提示。
训练脚本初始化 engine 的部分：
import deepspeed ds_config = json.load(open("ds_config.json")) model_engine, optimizer, _, scheduler = deepspeed.initialize( model=model, model_parameters=filter(lambda p: p.requires_grad, model.parameters()), config_params=ds_config )启动命令：
deepspeed --num_gpus=4 train_reward_model.py --deepspeed_config ds_config.jsonZeRO stage 的选择。对于 7B 级别全参微调，stage 2 基本够用。如果显存还是紧，可以开 stage 3，让模型参数也切分到多卡。stage 3 的通信开销更大，训练速度会慢一些，所以能 stage 2 就不要盲目上 stage 3。CPU offload 优化器状态能进一步压显存，但会明显拖慢迭代速度，建议只在 batch size 提不上去的时候打开。
4. 训练过程中的监控与调参
4.1 关键指标怎么看
训练奖励模型不是只看 loss。我平时重点监控四个指标：
| 指标 | 计算方式 | 健康参考范围 |
|---|---|---|
| pairwise accuracy | chosen_rew > rejected_rew 的比例 | 0.6 ~ 0.85 |
| loss | -log(sigmoid(chosen_rew - rejected_rew)) | 稳定下降即可 |
| reward diff | chosen_rew - rejected_rew 的均值 | 1.0 ~ 3.0 |
| reward std | 所有样本 reward 的标准差 | 0.5 ~ 2.5 |
acc 超过 0.85 不一定说明模型好，可能是数据里有明显偏离的样本被吃进去了。reward diff 太小则说明模型区分度不足，后续做强化学习时模型很难根据微小分数差优化。reward std 是用来防漂移的，如果标准差持续涨到 3 以上，就要考虑加正则或降低学习率。
4.2 超参调整经验
学习率：全参数微调时，我推荐1e-5起步；如果是 14B 更大底座，降到5e-6。奖励模型的数据量一般不是海量，大部分场景在几十万条以内，学习率太大很容易让底座原有的语义表征被破坏。
训练轮数：1 到 2 个 epoch 就够。偏好数据相比预训练语料非常少，模型反复看同一个 pair，很容易把排序“背下来”，而不是学会泛化。如果你发现 eval acc 不再上升甚至下降，train acc 还在继续涨，那就是典型的过拟合信号。
Batch size：微调阶段 32 到 64 是比较稳的区间。batch 太小会让 pairwise 排序的目标函数方差很大，reward diff 起起伏伏，不好收敛。
数据截断长度：7B 模型推荐 2048，如果显存允许可以用 4096。但不建议为了强行覆盖长文本把所有样本 padding 到 4096。很多关于回答质量的信号其实集中在前几百个 token，过长序列反而把训练重点带偏。
4.3 奖励分数漂移的应对方案
我踩过最明显的一个坑：训练到一半，模型对几乎所有输入都打出越来越高的 reward，chosen 和 rejected 的分差却保持不变。这种情况通常是 reward head 权重在持续增大，底座的 hidden states 也在朝偏大的方向更新。
我处理漂移有两个办法：
- 给 reward head 的权重加 L2 约束，前面代码里的
reward_l2_reg干的就是这个。 - 定期打印 train 和 eval 的 reward 均值，如果两边同时持续抬升，就降低学习率，或者把 L2 系数调大。
不要指望模型自己把 reward 控制在合理区间，因为损失函数只关心排序，不关心分数绝对值。监管分数分布是训练者的责任。
5. 踩坑记录与问题排查实录
5.1 显存不足的解决顺序
训练 7B 奖励模型最常见的痛点是显存爆掉。我的排查顺序固定是这样：
| 处理手段 | 作用 | 代价 |
|---|---|---|
| 开启 gradient checkpointing | 减少激活显存 | 训练速度下降 20%~30% |
| 从 ZeRO stage 2 切到 stage 3 | 参数、梯度、优化器状态全切分 | 通信量变大 |
| 开启 CPU offload optimizer | 优化器状态放内存 | CPU 瓶颈，迭代变慢 |
| 降低 max length | 减少激活值 | 长文本信息丢失 |
| 减小 per_gpu micro batch | 最直接 | batch 太小影响稳定性 |
这五招按顺序从“影响小”到“影响大”尝试。很多时候开一个 gradient checkpointing 就够了，不用急着上 stage 3。
5.2 运行报错排查
训练过程中常见的报错和对应解法：
| 现象 | 大概率原因 | 对策 |
|---|---|---|
| DeepSpeed 初始化时报 train_batch_size mismatch | 配置里 train_batch_size 和 micro_batch × GPU数 × 累积步数不匹配 | 重新计算并修正 |
| 加载模型后 forward 报 device mismatch | 模型被 transformers 的 device_map 手动分配过，和 DeepSpeed 产生冲突 | 移除device_map，交给 DeepSpeed 管理 |
| 多卡训练出现 reward 指标不一致 | 日志里打印了 rank 0 的局部指标，没有聚合 | 打印平均指标时用分布式 allreduce |
| 训练中途出现 NaN | 学习率过高、head 初始化异常、数据有空 id | 调低 lr，确认 head 权重为零初始化，检查分词结果 |
| 验证时 reward 和训练时差很多 | tokenization 不一致，或 padding 方式变了 | 确保验证数据走同一个 tokenize 函数 |
其中一个容易被忽略的坑是：加载 Qwen3 时如果用了AutoModelForCausalLM，会把巨大的 lm_head 也加载进显存。奖励模型完全用不到这个 lm_head，白白占了几个 GB。改用AutoModel或者显式把 lm_head 丢掉，是很好的降显存技巧。
5.3 数据泄漏与质量盲点
数据问题比框架问题更容易让奖励模型“看着正常，实际很糟”。
位置偏置是典型问题。如果训练数据构造时总是 chosen 在前、rejected 在后，模型可能学到“这段文本在完整序列里靠前，分数高”，而不是真正理解内容。我一开始吃过这个亏，后来构造 pair 时会随机交换 chosen 和 rejected 的位置，并把交换标记存下来，loss 计算时根据标记重新对齐。
另一个盲点是重复内容。我遇到过一个数据集里大量回答只是模板加参数替换，比如“以下是推荐：A 产品，理由：...”这种结构。奖励模型很快就会发现模板特征和排序的关系，导致真实场景里换个说法就不会打分。
所以现在我会在清洗阶段做文本去重，并抽样可视化 reward 和文本长度的关系。如果发现 reward 显著偏向长回答，就要临时把长度作为一个特征做正则，或者在数据里手动平衡长中短回答的比例。
6. 部署与后续扩展
6.1 将训练好的奖励模型导出与推理
训练完成后的模型和普通模型一样，可以用save_pretrained把权重存下来。保存时我会单独保存 reward head 的权重，或者直接把整个 RewardModel 包装成 transformers 的 PreTrainedModel 格式。
推理时最关键的是和训练保持一致的 tokenize 方式。我用 bf16 加载模型，推理脚本很短：
import torch def score_pairs(prompt, response_list): texts = [build_text(prompt, r) for r in response_list] encodings = tokenizer(texts, padding="longest", truncation=True, max_length=2048, return_tensors="pt") encodings = {k: v.cuda() for k, v in encodings.items()} with torch.no_grad(): rewards = model(encodings["input_ids"], encodings["attention_mask"]).float() return rewards.tolist()输出可以接 Python 的排序服务，也可以把分数写入日志做离线分析。注意推理时不要把 padding_side 改掉，训练时右 padding，推理时改成左 padding，可能会导致 attention mask 定位的“最后一个非 padding token”不再准确。如果要换，两边必须同时换。
6.2 后续工程化扩展
奖励模型最有价值的场景是接强化学习。在 PPO 里，模型生成一个 response 后，奖励模型的输出和 KL penalty 结合起来：
final_reward = reward_model_score - beta * kl_penalty
这里的 beta 要先去小规模试跑几轮才能定，reward 分布太集中时，beta 稍微调大一点就会让模型变得极其保守。我建议先把奖励模型的输出保存一批样本，统计 mean 和 std，再设置 KL 惩罚的权重。
如果预算允许，还可以训练多个初始化不同的奖励模型做集成。单个奖励模型难免有偏向，几个模型对同一文本的分数做平均，能平滑掉一部分随机性和偏好噪声。这个技巧在数据量少的时候尤其有用。
最后再分享一点个人体会
跑完整个流程，我最大的感觉是：训练奖励模型的瓶颈不在 DeepSpeed，也不在 Qwen3，而在数据准备和验收标准。DeepSpeed 只是把显存和并行问题解决得很干净，真正决定模型质量的是我们交给它的 pair 样本是否一致、评测集是否干净。我现在每接手一批新数据，都会先画一张 chosen 和 rejected 的 reward 分布图，观察两者分界是否清晰，而不是只盯着 accuracy。这个习惯帮我抓住了很多隐藏的数据问题。如果你也在搭类似的奖励模型，我建议把“可视化 reward 分布”加入你的必做清单，它值得花这十几分钟。
