跳到主要内容

数据截至 (上游 commit 67dfbe211a07)

05 · 数据格式工具与奖励管线

这一章讲什么： 前几章反复出现的「数据格式约定」在这里收拢成完整图景：trl/data_utils.py 的判定与转换工具、每种 Trainer 期待的列、以及 trl/rewards/ 里现成的奖励函数。准备训练数据或写自定义奖励函数时，先读这章。

1. 它要解决的小问题​

后训练数据集在现实里有无数种长相：纯文本、指令-回答对、成对偏好、点赞点踩、ShareGPT 导出的对话、带工具定义的对话……如果每个 Trainer 各自处理，就会重复造轮子且互不一致。

TRL 的约定是两层：

  1. 格式判定层：任何数据先被分成两类——standard（值是字符串）或 conversational（值是消息列表）。
  2. 列约定层：每种方法只认固定的列名组合。所有清洗、转换都发生在数据进 Trainer 之前（或 Trainer 的 _prepare_dataset 第一步）。

2. 格式判定与转换：data_utils 的工具箱​

2.1 判定：看值，不看列名​

is_conversational（trl/data_utils.py:160）的规则极简：取 prompt/chosen/rejected/completion/messages 中任一存在的列，值是列表且首元素是带 role 键的字典，就是对话格式。

# 示意，非源码
is_conversational({"prompt": [{"role": "user", "content": "hi"}]})  # True
is_conversational({"prompt": "The sky is"})                         # False

这个「duck typing」判定让所有下游函数不必关心数据来自哪个数据集。

2.2 转换工具一览​

函数干什么位置
apply_chat_template消息列表 → token ids（走 tokenizer 的 Jinja 模板），可顺带产出 assistant_maskstrl/data_utils.py:201
maybe_apply_chat_template只有对话格式才应用模板，否则原样返回trl/data_utils.py:337
maybe_convert_to_chatml旧式 {"conversations": [{"from": "human", ...}]} → 标准 {"messages": [{"role": ...}]}trl/data_utils.py:975
extract_prompt偏好数据里 prompt 没单列时，取 chosen/rejected 的最长公共前缀当 prompttrl/data_utils.py:557
unpair_preference_dataset成对偏好（chosen/rejected）拆成两行 prompt/completion/label，给 KTO 用trl/data_utils.py:454
pack_dataset装箱（第 2 章 §4 详述）trl/data_utils.py:845

两个值得单独说的：

  • extract_prompt 的前缀是「对话轮次级」而非「字符串级」——chosen 和 rejected 共享的是前面的消息轮（trl/data_utils.py:557-600 的 docstring 示例）。这保证拆出的 prompt 一定终止在角色边界上。
  • unpair_preference_dataset 是「格式即算法」的好例：DPO 数据一行拆两行、chosen 行 label=True、rejected 行 label=False（docstring 示例见 trl/data_utils.py:483-493），偏好对数据就此变成 KTO 数据。

2.3 各方法期待的列（速查表）​

Trainer需要的列两种格式都支持？
SFTTrainertext；或 prompt+completion；或 messages是
DPOTrainerprompt + chosen + rejected（prompt 可缺省，用 extract_prompt 推）是
KTOTrainerprompt + completion + label(bool)是
RewardTrainerprompt + chosen + rejected（可带 margin 列做 margin loss，trl/trainer/reward_trainer.py:147-148）是
GRPOTrainer / RLOOTrainer只需 prompt；其余列原样透传给奖励函数是

3. 奖励管线：把「打分」开放给用户函数​

3.1 签名约定​

在线 RL（GRPO/RLOO）没有标注数据，标注由奖励函数现场产生。约定（RLOOTrainer docstring 写得最全，trl/trainer/rloo_trainer.py:141-165）：

# 示意，非源码：奖励函数的统一签名
def my_reward(prompts, completions, completion_ids=None, **kwargs) -> list[float | None]:
    # kwargs 里有数据集的全部其他列（如数学题的 solution 列）
    # 还有 trainer 注入的 trainer_state / log_extra / log_metric
    return [1.0 if good(c) else 0.0 for c in completions]

三种形态（模型 id 字符串 / PreTrainedModel / 普通函数，可混合成列表）已在第 4 章 §5.1 讲过，这里补数据侧的两个约定：

  • 数据集其他列会按名字透传（trl/trainer/grpo_trainer.py:1637-1638）：数据里有 solution 列，奖励函数就能收 solution 参数。列名即接口。
  • 返回 None = 此样本不归这个函数管（转 NaN，组内归一时跳过）：多任务混训时，数学奖励函数看到代码题就返回 None。

