人类偏好对齐训练文档

本文档提供了各种人类偏好对齐算法的训练脚本。若您希望深入了解更详尽的算法信息及其选择方法，请参考文档

目录

环境准备
数据集
DPO
RM
PPO
KTO
CPO
ORPO
SimPO

环境准备

# 设置pip全局镜像 (加速下载)
pip config set global.index-url https://mirrors.aliyun.com/pypi/simple/
# 安装ms-swift
git clone https://github.com/modelscope/swift.git
cd swift
pip install -e '.[llm]'

# 环境对齐 (通常不需要运行. 如果你运行错误, 可以跑下面的代码, 仓库使用最新环境测试)
pip install -r requirements/framework.txt  -U
pip install -r requirements/llm.txt  -U

数据集

人类偏好对齐训练一般需要 $(x,y_w,y_l)$ 格式的数据，其中 $x$ 表示模型输入，$y_w,y_l$ 分别表示符合人类偏好的偏好回答和不符合人类偏好的拒绝回答,比如 dpo_data

其中KTO算法的数据比较特殊，只需要 $(x,y,\text{label})$ 格式的数据，其中 $x$ 表示模型输入，$y$ 表示模型输出，label表示回答是否符合人类偏好比如 kto_data

训练提示:

如果用带有history的数据训练base模型，需要指定支持多轮对话的template(base模型往往不支持多轮对话)，对于这种情况我们默认设置了chatmltemplate，你也可以使用--model_type 来选择训练模型的template
使用自定义数据集进行训练请参考自定义与拓展
下面的训练脚本使用--lora_target_modules ALL来训练模型的全部线性层，你也可以设置--lora_target_modules DEFAULT只训练模型的QKV矩阵

DPO

论文arvix

超参

beta：KL正则系数，值越大表示对偏离参考模型的惩罚越大。默认为0.1

建议在开始DPO训练之前，使用偏好数据集中的偏好回答部分进行SFT训练，以确保数据符合DPO算法的分布要求。我们也在DPO loss中混合了sft loss来稳定训练，你可以通过设置超参rpo_alpha来调整sft loss的系数，默认为1..

训练脚本, 这里我们提供单卡/多卡device map/多卡ddp的版本，简洁起见，后续算法只给出单卡版本。

# Experimental environment: A100
# Memory usage: 40G
CUDA_VISIBLE_DEVICES=0 \
swift rlhf \
    --rlhf_type dpo \
    --model_type  llama3-8b-instruct \
    --beta 0.1 \
    --rpo_alpha 0.1 \
    --sft_type  lora \
    --dataset shareai-llama3-dpo-zh-en-emoji \
    --num_train_epochs  2  \
    --lora_target_modules  ALL  \
    --gradient_checkpointing  true  \
    --batch_size  1  \
    --learning_rate  5e-5  \
    --gradient_accumulation_steps  16  \
    --warmup_ratio  0.03  \
    --save_total_limit  2

# MP(device map)
# Memory usage: 2*24G
CUDA_VISIBLE_DEVICES=0,1 \
swift rlhf \
    --rlhf_type dpo \
    --model_type  llama3-8b-instruct \
    --beta 0.1 \
    --rpo_alpha 0.1 \
    --sft_type  lora \
    --dataset shareai-llama3-dpo-zh-en-emoji \
    --num_train_epochs  2  \
    --lora_target_modules  ALL  \
    --gradient_checkpointing  true  \
    --batch_size  1  \
    --learning_rate  5e-5  \
    --gradient_accumulation_steps  16  \
    --warmup_ratio  0.03  \
    --save_total_limit  2

# DDP + MP
# Memory usage: 4*24G
nproc_per_node=2

CUDA_VISIBLE_DEVICES=0,1,2,3 \
NPROC_PER_NODE=$nproc_per_node \
MASTER_PORT=29500 \
swift rlhf \
    --rlhf_type dpo \
    --model_type  llama3-8b-instruct \
    --beta 0.1 \
    --rpo_alpha 0.1 \
    --sft_type  lora \
    --dataset shareai-llama3-dpo-zh-en-emoji \
    --num_train_epochs  2  \
    --lora_target_modules  ALL  \
    --gradient_checkpointing  true  \
    --batch_size  1  \
    --learning_rate  5e-5  \
    --gradient_accumulation_steps  $(expr 16 / $nproc_per_node)  \
    --warmup_ratio  0.03  \
    --save_total_limit  2

训练后的模型推理和部署可以参考LLM推理文档和VLLM推理加速与部署文档

RM

论文arvix

RLHF中的Reward Modeling阶段

使用sft训练后的base model或者instruct model作为基底模型, 增加value head, 使用偏好数据集训练为reward model

CUDA_VISIBLE_DEVICES=0 \
swift rlhf \
    --rlhf_type rm \
    --model_type  llama3-8b-instruct \
    --sft_type  lora \
    --dataset hh-rlhf-cn-harmless-base-cn \
    --num_train_epochs  2  \
    --lora_target_modules  ALL  \
    --gradient_checkpointing  true  \
    --batch_size  1  \
    --learning_rate  5e-5  \
    --gradient_accumulation_steps  16  \
    --warmup_ratio  0.03  \
    --save_total_limit  2

增加的value head权重会保存在value_head.safetensors 或 value_head.bin文件中

PPO

论文arvix

RLHF中的PPO(proximal policy optimization)阶段, 涉及到四个模型

