TRL RewardTrainer 实战指南从偏好数据到 Bradley-Terry 奖励模型的完整训练方案【免费下载链接】trlTrain transformer language models with reinforcement learning.项目地址: https://gitcode.com/GitHub_Trending/tr/trl导读本文围绕 TRL 的RewardTrainerOutcome-supervised Reward ModelingORM 训练器展开完整覆盖其偏好数据集格式要求、预处理与 tokenization 流程、基于 Bradley-Terry 模型的损失函数实现、RewardConfig关键参数、PEFT 适配器训练以及工具调用tool calling微调等核心内容。读完本文你可以独立复现从原始偏好数据集到可打分奖励模型的完整训练链路并能结合仓库源码理解每一步在trl/trainer/reward_trainer.py中的实际落地方式。什么是 Reward Modeling强化学习后训练RLHF 类流程通常需要一个奖励模型Reward ModelRM它在成对的偏好/非偏好回复数据上通过监督学习学会给偏好回复打更高分从而能够对模型的输出进行排序与打分供下游 PPO 或 DPO 等算法使用。TRL 提供了 ORM 训练器RewardTrainer来训练这类奖励模型。从实现结构看RewardTrainer继承自transformers.Trainer经由trl/trainer/base_trainer.py中的_BaseTrainer并做了三处关键改造模型头改造以AutoModelForSequenceClassification加载因果语言模型num_labels被强制设为 1使模型对整条序列输出一个标量奖励见 reward_trainer.py数据管道改造内置DataCollatorForPreference把 chosen/rejected 两条序列动态打包进同一个 batch见 reward_trainer.py损失函数改造覆写compute_loss实现 Bradley-Terry 偏好损失与可选的中心化正则项见 reward_trainer.py。快速开始最小示例用RewardTrainer在 UltraFeedback 二分类偏好数据集上训练 Qwen3-0.6Bfrom trl import RewardTrainer from datasets import load_dataset trainer RewardTrainer( modelQwen/Qwen3-0.6B, train_datasetload_dataset(trl-lib/ultrafeedback_binarized, splittrain), ) trainer.train()训练完成后可通过trainer.save_model(output_dir)保存模型或通过trainer.push_to_hub()推送到 Hub。_save_checkpoint中还会自动生成模型卡片模板位于 rm_model_card.md随 checkpoint 一起保存方便标注模型用途。数据集类型与格式要求RewardTrainer支持preference偏好类型数据集同时兼容standard纯文本与conversational对话消息两种格式且 prompt 可以是隐式的直接包含在 chosen/rejected 中或显式的单独的prompt字段。四种组合如下# Standard preference (implicit prompt) {chosen: The sky is blue., rejected: The sky is green.} # Conversational preference (implicit prompt) {chosen: [{role: user, content: What color is the sky?}, {role: assistant, content: It is blue.}], rejected: [{role: user, content: What color is the sky?}, {role: assistant, content: It is green.}]} # Standard preference (explicit prompt) {prompt: The sky is, chosen: blue., rejected: green.} # Conversational preference (explicit prompt) {prompt: [{role: user, content: What color is the sky?}], chosen: [{role: assistant, content: It is blue.}], rejected: [{role: assistant, content: It is green.}]}当传入 conversational 格式时trainer 会自动对数据集应用 chat templateprompt与 completion 分开给出时会在 tokenization 前拼接。更完整的格式定义包括 tool calling、vision 等参见 dataset_formats.md。把任意数据集转换成偏好格式如果你的数据不是上述格式例如 response_a/response_b winner 列的标注数据可以先做预处理。以下示例以 lmarena-ai/arena-human-preference-55k 数据集为例分三步完成转换from datasets import load_dataset import json dataset load_dataset(lmarena-ai/arena-human-preference-55k) # 1. 过滤平局样本 dataset dataset.filter(lambda example: example[winner_tie] 0) # 2. 根据 winner 列构造 chosen / rejected 字段 def response_a_b_to_chosen_rejected(example): if example[winner_model_a] 1: example[chosen] example[response_a] example[rejected] example[response_b] else: example[chosen] example[response_b] example[rejected] example[response_a] return example dataset dataset.map(response_a_b_to_chosen_rejected) # 3. 转换为 conversational 格式 def make_conversation(example): prompt json.loads(example[prompt])[0] chosen json.loads(example[chosen])[0] rejected json.loads(example[rejected])[0] return { chosen: [{role: user, content: prompt}, {role: assistant, content: chosen}], rejected: [{role: user, content: prompt}, {role: assistant, content: rejected}], } dataset dataset.map(make_conversation) # 只保留必要列 dataset dataset.select_columns([chosen, rejected]) print(next(iter(dataset[train])))转换后的样本形如{ chosen: [ {role: user, content: Is it morally right to try to have a certain percentage of females on managerial positions?}, {role: assistant, content: The question of whether it is morally right to aim for a certain percentage of females...} ], rejected: [ {role: user, content: Is it morally right to try to have a certain percentage of females on managerial positions?}, {role: assistant, content: As an AI, I dont have personal beliefs or opinions. However, ...} ] }训练方法深入剖析预处理与 tokenization训练时每条样本必须包含chosen与rejected字段。trainer 在初始化时调用_prepare_datasetreward_trainer.py完成整个预处理流程如下已 tokenized 数据直通若数据集已含chosen_ids与rejected_ids列跳过处理旧列名chosen_input_ids/rejected_input_ids会触发弃用警告并自动重命名。追加 EOS对 non-conversational 数据若chosen/rejected未以 EOS token 结尾则自动追加保证奖励模型能感知序列结束位置。tokenization显式 prompt 场景下先执行example[chosen] example[prompt] example[chosen]完成拼接再用模型 tokenizer 编码conversational 数据走_tokenize内部应用 chat template并支持样本级chat_template_kwargs覆盖。长度过滤max_length默认 1024不为None时chosen 或 rejected 任一超过该长度的样本会被整体过滤掉而非截断。分布式环境下这些map/filter操作在PartialState().main_process_first()保护下执行即只有主进程实际处理其余进程等待避免多进程重复计算。数据打包DataCollatorForPreference与常见的按序列 collate 不同DataCollatorForPreferencereward_trainer.py把 batch 内每个样本的 chosen 与 rejected展平为 2N 条序列前半 batch 是 chosen后半 batch 是 rejected并动态 pad 到 batch 内最长长度右填充pad_token_id取自 tokenizer。若样本携带margin字段collator 会额外输出margin张量。batch_size因此实际是偏好对数量显存占用约为单序列训练的 2 倍规划 batch size 时需注意。损失函数Bradley-Terry 模型设输入序列prompt为 ( x )chosen 与 rejected 序列分别为 ( y^ ) 与 ( y^- )。在 Bradley-Terry 模型Bradley Terry, 1952下偏好回复 ( y^ ) 胜过 ( y^- ) 的概率为$$ p(y^ \succ y^- \mid x) \sigma\big(r(x, y^) - r(x, y^-)\big) $$其中 ( \sigma ) 为 sigmoid 函数。训练目标是最大化观测偏好的对数似然即最小化损失$$ \mathcal{L}(\theta) - \mathbb{E}_{(x, y^, y^-) \sim \mathcal{D}} \left[ \log \sigma\big(r_\theta(x, y^) - r_\theta(x, y^-)\big) \right] $$对应源码在compute_lossreward_trainer.py# 模型输出按 batch 拆成 chosen / rejected 两半 rewards_chosen, rewards_rejected torch.chunk(outputs.logits.squeeze(-1), chunks2) # Bradley-Terry 损失可选 margin 调制 if margin in inputs: loss -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - inputs[margin]).mean() else: loss -nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean() # 可选的奖励中心化正则 if self.args.center_rewards_coefficient is not None: loss self.args.center_rewards_coefficient * torch.mean((rewards_chosen rewards_rejected) ** 2)三个实现细节值得注意margin 支持若数据集提供margin列损失变为-log σ(r_chosen - r_rejected - margin)即要求 chosen 的奖励至少比 rejected 高出 margin。该特性在 tests/test_reward_trainer.py 的test_train_with_margin中有回归测试覆盖。奖励中心化正则Bradley-Terry 模型是欠定的——给所有奖励加上同一常数并不改变偏好概率。论文Helping or Herding? Reward Model Ensembles Mitigate but do not Eliminate Reward Hacking提出附加一项鼓励奖励均值归零的辅助损失。TRL 中由center_rewards_coefficient控制推荐值1e-2对应源码中0.01 * mean((r_chosen r_rejected)^2)这一项其训练路径由 test_train_with_center_rewards_coefficient 验证。梯度检查点默认开启RewardConfig覆写了若干默认值gradient_checkpointingTrue、bf16True、logging_steps10、learning_rate1e-4见 reward_config.py 的 NOTE 说明因此默认配置下训练即启用 checkpointing 与 bf16长序列偏好对训练的显存压力可控。RewardConfig 关键参数详解RewardConfig继承自 TRL 的_BaseConfig后者继承transformers.TrainingArguments因此除下表列出的奖励训练专用参数外还支持全部标准训练参数output_dir、per_device_train_batch_size、eval_strategy、gradient_accumulation_steps等。专用参数如下依据 reward_config.py参数默认值说明learning_rate1e-4初始学习率AdamW覆盖TrainingArguments的5e-5model_init_kwargsNone传给AutoModelForSequenceClassification.from_pretrained的关键字参数当model以字符串传入时生效其中revision同时用于加载 tokenizertrust_remote_codeFalse是否允许加载携带自定义 Python 代码的 Hub 模型与 tokenizerchat_template_pathNone指定 chat 模板来源tokenizer 路径或 Jinja 模板文件路径disable_dropoutTrue训练前是否禁用模型 dropoutdataset_num_procNone数据集处理map/filter使用的进程数eos_tokenNone回合/序列结束 token默认取processing_class.eos_tokenmax_length1024tokenized 序列最大长度超长样本被过滤非截断设为None则不过滤pad_to_multiple_ofNone序列 pad 到该值的倍数center_rewards_coefficientNone奖励中心化正则系数推荐0.01activation_offloadingFalse是否将激活值 offload 到 CPU与TrainingArguments相比另有默认值差异需要注意logging_steps为10而非500gradient_checkpointing为Truebf16在fp16未设置时默认为True。模型加载相关约定来自 reward_trainer.py 的初始化逻辑字符串模型名默认以float32加载model_init_kwargs未指定dtype时与 Transformers v5 的 config 推断行为不同传入已实例化的模型时num_labels必须为 1否则抛出ValueError——奖励模型每条序列只输出一个标量分数分布式训练多 GPU / DeepSpeed下device_map会被强制为None。定制训练通过 model_init_kwargs 控制模型初始化RewardConfig.model_init_kwargs会原样转发给AutoModelForSequenceClassification.from_pretrained例如指定精度from trl import RewardConfig training_args RewardConfig( model_init_kwargs{dtype: torch.bfloat16}, )除num_labels外from_pretrained的所有关键字参数均支持num_labels被自动设为 1。若模型已经实例化再传入 trainermodel_init_kwargs会被忽略并给出警告。使用 PEFT 训练适配器RewardTrainer与 PEFT 深度集成可以只训练 LoRA 适配器而非全量模型from datasets import load_dataset from trl import RewardTrainer from peft import LoraConfig dataset load_dataset(trl-lib/ultrafeedback_binarized, splittrain) trainer RewardTrainer( Qwen/Qwen3-4B, train_datasetdataset, peft_configLoraConfig(modules_to_save[score]) # 基座不是序列分类模型时务必包含 score 头 ) trainer.train()modules_to_save[score]是关键从因果 LM checkpoint 加载时新的score分类头是随机初始化的必须让它参与训练否则模型只靠 LoRA 分支学习而奖励头保持随机。继续训练已有适配器则直接加载PeftModel并以is_trainableTrue传入from datasets import load_dataset from trl import RewardTrainer from peft import AutoPeftModelForCausalLM model AutoPeftModelForCausalLM.from_pretrained(trl-lib/Qwen3-4B-Reward-LoRA, is_trainableTrue) dataset load_dataset(trl-lib/Capybara, splittrain) trainer RewardTrainer( modelmodel, train_datasetdataset, ) trainer.train()提示训练适配器时通常应使用更高的学习率约 1e-3因为实际被优化的是新参数RewardConfig(learning_rate1e-3, ...)。PEFT 路径还有若干源码层面的自动处理reward_trainer.pyDeepSpeed ZeRO-3 非量化 PEFT 时自动关闭autocast_adapter_dtype以避免混合精度 buffer 报错PEFT 梯度检查点时自动调用enable_input_require_grads()QLoRA 场景下可训练参数会被转为 bf16。训练工具调用模型RewardTrainer完整支持对具备 tool calling 能力的模型做奖励微调。此时数据集每条样本需包含对话消息其中可包含工具调用tool_calls字段与工具返回tool角色的消息一个tools列通常为 JSON Schema 格式的可用工具列表。tokenization 阶段tokenize_fn会读取样本的tools列字符串会json.loads解析并将其传入 chat template 应用过程reward_trainer.py因此 Qwen2.5 等支持工具模板的模型可以直接对带工具调用的偏好对打奖励分。格式细节见 dataset_formats.md 中 Tool Calling 相关章节。该能力在 tests/test_reward_trainer.py 中由test_train_toolcall_data与test_train_toolcall_data_as_json两个用例验证前者确认偏好 工具数据可正常训练后者确认tools以 JSON 字符串形式存储时的兼容性。训练与评估指标训练和评估期间trainer 除记录 Transformers 标准日志global_step、epoch、loss、learning_rate、grad_norm外还记录以下奖励模型专属指标累加值在 log() 中取区间平均后输出指标含义num_tokens累计处理的 token 总数accuracy打分的偏好准确率即 chosen 得分高于 rejected 的比例min_reward/mean_reward/max_reward模型输出奖励的最小 / 平均 / 最大值跨进程 gather 后统计marginchosen 与 rejected 奖励差的平均值监控accuracy与margin是判断奖励模型质量的最直接方式训练收敛时 accuracy 应稳定上升、margin 持续拉大若mean_reward持续漂移远离 0建议启用center_rewards_coefficient。命令行训练脚本除 Python API 外仓库还提供命令行入口 trl/scripts/reward.py支持以trl reward子命令方式运行CLI 总览。该脚本通过TrlParser解析ScriptArguments、RewardConfig、ModelConfig、DatasetMixtureConfig四类参数将模型侧参数model_revision、trust_remote_code、attn_implementation、dtype组装进model_init_kwargs支持 QLoRA 量化quantization_config与 PEFTpeft_config训练结束后自动save_model并可选push_to_hub。数据集既可通过--dataset_name从 Hub 加载也可通过--datasets参数指定混合数据集。小结RewardTrainer的完整链路是偏好格式数据集 → 自动 EOS 追加与 tokenizationconversational 走 chat template支持 tools→ max_length 过滤 → DataCollatorForPreference 动态打包 → Bradley-Terry 损失可选 margin、奖励中心化正则→ 专属指标日志。所有环节均由 trl/trainer/reward_trainer.py 与 trl/trainer/reward_config.py 中的实现背书回归测试集中在 tests/test_reward_trainer.py。掌握了本文的参数表与源码对照你可以从最小示例平滑扩展到 LoRA/QLoRA 训练、工具调用偏好数据以及分布式环境下的奖励模型训练。【免费下载链接】trlTrain transformer language models with reinforcement learning.项目地址: https://gitcode.com/GitHub_Trending/tr/trl创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考