准备训练一个自己的奖励模型我踩过不少坑。去年某段时间我负责一个对齐项目需要给生成模型组里的回复做质量排序最开始图省事直接调外部打分接口结果成本高、延迟大而且接口返回的分数分布和业务场景完全不匹配。后来决定自己训一个基于偏好数据的奖励模型技术栈选了 DeepSpeed 加 Qwen3 底座整个过程从数据构造、环境对齐到 Loss 收敛一路踩坑一路解决。这篇文章就把我从零构建 DeepSpeed 奖励模型的全过程拆开讲包括环境配置、数据细节、训练脚本里的关键逻辑以及部署时那些容易忽略的坑。1. 从一次失败的对齐实验说起为什么要单独训练一个奖励模型1.1 奖励模型在 RLHF 链路中的真实位置很多刚开始接触强化学习的人对奖励模型有个误解以为它就是给文本打个分或者以为直接用大模型当裁判就行。实际在 RLHF 链路里奖励模型是策略模型和对齐目标之间的翻译器它的任务是把人类偏好变成稳定的标量信号。强化学习阶段策略模型每生成一个回复奖励模型就要给这个回复打分这个分数直接决定策略梯度往哪个方向走。如果奖励模型不稳定或者有偏后面的一切都会跟着歪掉。用现成的大模型 API 当奖励模型确实省事但我在实际项目里遇到两个硬伤第一接口延迟不稳定强化学习阶段每轮要采集成千上万个样本一个样本来回几百毫秒训练节奏全被打乱。第二外部模型的打分分布和我们的业务数据分布不一致比如某个回答在我们场景里明显是安全但无用的外部打分可能给高分。这就是为什么很多团队最终都会走到自己训奖励模型这条路。自己训奖励模型的本质就是用人类标注好的偏好数据训练一个带回归头的语言模型让它学会哪个回复更好这件事。这个任务的样本格式通常是三元组提示、好回复、差回复。模型需要学会给好回复打高分给差回复打低分。这个任务看起来比生成任务简单但对数据质量和训练细节的要求非常高。1.2 为什么底座选 Qwen3底座模型的选择决定了奖励模型的上限。我用过几个候选最终选 Qwen3 主要出于三点考虑。第一是基础能力足够强。奖励模型不是从零学习语言而是在已有语言理解能力之上学习判断。Qwen3 的指令遵循能力和长文本处理能力在同量级底座里表现稳定尤其在中文数据上优势明显我们的业务数据恰好是中文为主。第二是开源生态完整。加载权重、做对话模板、微调脚本这些都有现成实现相比一些冷门底座省去大量造轮子的时间。更关键的是它在 Hugging Face 和 ModelScope 都有权重下载和部署都比较方便。第三是模型结构适合加回归头。Qwen3 的 hidden state 维度适中在顶部加一个线性层做 scalar 输出参数量增加很少训练和推理的开销都可控。我做了一个简单对比同样数据量下不同底座的奖励模型在验证集上的准确率表现底座参数量验证集准确率训练显存占用batch4某 1.5B 底座1.5B62.3%约 24GB某 3B 底座3B66.1%约 36GBQwen3-4B4B69.8%约 38GB这个结果不是绝对结论因为不同底座在不同数据分布上的表现差异很大但它至少说明 Qwen3 在我们的业务数据上具备足够优势。最终我选择了 4B 规模在效果和资源消耗之间取得了平衡。当然如果显存充足8B 底座的奖励模型上限会更高训练收敛也通常更平稳。2. 环境准备与依赖版本对齐最容易翻车的一步2.1 硬件选型与显存预估不少人一上来就卡在环境上。奖励模型本身不算重但由于训练需要同时加载底座、优化器状态和梯度显存压力比纯推理高一个量级。以 Qwen3-4B 为例底座权重大概占 8GBBF16AdamW 优化器状态需要额外 12GB 左右梯度还需要 8GB单卡 40GB 勉强能跑一个小 batch想跑得舒服一点还是得靠多卡。我在实际项目里用的是两张 24GB 显存的消费级显卡做数据并行配合 DeepSpeed 的 ZeRO Stage 2 把优化器状态切分到多卡。有人问为什么不用 ZeRO Stage 3我的经验是奖励模型训练通常 batch size 不需要特别大Stage 2 已经能把单卡显存占用压到可控范围而 Stage 3 的通信开销更大在小规模集群上反而拖慢速度。如果只有单卡 24GB 显存也可以跑但 batch size 只能压到 1配合梯度累积到 8 步来模拟更大的有效 batch。2.2 依赖安装的版本陷阱训练代码出现莫名其妙的报错十有八九是版本没对齐。我整理了一份经过验证的版本组合这组组合在我的环境里跑了完整训练流程没有报错依赖推荐版本说明Python3.10版本太低会导致某些新库不支持PyTorch2.1.2必须支持 BF16 自动混合精度Transformers4.44低版本可能不认识 Qwen3 的模型结构DeepSpeed0.14.40.15 以后接口有变化遇到报错先回退Accelerate0.33.0与 DeepSpeed 集成比较稳定FlashAttention2.5.8可选但能显著降低显存占用值得单独说的是 FlashAttention 的安装。组件本身能省不少显存但如果你的显卡是旧架构或者驱动版本偏低FlashAttention 编译会直接报错。我的做法是先在纯 PyTorch 环境下把训练流程跑通确认逻辑没问题之后再考虑加 FlashAttention 提速避免一上来就是环境问题根本分不清是代码错了还是环境错了。2.3 基于 DeepSpeed 的加载脚本骨架开始训练之前需要确认模型能正常加载并能输出 hidden state。一个常规做法是先把底座加载出来跑一次简单的前向传播确认 tokenzier 和模型结构匹配再进入正式的训练脚本。这个验证步骤看起来多余但能省下后面至少半小时的调试时间。import torch from transformers import AutoModelForCausalLM, AutoTokenizer model_name your_qwen3_model_path tokenizer AutoTokenizer.from_pretrained(model_name, trust_remote_codeTrue) model AutoModelForCausalLM.from_pretrained( model_name, torch_dtypetorch.bfloat16, device_mapauto, trust_remote_codeTrue, ) inputs tokenizer([这是一个测试], return_tensorspt) outputs model(**inputs, output_hidden_statesTrue) print(outputs.hidden_states[-1].shape)这一步能确认三件事tokenizer 的词汇表是否正常、模型能否自动把输入切分到多卡、hidden state 的形状是否符合预期。hidden state 的形状一般是 [batch_size, seq_len, hidden_dim]Qwen3-4B 的 hidden_dim 是 2560。这个数字后面会被用到因为回归头的输入维度就是它。3. 偏好数据的组织细节pairwise 格式、长度对齐与 tokenizer3.1 三元组数据的来源与质量过滤奖励模型的训练数据是偏好三元组形式是提示、好回复、差回复。这块数据的质量直接决定奖励模型的天花板数据构造上的偷懒后期无论怎么调 Loss 都补不回来。数据来源一般有三种渠道。第一种是人工标注找标注人员对同一个提示写两个回复再让标注员选择哪个更好。这种数据质量最高但成本也最高。第二种是线上日志把用户点击、点赞、采纳等行为作为隐式反馈这种方式量大但噪声大。第三种是先用一个强模型生成多个回复再用成对比较的方式筛选类似 AI-assisted 标注。第三种方式适合冷启动我一开始就是用这种方式构建了第一批训练数据。数据质量过滤是关键步骤。我在清洗时重点处理了三类坏样本一是两个回复长度差异极大比如一个只有 5 个 token另一个有 200 个 token这种样本会让模型学会长就是好而不是好就是好。二是两个回复内容高度重复模型学不到任何区分信号。三是两个回复都是明显的低质量回复比如全是重复的话这种样本即便标注了偏好对模型也没有正向帮助。最终有效数据占比通常在 60% 到 70% 之间剩下的大部分是长度失衡或者内容重复。3.2 训练样本的拼接策略奖励模型和生成模型对输入组织方式不同。生成模型只需要一个单向序列而奖励模型需要同时处理 prompt、chosen好回复和 rejected差回复。这里有两个常用策略。第一种是双序列策略把 prompt 和 chosen 拼接成一个序列prompt 和 rejected 拼接成另一个序列每个样本在 batch 里是两条独立数据最后把两个分数做差得到 loss。这种策略实现简单很多开源代码都用这种方式。第二种是单序列策略把 prompt、chosen、rejected 三部分拼成一个序列用特殊 token 分隔通过一个回归头输出两个分数。后者节省了重复计算 prompt 的开销在推理阶段更高效但训练实现更复杂。我在训练阶段用的是第一种方案推理阶段用第二种。如果你刚开始做建议直接用第一种代码实现最直观。3.3 Tokenizer 对齐与长度截断的细节数据预处理里最容易忽略的是 tokenizer 长度对齐问题。chosen 和 rejected 的长度往往不一致如果直接按固定长度截断可能导致一边已经被截没了另一边还很完整。我的处理方式是先分别 tokenize chosen 和 rejected然后设置最大长度为 2048超长部分优先保留尾部。这是因为模型的判断依据通常集中在回复的中后段开头部分主要是上下文。max_length 2048 def preprocess_pair(prompt, chosen, rejected): prompt_ids tokenizer.encode(prompt, add_special_tokensFalse) chosen_ids tokenizer.encode(chosen, add_special_tokensFalse) rejected_ids tokenizer.encode(rejected, add_special_tokensFalse) chosen_input prompt_ids chosen_ids [tokenizer.eos_token_id] rejected_input prompt_ids rejected_ids [tokenizer.eos_token_id] chosen_input chosen_input[:max_length] rejected_input rejected_input[:max_length] chosen_input chosen_input [tokenizer.pad_token_id] * (max_length - len(chosen_input)) rejected_input rejected_input [tokenizer.pad_token_id] * (max_length - len(rejected_input)) return chosen_input, rejected_input这段代码里 eos_token_id 用来标记回复结束pad_token_id 用来对齐长度。训练时通过 attention mask 区分真实 token 和 padding token这样 Loss 计算时才不会把 padding 部分算进去。实际情况中 Qwen3 的 tokenizer 没有默认 pad token需要在加载后手动设置通常把 eos token 作为 pad token 使用。4. DeepSpeed 训练主流程从配置文件到 Loss 实现4.1 DeepSpeed 配置文件的核心参数逻辑进入正式训练前需要写好 DeepSpeed 配置文件。我不建议用命令行参数硬塞维护一个 JSON 文件更清晰也方便多卡复用。下面这个配置是我在实际项目里验证过的{ train_batch_size: 16, train_micro_batch_size_per_gpu: 2, gradient_accumulation_steps: 4, optimizer: { type: AdamW, params: { lr: 5e-6, weight_decay: 0.01 } }, fp16: { enabled: false, loss_scale: 0, loss_scale_window: 1000 }, bf16: { enabled: true }, zero_optimization: { stage: 2, allgather_partitions: true, reduce_scatter: true, overlap_comm: true, contiguous_gradients: true }, gradient_clipping: 1.0 }几个关键参数我解释一下。train_batch_size 是全局 batch size等于 micro batch size 乘以卡数再乘以梯度累积步数这里 2 x 2 x 4 16。学习率用 5e-6 而不是常见的 5e-5因为奖励模型的底座已经具备很强的语言理解能力学习率太大会把底座学坏这个教训是我用 5e-5 跑了一次之后发现的模型很快就出现了对训练集过拟合但验证集准确率不涨的迹象。fp16 和 bf16 的区别值得说清楚。fp16 的动态范围小loss 一放大就溢出经常出现 loss 直接变 NaN 的情况。bf16 的指数范围和 fp32 一样只是精度低一点但训练奖励模型这种对梯度精度要求不高的任务完全够用。如果显卡支持 bf16优先用 bf16。4.2 回归头的设计把生成模型改成评分模型Qwen3 本身没有回归头需要自己在模型顶部加一个线性层。这里的关键是回归头的初始化。我一开始直接随机初始化结果前几个 step 的 loss 波动特别大。后来改成从零初始化并降低学习率让回归头在训练早期慢慢适应底座的特征分布收敛才平稳下来。import torch.nn as nn class RewardModel(nn.Module): def __init__(self, base_model, hidden_dim2560): super().__init__() self.base_model base_model self.value_head nn.Linear(hidden_dim, 1) def forward(self, input_ids, attention_mask): outputs self.base_model( input_idsinput_ids, attention_maskattention_mask, output_hidden_statesTrue, ) hidden outputs.hidden_states[-1] masked_hidden hidden * attention_mask.unsqueeze(-1) pooled masked_hidden.sum(dim1) / attention_mask.sum(dim1, keepdimTrue) reward self.value_head(pooled) return reward这里的池化方式用了 mean pooling取所有非 padding token 的 hidden state 的平均值再映射到标量。有人用最后一个 token 的 hidden state也有人用 first token 的 hidden state各有道理。我的经验是 mean pooling 对短回复和长回复之间的不公平有较好的抑制作用训练出来的分数更稳健。4.3 Pairwise Ranking Loss 的原理与实现奖励模型的核心 Loss 是 pairwise ranking loss公式表达为模型对好回复的分数要高于差回复的分数且差值越大越好。通俗地说我们不关心模型打分的绝对值只关心相对高低是否正确。这和普通回归任务的 Loss 有本质区别。def compute_loss(chosen_reward, rejected_reward): logsigmoid_diff torch.nn.functional.logsigmoid(chosen_reward - rejected_reward) loss -logsigmoid_diff.mean() return loss这段代码看起来很简单但背后的直觉是当 chosen_reward 明显大于 rejected_reward 时log 部分接近 0loss 趋近于 0当大小关系搞反时loss 会很大梯度会推动模型修正。这个 Loss 对所有样本对求平均因此 batch 内不同样本之间的难度会被平衡掉。实际操作中还有一个细节除了这个 Loss有人还会加一个辅助的 margin loss强制 chosen 和 rejected 的分数差异至少达到某个阈值。我实验发现对 Qwen3-4B 底座来说加 margin 收益不大反而多了一个超参数要调。建议先不加跑一个 baseline如果发现 chosen 和 rejected 分数差异很小再加 margin 不迟。4.4 完整的训练循环与梯度累积把前面所有环节串起来完整的训练循环大概长这样from deepspeed.ops.adam import FusedAdam def train_step(batch, model, optimizer): chosen_input_ids batch[chosen_input_ids].cuda() chosen_attention_mask batch[chosen_attention_mask].cuda() rejected_input_ids batch[rejected_input_ids].cuda() rejected_attention_mask batch[rejected_attention_mask].cuda() chosen_reward model(chosen_input_ids, chosen_attention_mask) rejected_reward model(rejected_input_ids, rejected_attention_mask) loss compute_loss(chosen_reward, rejected_reward) model.backward(loss) model.step() return loss.item()这里用的是 DeepSpeed 封装的model.backward()和model.step()它内部会处理梯度累积、梯度裁剪和 ZeRO 同步不需要手写optimizer.zero_grad()、loss.backward()这些操作。训练过程中我会把每个 step 的 loss 打印出来同时每隔一定步数在验证集上算一次准确率。准确率定义为 chosen_reward rejected_reward 的比例这个指标比 loss 值更直观。5. 收敛效果与实用踩坑记录从训练到推理部署的完整闭环5.1 训练过程中的显存与 loss 监控训练前先算好显存预算避免训练途中 OOM。以 Qwen3-4B 为例两个序列拼接后的长度 2048micro batch size 2ZeRO Stage 2双卡 24GB单卡显存占用实测在 18GB 左右留有一定的余量给临时张量。如果 OOM先把 micro batch size 减半或者把 max length 缩短到 1536优先保 batch 正常跑通。训练的正常 loss 走势我是这样观察的最开始几个 steploss 应该在 0.7 到 0.75 之间波动这个值接近-log(0.5)说明模型刚开始基本是瞎猜。大约训练 200 到 300 个 step 之后loss 会下降到 0.5 以下对应准确率从 50% 提升到 70% 以上。如果 loss 下降很慢优先检查数据质量九成情况是数据里混入了大量两个回复都很差的坏样本。更新完每个 epoch 的参数之后我通常会在验证集上算指标。一个简单有效的做法是每次保存 checkpoint 之后就批量跑一遍验证集记录准确率、chosen 平均分、rejected 平均分画成表格观察趋势。如果特定 epoch 准确率反而掉了说明模型开始过拟合回退到之前的 checkpoint 即可。5.2 推理部署时省略底座加载的坑训完奖励模型之后的部署环节有个非常隐蔽的坑模型保存时不小心把完整的底座参数也存了下来导致推理时加载了多余参数或是在推理脚本里没有屏蔽掉底层的生成模块导致打分速度极慢。正确的做法是只保存 value head 的权重或者只保存底座的 LoRA 权重。推理时用一个统一的模型类把底座和回归头加载起来但关闭梯度计算并确保 tokenizer 里不会生成多余内容。我经常遇到的一个问题是想当然地调model.generate()奖励模型根本没有生成能力直接报错或者输出空内容。奖励模型的输出就是一个标量分数拿这个分数排序才是正确的用法。推理阶段还有一个小优化因为奖励模型只需要评分不需要生成所以可以一次性把 batch 拉到足够大。我实测在单卡 24GB 上把 batch size 拉到 32 推理耗时也不会显著增加因为瓶颈在底座前向传播而前向传播对 batch 的利用率很高。所以奖励模型部署的成本远低于生成模型部署的成本这也是奖励模型适合自建的原因之一。5.3 实际应用中遇到的三个隐藏问题训练完成并不代表一劳永逸真正上线后我才发现奖励模型还有很多边界情况要处理。第一个是上下文长度敏感。我们的业务提示里经常有很长的背景说明如果训练时 max_length 设成 2048线上出现超过这个长度的输入tokenizer 截断后会把关键信息截掉打分明显异常。这个问题的做法我是在数据处理时就按线上分布来设置 max_length宁可牺牲一些训练效率也要保证长度分布一致不要只图快图省事。第二个是套话偏好问题。奖励模型很容易学会给某种固定话术打高分比如包含特定措辞、特定结构的回复总是拿到高分。这其实是数据分布的偏置。解决思路是在数据里刻意加入一些好话术但坏内容的反例用模型学句子漂亮不等于内容好。尽管没法完全消除但加上这些样本之后模型在人工评测里的误判率显著下降。第三个是校准问题。奖励模型输出 0.8 分和 0.9 分之间的差值到底表示多大的偏好强度这其实没有严格定义。如果要拿奖励分数去做阈值决策需要额外做一步校准比如收集一些线上样本让标注员排序找到分差和真实偏好概率之间的映射。这一步看起来繁琐但对业务落地很重要否则奖励分数只能做相对排序不能做绝对判断。5.4 后续可以继续探索的方向在现有的 DeepSpeed Qwen3 奖励模型基础上后续有几个方向值得尝试。一个是把训练好的奖励模型接进强化学习管线用 PPO 或者 GRPO 算法继续对齐策略模型这是奖励模型最有价值的应用。另一个方向是尝试更高质量的数据构造方案用迭代式标注让奖励模型和生成模型互相促进。还有一个方向是试一下多任务学习同时让模型学会打分和给出解释输出不仅能告诉策略模型哪里不好还能告诉它为什么不好。我在实际使用中发现奖励模型的调优过程有点像给一个资深评审员做入职培训它不需要自己会写文章但要有一双看得出好坏的眼睛。数据喂得好它判断就准数据喂得杂它判断就会飘。这也是整个训练过程中最需要耐心的地方。如果做好了你会发现后续所有对齐实验的迭代速度都会明显变快很多方向性问题在奖励模型这一关就能被挡下来。