model: 训练模型, sft训练后的base model或者instruct model
ref_model: 参考模型, 默认为 model
reward_model: 奖励模型, 由RM阶段训练得到
value_model: 价值模型, 由reward_model初始化, 在训练中同步更新

超参

local_rollout_forward_batch_size: 每次数据采样的批量大小, 默认为64
whiten_rewards: 对奖励进行归一化处理, 默认为False
kl_coef: KL散度项的系数, 默认为0.05
cliprange: PPO策略损失函数中的clip范围, 默认为0.2
vf_coef: 价值损失函数系数, 默认为0.1
cliprange_value: PPO价值损失函数中的clip范围, 默认为0.2
gamma: 累计奖励的折扣因子, 默认为1.0
lam: GAE中的lambda系数, 默认为0.95
num_sample_generations: 训练过程中生成的调试样本数量, 默认为10

CUDA_VISIBLE_DEVICES=0 \
swift rlhf \
    --rlhf_type ppo \
    --model_type  llama3-8b-instruct \
    --sft_type  lora \
    --dataset hh-rlhf-cn-harmless-base-cn \
    --reward_model_id_or_path path/to/reawrd_model \
    --reward_model_type  llama3-8b-instruct \
    --num_train_epochs  2  \
    --lora_target_modules  ALL  \
    --gradient_checkpointing  true  \
    --batch_size  1  \
    --learning_rate  5e-5  \
    --gradient_accumulation_steps  16  \
    --warmup_ratio  0.03  \
    --save_total_limit  2

注意: 训练base model时, 需要先sft后再进行rlhf, 指定chat template, sft_type建议使用full

训练中的指标解释参考文档

KTO

论文arvix

超参

beta： KL正则系数，值越大表示对偏离参考模型的惩罚越大。默认为0.1
desirable_weight ：损失函数中的$\lambda_D$项，偏好回答样本的损失权重, 默认为1.0
undesirable_weight ：损失函数中的$\lambda_U$项，拒绝回答样本的损失权重，默认为1.0

用 $n_D$ 和 $n_U$ 分别表示数据集中偏好回答和拒绝回答的样本数量，对于超参 $\lambda_D$ 和 $\lambda_U$ ，作者推荐设置 $\frac{\lambda_Dn_D}{\lambda_Un_U}\in[1,\frac{4}{3}]$

训练脚本使用 $(x,y,\text{label})$ 格式数据训练

CUDA_VISIBLE_DEVICES=0 \
swift rlhf \
    --rlhf_type kto \
    --model_type  llama3-8b-instruct \
    --beta 0.1 \
    --desirable_weight 1.0 \
    --undesirable_weight 1.0 \
    --sft_type  lora \
    --dataset ultrafeedback-kto \
    --num_train_epochs  2  \
    --lora_target_modules  ALL  \
    --gradient_checkpointing  true  \
    --batch_size  1  \
    --learning_rate  5e-5  \
    --gradient_accumulation_steps  16  \
    --warmup_ratio  0.03  \
    --save_total_limit  2

CPO

论文arvix 超参

beta：隐含奖励前的系数，默认为0.1
cpo_alpha: nll loss系数, 默认为1.0

训练脚本

CUDA_VISIBLE_DEVICES=0 \
swift rlhf \
    --rlhf_type cpo \
    --model_type  llama3-8b-instruct \
    --beta 0.1 \
    --sft_type  lora \
    --dataset shareai-llama3-dpo-zh-en-emoji \
    --num_train_epochs  2  \
    --lora_target_modules  ALL  \
    --gradient_checkpointing  true  \
    --batch_size  1  \
    --learning_rate  5e-5  \
    --gradient_accumulation_steps  16  \
    --warmup_ratio  0.03  \
    --save_total_limit  2

ORPO

论文arvix

超参

lambda: Odds Ratio loss系数

注意：ORPO使用参数--beta传入超参lambda

CUDA_VISIBLE_DEVICES=0 \
swift rlhf \
    --rlhf_type orpo \
    --model_type  llama3-8b-instruct \
    --beta 0.1 \
    --sft_type  lora \
    --dataset shareai-llama3-dpo-zh-en-emoji \
    --num_train_epochs  2  \
    --lora_target_modules  ALL  \
    --gradient_checkpointing  true  \
    --batch_size  1  \
    --learning_rate  5e-5  \
    --gradient_accumulation_steps  16  \
    --warmup_ratio  0.03  \
    --save_total_limit  2

SimPO

论文arvix 超参

beta：隐含奖励前的系数，默认为2.0
simpo_gamma：reward margin项，默认为1.0
cpo_alpha: 混合CPO nll loss提高训练稳定性, 默认为1.0, 设置0.0使用原始SimPO算法

CUDA_VISIBLE_DEVICES=0 \
swift rlhf \
    --rlhf_type simpo \
    --model_type  llama3-8b-instruct \
    --beta 2.0 \
    --simpo_gamma 1.0 \
    --cpo_alpha 1.0 \
    --sft_type  lora \
    --dataset shareai-llama3-dpo-zh-en-emoji \
    --num_train_epochs  2  \
    --lora_target_modules  ALL  \
    --gradient_checkpointing  true  \
    --batch_size  1  \
    --learning_rate  5e-5  \
    --gradient_accumulation_steps  16  \
    --warmup_ratio  0.03  \
    --save_total_limit  2