本文档提供了各种人类偏好对齐算法的训练脚本。若您希望深入了解更详尽的算法信息及其选择方法,请参考文档
# 设置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
人类偏好对齐训练一般需要
其中KTO算法的数据比较特殊,只需要
训练提示:
- 如果用带有history的数据训练base模型,需要指定支持多轮对话的template(base模型往往不支持多轮对话),对于这种情况我们默认设置了
chatml
template,你也可以使用--model_type
来选择训练模型的template - 使用自定义数据集进行训练请参考自定义与拓展
- 下面的训练脚本使用
--lora_target_modules ALL
来训练模型的全部线性层,你也可以设置--lora_target_modules DEFAULT
只训练模型的QKV矩阵
超参
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推理加速与部署文档
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
文件中
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
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
训练中的指标解释参考文档
超参
- beta: KL正则系数,值越大表示对偏离参考模型的惩罚越大。默认为0.1
- desirable_weight :损失函数中的$\lambda_D$项,偏好回答样本的损失权重, 默认为1.0
- undesirable_weight :损失函数中的$\lambda_U$项,拒绝回答样本的损失权重,默认为1.0
用
训练脚本
使用
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
论文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
超参
- 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
论文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