kai119's picture
Upload folder using huggingface_hub
22fb4ec verified

A newer version of the Streamlit SDK is available: 1.43.2

Upgrade

Reward Model 快速上手

在本章节中,我们将介绍如何使用 XTuner 训练 1.8B 的 Reward Model,以帮助您快速上手。

准备预训练模型权重

依据 Training language models to follow instructions with human feedback 论文中的描述,我们使用进过 SFT 的语言模型作为 Reward Model 的初始化模型。这里我们使用InternLM2-chat-1.8b-sft作为初始化模型。

在训练配置文件中设置pretrained_model_name_or_path = 'internlm/internlm2-chat-1_8b-sft',则会在启动训练时自动下载模型文件。若您需要手动下载模型权重,那么请参考准备预训练模型权重章节,其中详细说明了如何从 Huggingface 或者是 Modelscope 下载模型权重的方法。这里我们附上模型的 HuggingFace 链接与 ModelScope 链接:

准备训练数据

在本教程中使用 UltraFeedback 数据集作为演示,为了方便起见,我们使用 huggingface 上已经预处理过的 argilla/ultrafeedback-binarized-preferences-cleaned 数据集,

train_dataset = dict(
    type=build_preference_dataset,
    dataset=dict(
        type=load_dataset,
        path='argilla/ultrafeedback-binarized-preferences-cleaned'),
    dataset_map_fn=orpo_dpo_mix_40k_map_fn,
    is_dpo=False,
    is_reward=True,
)

在配置文件中使用以上配置,即可自动下载并处理该数据集。如果您希望使用其他 huggingface 上的开源数据集或是使用自定义的数据集,请参阅偏好数据集章节。

准备配置文件

XTuner 提供了多个开箱即用的配置文件,可以通过 xtuner list-cfg 查看。我们执行如下指令,以复制一个配置文件到当前目录。

xtuner copy-cfg internlm2_chat_1_8b_reward_full_ultrafeedback .

打开复制后的配置文件,如果您选择自动下载模型和数据集,则无需修改配置。若您希望填入您预先下载的模型路径和数据集路径,请修改配置中的 pretrained_model_name_or_path 以及 train_datasetdatasetpath 参数。

更多的训练参数配置,请参阅修改Reward训练配置章节。

启动训练

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

# 单机单卡
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

正确的训练日志应当如下所示(在单卡 A800 上运行):

06/06 16:12:11 - mmengine - INFO - Iter(train) [   10/15230]  lr: 3.9580e-07  eta: 2:59:41  time: 0.7084  data_time: 0.0044  memory: 18021  loss: 0.6270  acc: 0.0000  chosen_score_mean: 0.0000  rejected_score_mean: 0.0000  num_samples: 4.0000  num_tokens: 969.0000
06/06 16:12:17 - mmengine - INFO - Iter(train) [   20/15230]  lr: 8.3536e-07  eta: 2:45:25  time: 0.5968  data_time: 0.0034  memory: 42180  loss: 0.6270  acc: 0.5000  chosen_score_mean: 0.0013  rejected_score_mean: 0.0010  num_samples: 4.0000  num_tokens: 1405.0000
06/06 16:12:22 - mmengine - INFO - Iter(train) [   30/15230]  lr: 1.2749e-06  eta: 2:37:18  time: 0.5578  data_time: 0.0024  memory: 32121  loss: 0.6270  acc: 0.7500  chosen_score_mean: 0.0016  rejected_score_mean: 0.0011  num_samples: 4.0000  num_tokens: 932.0000
06/06 16:12:28 - mmengine - INFO - Iter(train) [   40/15230]  lr: 1.7145e-06  eta: 2:36:05  time: 0.6033  data_time: 0.0025  memory: 42186  loss: 0.6270  acc: 0.7500  chosen_score_mean: 0.0027  rejected_score_mean: 0.0016  num_samples: 4.0000  num_tokens: 994.0000
06/06 16:12:35 - mmengine - INFO - Iter(train) [   50/15230]  lr: 2.1540e-06  eta: 2:41:03  time: 0.7166  data_time: 0.0027  memory: 42186  loss: 0.6278  acc: 0.5000  chosen_score_mean: 0.0031  rejected_score_mean: 0.0032  num_samples: 4.0000  num_tokens: 2049.0000
06/06 16:12:40 - mmengine - INFO - Iter(train) [   60/15230]  lr: 2.5936e-06  eta: 2:33:37  time: 0.4627  data_time: 0.0023  memory: 30238  loss: 0.6262  acc: 1.0000  chosen_score_mean: 0.0057  rejected_score_mean: 0.0030  num_samples: 4.0000  num_tokens: 992.0000
06/06 16:12:46 - mmengine - INFO - Iter(train) [   70/15230]  lr: 3.0331e-06  eta: 2:33:18  time: 0.6018  data_time: 0.0025  memory: 42186  loss: 0.6247  acc: 0.7500  chosen_score_mean: 0.0117  rejected_score_mean: 0.0055  num_samples: 4.0000  num_tokens: 815.0000

模型转换

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 格式的模型。

需要注意的是,由于 Reward Model 的类型并未在 transformers 官方库中集成,因此目前只有InternLM2模型训练得到的 Reward Model 会被转换为 InternLM2ForRewardModel 类型,而其他模型则会默认转换为 SequenceClassification 类型(例如 LLaMa3 会被转换为 LlamaForSequenceClassification 类型),但这并不影响其在 XTuner PPO 训练中的使用。