奖励建模
训练奖励模型来评分文本响应,用于 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-column | prompt | 包含提示词的列 |
text_column | --text-column | text | 包含选择响应的列 |
rejected_text_column | --rejected-text-column | rejected | 包含拒绝响应的列 |
输出模型
训练后的模型是一个 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,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")
下一步