PRM Trainer
AI summary · Qwen Summary in Chinese Indexed
Original source 该页面介绍 TRL 中实验性 PRM Trainer 的论文背景、示例训练与评估、数据格式以及 PRMTrainer 和 PRMConfig 的主要参数。
概述与来源
- PRM Trainer 是 TRL 文档中的实验性 API，可能随时变化。
- 过程监督奖励模型（PRM）源自论文《Solving math word problems with process- and outcome-based feedback》。
- 该后训练方法由 Gaetan Lopez、Lewis Tunstall、Quentin Gallouédec 和 Agustín Piqueres 贡献。
方法效果
- 论文比较了过程监督与结果监督在 GSM8K 上的表现。
- 纯结果监督能以较少标签获得相近的最终答案错误率。
- 对正确推理步骤，需要过程监督或模拟过程反馈的学习奖励模型。
- 论文将最终答案错误从 16.8% 降至 12.7%，并将最终答案正确解中的推理错误从 14.0% 降至 3.4%。
示例训练
- 示例以 Qwen2-0.5B 为基础模型，使用 Math Shepherd 数据集的步骤监督数据。
- 训练脚本加载 AutoModelForTokenClassification、tokenizer 和 train[:10%] 数据，并用 PRMConfig 与 PRMTrainer 训练。
- 使用 accelerate launch train_prm.py 运行；文档称在 8 个 GPU 上约需 1 小时。
示例评估
- 评估脚本使用 token-classification pipeline 加载 trl-lib/Qwen2-0.5B-Reward-Math-Sheperd。
- 示例对一道 45 名学生分组数学题的三步推理逐步预测。
- 示例输出为 Step 1 预测 True、Step 2 预测 False、Step 3 预测 False。
数据格式
- PRM 需要步骤监督数据，包含 prompt、completions 和 labels 列。
- completions 是推理步骤列表，labels 是表示每步正确性的布尔值或浮点数。
- PRMTrainer 仅支持 standard 数据集格式。
API 与配置
- PRMTrainer 接受 model、args、data_collator、train_dataset、eval_dataset、processing_class、compute_metrics、peft_config 等参数。
- model 建议使用 AutoModelForTokenClassification；默认 data_collator 为 DataCollatorForTokenClassification，默认 compute_metrics 为 compute_accuracy。
- PRMConfig 的 PRM 特定参数包括 max_length、max_completion_length、disable_dropout、step_separator、train_on_last_step_only 和 dataset_num_proc。
- 与 TrainingArguments 相比，PRMConfig 默认 logging_steps 为 10、gradient_checkpointing 为 True、learning_rate 为 1e-5，且 bf16 在 fp16 未设置时默认为 True。