聊聊RLHF的奖励模型——上海人工智能书生大模型的RW实践

概述

在LLM对齐问题上，OpenAI提出的RLHF训练范式最为人熟知，同时也是ChatGPT行之有效的对齐方案。

RLHF通常包含三个步骤：SFT, Reward Model, PPO, 该方案优点不需多说，缺点也很明显：训练流程繁琐、算法复杂、超参数多和计算量大，因此RLHF替代方案层出不穷。

DPO（Direct Preference Optimization）是一种非常高效的RLHF算法。它巧妙地绕过了构建奖励模型和强化学习这两个的繁琐过程，直接通过偏好数据进行微调，效果简单粗暴，在使模型输出更符合人类偏好的同时，极大地缩短了训练时间和难度。

Reward Model（奖励模型）是强化学习过程中一个关键的组成部分。它的主要任务是根据给定的输入和反馈来预测奖励值，从而指导学习算法的方向。在RLHF（Reinforcement Learning from Human Feedback）中，Reward Model 通过整合人类反馈，帮助强化学习算法更有效地优化策略。

在大语言模型训练中，Reward Model 通常指的是偏好模型（Preference Model）。通过在训练时提供相同提示词的好与坏（chosen&rejected）的回复来拟合人类的偏好，并在推理时预测出一个奖励值，以指导 RLHF 过程中 Actor 模型的优化过程。

Reward Model的应用场景包括但不限于：

RLHF训练：在使用 Proximal Policy Optimization（PPO）算法进行 RLHF 训练时，Reward Model提供奖励信号，指导模型优化策略，提高生成内容的质量并使其更贴近人类偏好。

BoN采样：在 Best-of-N（BoN）采样过程中，用户可以使用 Reward Model 对同一个提示词的多条回复进行打分，并选择奖励得分最高的生成结果，从而提升模型的输出效果。

数据构造：Reward Model 可以用于评估和过滤训练数据，或者也可以使用 Reward Model 替代人工标注来构造 DPO 训练数据。

Reward的训练流程一般如下：

基于SFT后的模型，在偏好数据集上做训练，最后得到Reward Model。

目前开源的Reward Model不多，之前qwen1开源了sft模型后，又推出了rw模型，但后面的版本没有再推出。internlm系列，在2.0版本后没有推出rw模型（最新的2.5版本没有）。既可以自己做rw，按流程来就行（如基于llama）；也可以基于某个系列模型的rw版本继续做（如internlm）。不少基于开源模型llama系列、gemma系列做rm，如下：

昆仑万维发布全新大模型奖励模型Skywork-Reward，登顶RewardBench

InternLM2-Reward Model

InternLM2提供了 Reward Model^[1] 模型及相关的说明文档，训练工具是XTuner。

训练优势

XTuner 中的 Reward Model 训练具备以下显著优势：

1. 使用最新的训练技巧：XTuner 中集成了 InternLM2 中的 Reward Model 训练损失函数，可以稳定奖励得分的数值范围，也可以减少在简单样本上的过拟合（具体可参考 InternLM2 技术报告^[2]）。
2. 减少显存浪费：由于偏好数据中的 chosen 和 rejected 数据通常存在长度上的差异，因此在训练数据的拼接时会存在填充（padding token）,造成显存浪费。在 XTuner 中，基于 Flash Attention2 中的变长注意力功能，我们在训练过程中通过将偏好数据打包到同一个序列中，显著减少了由于 padding token 带来的显存浪费。这不仅提高了显存的利用效率，还使得在相同硬件条件下可以训练更大的模型或处理更多的数据。

1. 高效训练：借助 XTuner 的 QLoRA 训练功能，我们能够仅对 Reward Model 的 Value Head 进行全参数训练，而对语言模型本身使用 QLoRA 微调，大幅降低了模型训练的显存开销。
2. 长文本训练: 借助 XTuner 的序列并行功能，能够对长文本数据进行训练。

开始训练

基于InternLM2-chat-1.8b-sft^[3]作为初始化模型，使用 UltraFeedback^[4] 数据集，为了方便起见，使用 huggingface 上已经预处理过的 argilla/ultrafeedback-binarized-preferences-cleaned^[5] 数据集。

XTuner 提供了多个开箱即用的配置文件，可以通过 xtuner list-cfg 查看。也可以直接查看github上的代码文件^[6]，复制internlm2_chat_1_8b_reward_full_ultrafeedback.py文件。

调整path和dataset等相关参数。

在完成上述操作后，便可以使用下面的指令启动训练任务了。

# 单机单卡
xtuner train ./internlm2_chat_1_8b_reward_full_ultrafeedback_copy.py
# 单机多卡
NPROC_PER_NODE=${GPU_NUM} xtuner train ./internlm2_chat_1_8b_reward_full_ultrafeedback_copy.py
# slurm 集群
srun ${SRUN_ARGS} xtuner train ./internlm2_chat_1_8b_reward_full_ultrafeedback_copy.py --launcher slurm

XTuner 已经集成好了将模型转换为 HuggingFace 格式的工具，我们只需要执行：

# 创建存放 hf 格式参数的目录
mkdir work_dirs/internlm2_chat_1_8b_reward_full_ultrafeedback_copy/iter_15230_hf

# 转换格式
xtuner convert pth_to_hf internlm2_chat_1_8b_reward_full_ultrafeedback_copy.py \
work_dirs/internlm2_chat_1_8b_reward_full_ultrafeedback_copy.py/iter_15230.pth \
work_dirs/internlm2_chat_1_8b_reward_full_ultrafeedback_copy.py/iter_15230_hf

便能够将 XTuner 的 ckpt 转换为 Huggingface 格式的模型。

小结

在很多时候，比较好的选择是基于internlm等这类开源模型的rw模型继续训练；如果有条件的话，可以重新基于llama等模型训练SFT后，再训练rw模型。

rw模型对训练结果的加持是很大的，而且在很多的场景下是非常有用的，尤其是应用于打分场景，也是RLHF的一个关键环节，虽然是可以用DPO来替换，但在很多场景下还是可以应用打分模型。