3.2 内置奖励函数​

trl/rewards/ 提供开箱即用的几个：

函数干什么位置
accuracy_reward数学答案核对：抽 \boxed{} 或最终答案，与 solution 列做符号级比对（解析不了就返回 None）trl/rewards/accuracy_rewards.py:28
reasoning_accuracy_reward带 <think> 推理段模型的答案核对trl/rewards/accuracy_rewards.py:223
think_format_reward检查输出是否符合 think/answer 格式约定trl/rewards/format_rewards.py:18
get_repetition_penalty_rewardn-gram 重复惩罚（返回惩罚函数的工厂）trl/rewards/other_rewards.py:18
get_soft_overlong_punishment超长软惩罚：超过预算越长扣越多trl/rewards/other_rewards.py:83

accuracy_reward 是「可验证奖励」的范本：不用模型、纯程序判定（math-verify 库做 LaTeX 等价判断），零成本零噪声。DeepSeek-R1 式训练的主要奖励就长这样。

3.3 奖励模型在哪​

如果奖励无法程序化判定（开放域偏好），就得训一个奖励模型——这是 RewardTrainer（trl/trainer/reward_trainer.py:227）的活：吃 chosen/rejected 偏好对，把序列分类模型训成「chosen 分高于 rejected 分」（可选 margin）。训出来的模型可以直接作为字符串 id 传给 GRPOTrainer(reward_funcs="org/my-rm")，闭环。

4. 一图收拢：从原始数据到 loss​

原始数据集（各种长相）
   │  is_conversational 判定
   ▼
standard / conversational  ──maybe_convert_to_chatml──▶ 统一消息格式
   │
   │  按目标方法整形:
   │   SFT → text / prompt+completion / messages
   │   DPO → prompt+chosen+rejected（extract_prompt 补 prompt）
   │   KTO → prompt+completion+label（unpair_preference_dataset 拆对）
   │   GRPO/RLOO → prompt + 任意透传列（喂奖励函数）
   ▼
各 Trainer._prepare_dataset: apply_chat_template → tokenize → 掩码 → labels
   │
   ▼
collator 拼 batch → compute_loss

一句话： 数据侧的一切复杂度都被推到「进 Trainer 之前」；Trainer 内部只见 token 和掩码。

5. 关键细节与坑​

  • 列名就是契约。 GRPO 数据集里多余的列不会报错，会原样进奖励函数的 kwargs——拼错列名（如 solutions vs solution）的代价是奖励函数收到意外参数。
  • assistant_masks 依赖模板带 {% generation %}（见第 2 章 §3）；Qwen/Llama 官方模板支持程度不一，用前先小数据验证。
  • v1 起不再自动剥 None 值。 旧版 TRL 会对嵌套列里的 None 做清理，v1 移除了这个行为（MIGRATION.md:18-38）；老数据集遇到嵌套 None 问题需自己 dataset.with_transform(remove_none_values)（trl/trainer/utils.py:1110）。
  • 偏好数据 prompt 省略是双刃剑。 extract_prompt 按轮次前缀推断，chosen/rejected 第一轮就分叉的数据会推出空 prompt——检查数据集时先看这个边界情形。
  • 奖励函数里做重度计算要慎重。 同步函数会阻塞整个 rollout 打分阶段；网络/沙箱类判定请写成 async 函数（第 4 章 §5.1）。

6. 代码地图（本章）​

主题文件路径符号名
格式判定trl/data_utils.pyis_conversational、is_conversational_from_value
chat template 应用trl/data_utils.pyapply_chat_template、maybe_apply_chat_template
旧格式转换trl/data_utils.pymaybe_convert_to_chatml
偏好数据整形trl/data_utils.pyextract_prompt、unpair_preference_dataset
None 值清理（v1 需手动）trl/trainer/utils.pyremove_none_values
数学/推理奖励trl/rewards/accuracy_rewards.pyaccuracy_reward、reasoning_accuracy_reward
格式奖励trl/rewards/format_rewards.pythink_format_reward
重复/超长惩罚trl/rewards/other_rewards.pyget_repetition_penalty_reward、get_soft_overlong_punishment
奖励模型训练trl/trainer/reward_trainer.pyRewardTrainer
奖励函数签名约定trl/trainer/grpo_trainer.py_calculate_rewards