跳转到主要内容

奖励建模

训练奖励模型来评分文本响应,用于 PPO/RLHF 训练。
重要: 奖励模型不是文本生成器。它们为给定文本输出标量分数,用于在 PPO 训练期间提供奖励。您不能将奖励模型用作普通 LLM 进行文本生成。

快速开始

aitraining llm --train \
  --model google/gemma-3-270m \
  --data-path ./preferences.jsonl \
  --project-name reward-model \
  --trainer reward \
  --prompt-text-column prompt \
  --text-column chosen \
  --rejected-text-column rejected

Python API

from autotrain.trainers.clm.params import LLMTrainingParams
from autotrain.project import AutoTrainProject

params = LLMTrainingParams(
    model="google/gemma-3-270m",
    data_path="./preferences.jsonl",
    project_name="reward-model",

    trainer="reward",

    # Column mappings (required for reward training)
    prompt_text_column="prompt",
    text_column="chosen",
    rejected_text_column="rejected",

    epochs=1,
    batch_size=4,
    lr=2e-5,
)

project = AutoTrainProject(params=params, backend="local", process=True)
project.create()

数据格式

奖励训练需要包含三列的偏好数据:
描述
prompt输入提示词/问题
chosen首选/更好的响应
rejected次选/更差的响应

示例数据

{"prompt": "Explain gravity", "chosen": "Gravity is a fundamental force...", "rejected": "gravity makes stuff fall down"}
{"prompt": "What is Python?", "chosen": "Python is a high-level programming language...", "rejected": "its a snake"}
{"prompt": "Write a greeting", "chosen": "Hello! How can I assist you today?", "rejected": "hey"}

必需参数

奖励训练需要指定所有三个列参数:
  • --prompt-text-column
  • --text-column(用于选择的响应)
  • --rejected-text-column

参数

参数CLI 标志默认值描述
prompt_text_column--prompt-text-columnprompt包含提示词的列
text_column--text-columntext包含选择响应的列
rejected_text_column--rejected-text-columnrejected包含拒绝响应的列

输出模型

训练后的模型是一个 AutoModelForSequenceClassification,它:
  • 接受文本输入
  • 返回标量奖励分数
  • 更高分数表示更好的响应
  • 通过 --rl-reward-model-path 用作 PPO 训练的输入

使用奖励模型

与 PPO 训练一起使用

aitraining llm --train \
  --model google/gemma-3-270m \
  --data-path ./prompts.jsonl \
  --project-name ppo-model \
  --trainer ppo \
  --rl-reward-model-path ./reward-model

直接推理

from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

# 加载奖励模型
model = AutoModelForSequenceClassification.from_pretrained("./reward-model")
tokenizer = AutoTokenizer.from_pretrained("./reward-model")

# 评分响应
text = "What is AI? AI is artificial intelligence, a field of computer science..."
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)

with torch.no_grad():
    outputs = model(**inputs)
    score = outputs.logits.item()

print(f"Reward score: {score}")

最佳实践

  1. 高质量偏好数据 - 奖励模型的质量取决于您的标注质量
  2. 多样化示例 - 包含各种提示词和响应质量级别
  3. 清晰的偏好信号 - 选择的应该明显优于拒绝的
  4. 平衡数据集 - 避免偏向某些响应类型
  5. 充足数据 - 目标至少 1,000+ 偏好对

示例:构建偏好数据

# 创建偏好数据的示例脚本
import json

preferences = [
    {
        "prompt": "Summarize machine learning",
        "chosen": "Machine learning is a subset of AI that enables systems to learn from data...",
        "rejected": "ml is computers learning stuff"
    },
    # 添加更多示例...
]

with open("preferences.jsonl", "w") as f:
    for item in preferences:
        f.write(json.dumps(item) + "\n")

下一步