奖励头设计全指南如何在SFT骨干上添加一个标量头3步训出RLHF奖励模型【免费下载链接】train-llm-from-scratchA straightforward method for training your LLM, from downloading data to generating text.项目地址: https://gitcode.com/GitHub_Trending/tr/train-llm-from-scratch在开源项目train-llm-from-scratch中奖励头Reward Head是 RLHF 后训练流程里最关键的一步它把 SFT 骨干网络改造成一个给回复打分的评委只需用一个标量 Linear 层就能为每条回复输出一个分数。本文带你完整理解奖励模型Reward Model设计为什么奖励要加在 SFT 骨干上、为什么从最后一个真实 token读数、以及 Bradley-Terry 成对损失如何训练它让你快速掌握从sft.pt到reward.pt的完整链路 奖励模型在整条流水线中的位置先建立全局视角本项目从 The Pile 数据预训练出基座经过 SFT 指令微调后奖励模型作为分支之一被训练出来再为 PPO 提供标量奖励信号如果没有奖励头PPO 就失去了什么是更好的回复的判据——它是经典 RLHF 里连接人类偏好与强化学习的桥梁。 奖励头结构一个 Linear 层就够了核心实现在 src/post_training/reward_model.pyRewardModel做的事非常克制包裹一个 Transformer 骨干直接复用 SFT 检查点里的权重丢弃lm_head不再预测下一个 token加一个标量头nn.Linear(n_embed, 1, biasFalse)把隐藏状态压成一个数。整个奖励模型 SFT 骨干 一个线性投影。这正是 InstructGPT 的配方self.reward_head nn.Linear(n_embed, 1, biasFalse) nn.init.zeros_(self.reward_head.weight) # 初始奖励接近 0零初始化是个小心机训练初期所有回复的奖励都接近 0损失从抛硬币水平起步梯度信号干净训练更稳定。 从哪读数最后一个真实 token 的隐藏状态为什么不取整个序列的平均项目遵循 InstructGPT 惯例取最后一个真实 token非 padding的隐藏状态。原因很巧妙——注意力是因果的最后一个真实 token 已经看过整条序列且永远不会关注它右侧的 padding因此不需要额外的 attention mask。读数逻辑在 src/post_training/utils.py 的gather_last中对每个 batch 行按真实长度seq_lengths索引出对应位置的值即可一行搞定idx (seq_lengths - 1).clamp(min0).long() return values[torch.arange(B), idx]偏好数据由 data_loader/preference_dataset.py 提供每行 JSONL 包含{prompt, chosen, rejected}三元组右填充安全无副作用因为因果注意力天然屏蔽了尾部 padding。 Bradley-Terry 损失整个训练信号只有两行训练目标在 src/post_training/reward_train.py核心只有一个公式L -log sigmoid(r_chosen - r_rejected)即让被偏好chosen回复的分数高于被拒绝rejected的分数。配套的三个监控指标指标含义健康范围lossBradley-Terry 损失起点为-log σ(0) ≈ 0.693抛硬币持续下降preference_accr_chosen r_rejected的比例真实噪声数据约0.65–0.75属正常margin平均分差r_chosen − r_rejected越大说明分离越好⚠️ 注意在真实、带噪声的偏好数据如 HH-RLHF / UltraFeedback上准确率到 0.7 左右就很好了——人类偏好本身就是噪声的不必追求 1.0。 一键训练chosen 与 rejected 一次前向跑完训练入口是 scripts/train_reward.py从sft.pt加载骨干 → 包装RewardModel→ 用 Bradley-Terry 损失迭代。一个高效细节每个 batch 把 chosen 和 rejected拼接成一个 2B 序列做单次前向再切开得分GPU 利用率翻倍# 单卡 PYTHONPATH. python scripts/train_reward.py # 双卡 PYTHONPATH. torchrun --standalone --nproc_per_node2 scripts/train_reward.py默认超参见 configs/reward.jsonlr1e-5、batch_size8、max_len768对应配置数据类是 config/post_training_config.py 中的RewardConfig。细节彩蛋DDP 包装时特意开了find_unused_parametersTrue——因为奖励模型从不使用lm_head那些参数拿不到梯度不开这个开关第一次反向传播就会报错。 训练之后奖励头如何喂给 PPO训好的reward.pt由 PPO 阶段通过load_reward_model加载见 src/post_training/reward_model.py加载后参数冻结、只用于打分PYTHONPATH. python scripts/train_ppo.py --reward_source rmreward_sourcerm表示使用训好的奖励模型也可以换成verifier如 GSM8K 正确性检查器做免奖励模型路径。完整后训练命令速查见 POST_TRAINING.md。 延伸阅读与源码导航阶段详解文档docs/04_reward_model.md含可编辑的 Mermaid 图源奖励模型源码src/post_training/reward_model.py训练目标源码src/post_training/reward_train.py训练脚本scripts/train_reward.py偏好数据加载器data_loader/preference_dataset.py如果你还不熟悉 Transformer 骨干如何产生隐藏状态建议先看 docs/foundations/transformer.md 打底再回来读奖励头会更顺畅——项目内置了循序渐进的基础教程一句话总结奖励头设计的全部精髓就是三件事——SFT 骨干 最后一个真实 token 一个零初始化的标量 Linear配上 Bradley-Terry 成对损失就能为 PPO 提供高质量的偏好奖励信号 ✨【免费下载链接】train-llm-from-scratchA straightforward method for training your LLM, from downloading data to generating text.项目地址: https://gitcode.com/GitHub_Trending/tr/train-llm-from-scratch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
