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 依赖版本选择训练奖励模型的前置环境并不复杂但版本搭配不对会踩各种隐蔽的坑。我自己实测过的组合如下组件推荐版本说明Python3.10对 DeepSpeed 兼容性最好PyTorch2.1 ~ 2.3支持 bf16、FlashAttentionDeepSpeed0.14 以上ZeRO、CPUAdam、bf16 支持完整Transformers4.43 以上适配 Qwen3 权重Accelerate0.30 以上部分启动流程会用到安装直接用 pippip install deepspeed0.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 分别构造成完整文本再 tokenizefrom transformers import AutoTokenizer tokenizer AutoTokenizer.from_pretrained(qwen3-7b, trust_remote_codeTrue) def build_text(prompt, response): messages [ {role: user, content: prompt}, {role: assistant, content: response} ] return tokenizer.apply_chat_template(messages, tokenizeFalse, add_generation_promptFalse)然后对两个完整文本分别做 tokenize加上 truncation 和 paddingmax_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, truncationTrue, max_lengthmax_len) rejected_enc tokenizer(rejected_text, truncationTrue, max_lengthmax_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 策略我建议用paddinglongest让同一个 batch 内对齐到最长样例而不是固定 max_len。固定长度会浪费大量显存尤其当多数样本都只有几百 token 时。训练时的动态 padding 交给 DataLoader 里的collate_fn统一做。2.3 训练集/验证集拆分原则奖励模型的验证集拆分和普通分类任务不一样。必须按 prompt 维度拆不能按 sample 随机拆。举个例子如果同一条 prompt 的 10 对数据里8 对进了 train2 对进了 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_size3584): super().__init__() config AutoConfig.from_pretrained(model_path, trust_remote_codeTrue) self.base_model AutoModel.from_pretrained( model_path, configconfig, torch_dtypetorch.bfloat16, trust_remote_codeTrue, use_cacheFalse, device_mapNone ) self.reward_head nn.Linear(hidden_size, 1, biasFalse) 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_idsinput_ids, attention_maskattention_mask, output_hidden_statesTrue ) hidden_states outputs.hidden_states[-1] # [B, S, H] last_index attention_mask.sum(dim1) - 1 # [B] batch_index torch.arange(input_ids.size(0), deviceinput_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 级别一般是 358414B 级别可能是 512014B 一般 5120。最稳的办法是从config.hidden_size读取。初始化 reward_head 权重为 0 是我习惯的做法。这样初始阶段所有输入的 reward 都是 0chosen 和 rejected 的差也是 0不会刚开始训练就给模型一个错误的方向。另外注意这里没有device_mapauto。 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(fstep {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( modelmodel, model_parametersfilter(lambda p: p.requires_grad, model.parameters()), config_paramsds_config )启动命令deepspeed --num_gpus4 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 accuracychosen_rew rejected_rew 的比例0.6 ~ 0.85loss-log(sigmoid(chosen_rew - rejected_rew))稳定下降即可reward diffchosen_rew - rejected_rew 的均值1.0 ~ 3.0reward std所有样本 reward 的标准差0.5 ~ 2.5acc 超过 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 奖励分数漂移的应对方案我踩过最明显的一个坑训练到一半模型对几乎所有输入都打出越来越高的 rewardchosen 和 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, paddinglongest, truncationTrue, max_length2048, return_tensorspt) 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 分布”加入你的必做清单它值得花这十几分钟。
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 依赖版本选择训练奖励模型的前置环境并不复杂但版本搭配不对会踩各种隐蔽的坑。我自己实测过的组合如下组件推荐版本说明Python3.10对 DeepSpeed 兼容性最好PyTorch2.1 ~ 2.3支持 bf16、FlashAttentionDeepSpeed0.14 以上ZeRO、CPUAdam、bf16 支持完整Transformers4.43 以上适配 Qwen3 权重Accelerate0.30 以上部分启动流程会用到安装直接用 pippip install deepspeed0.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 分别构造成完整文本再 tokenizefrom transformers import AutoTokenizer tokenizer AutoTokenizer.from_pretrained(qwen3-7b, trust_remote_codeTrue) def build_text(prompt, response): messages [ {role: user, content: prompt}, {role: assistant, content: response} ] return tokenizer.apply_chat_template(messages, tokenizeFalse, add_generation_promptFalse)然后对两个完整文本分别做 tokenize加上 truncation 和 paddingmax_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, truncationTrue, max_lengthmax_len) rejected_enc tokenizer(rejected_text, truncationTrue, max_lengthmax_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 策略我建议用paddinglongest让同一个 batch 内对齐到最长样例而不是固定 max_len。固定长度会浪费大量显存尤其当多数样本都只有几百 token 时。训练时的动态 padding 交给 DataLoader 里的collate_fn统一做。2.3 训练集/验证集拆分原则奖励模型的验证集拆分和普通分类任务不一样。必须按 prompt 维度拆不能按 sample 随机拆。举个例子如果同一条 prompt 的 10 对数据里8 对进了 train2 对进了 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_size3584): super().__init__() config AutoConfig.from_pretrained(model_path, trust_remote_codeTrue) self.base_model AutoModel.from_pretrained( model_path, configconfig, torch_dtypetorch.bfloat16, trust_remote_codeTrue, use_cacheFalse, device_mapNone ) self.reward_head nn.Linear(hidden_size, 1, biasFalse) 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_idsinput_ids, attention_maskattention_mask, output_hidden_statesTrue ) hidden_states outputs.hidden_states[-1] # [B, S, H] last_index attention_mask.sum(dim1) - 1 # [B] batch_index torch.arange(input_ids.size(0), deviceinput_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 级别一般是 358414B 级别可能是 512014B 一般 5120。最稳的办法是从config.hidden_size读取。初始化 reward_head 权重为 0 是我习惯的做法。这样初始阶段所有输入的 reward 都是 0chosen 和 rejected 的差也是 0不会刚开始训练就给模型一个错误的方向。另外注意这里没有device_mapauto。 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(fstep {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( modelmodel, model_parametersfilter(lambda p: p.requires_grad, model.parameters()), config_paramsds_config )启动命令deepspeed --num_gpus4 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 accuracychosen_rew rejected_rew 的比例0.6 ~ 0.85loss-log(sigmoid(chosen_rew - rejected_rew))稳定下降即可reward diffchosen_rew - rejected_rew 的均值1.0 ~ 3.0reward std所有样本 reward 的标准差0.5 ~ 2.5acc 超过 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 奖励分数漂移的应对方案我踩过最明显的一个坑训练到一半模型对几乎所有输入都打出越来越高的 rewardchosen 和 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, paddinglongest, truncationTrue, max_length2048, return_tensorspt) 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 分布”加入你的必做清单它值得花这十几分钟。
平方复杂度无法被常规优化消除｜TaoToken 统一 Key 通道实测)
