在线RLHF
简要说明:这是一个用于通过在线迭代RLHF调整大型语言模型(LLMs)的仓库。还请查看我们的技术报告和Huggingface仓库!
我们展示了在线迭代人类反馈强化学习(RLHF)的工作流程,最近的LLM文献广泛报告其大大优于离线RLHF。然而,现有的开源RLHF项目仍主要局限于离线学习设置。在此仓库中,我们旨在填补这一空白,并提供一个详细的食谱,易于复现用于在线迭代RLHF。特别地,使用我们的食谱,仅凭开源数据,我们就可以获得与LLaMA3-8B-instruct相当甚至更好的结果。
模型发布
安装指南
建议分别为推理和训练创建两个独立的环境。
注意,numpy版本应为numpy<2.0
。Numpy 2.0
会遇到意外问题!!!
推理环境
conda create -n vllm python=3.10.9
conda activate vllm
pip install datasets
# 以下代码已测试在CUDA12.0-12.2。你可能需要根据自己的CUDA版本更新torch和flash-attention的来源。
pip3 install torch==2.1.2 torchvision torchaudio
pip install https://github.com/vllm-project/vllm/releases/download/v0.4.0/vllm-0.4.0-cp310-cp310-manylinux1_x86_64.whl
pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.5.7/flash_attn-2.5.7+cu122torch2.1cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
pip install accelerate==0.27.2
pip install deepspeed
训练环境
conda create -n rlhflow python=3.10.9
conda activate rlhflow
git clone https://github.com/huggingface/alignment-handbook.git
cd ./alignment-handbook/
git checkout d17fd7cd3b71c6a7bf7af34d8dc73135bb7ea8e9
pip3 install torch==2.1.2 torchvision torchaudio
python -m pip install .
pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.5.7/flash_attn-2.5.7+cu122torch2.1cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
pip install accelerate==0.27.2
你还需要安装wandb以记录训练进度,并使用你的huggingface帐户登录,以便访问LLaMA3模型。
pip install wandb
wandb login
huggingface-cli login
快速入门
我们在本节中提供了逐步指导。
第一步 监督微调
首先,你应该将你的数据集预处理为标准格式。这里有一个示例的数据集。你可能需要根据你的计算资源调整超参数(批量大小、打包大小)。要运行SFT,可以使用以下命令。
# 你可以在./sft/sft.py中调整训练参数
accelerate launch ./sft/sft.py
# 用deepspeed stage3训练
# 你可能需要根据你的环境调整./configs/zero3.yaml,特别是num_processes(GPU的数量)
accelerate launch --config_file ./configs/zero3.yaml ./sft/sft.py
第二步 奖励建模
我们将感兴趣的读者引导至这个仓库,以获得详细的配方来训练开源最先进的奖励/偏好模型。我们已经训练了几个RMs并将他们准备在huggingface上,如sfairXC/FsfairX-LLaMA3-RM-v0.1和RLHFlow/pair-preference-model-LLaMA3-8B,截至2024年5月它们是最先进的开源RMs。
第三步 数据生成
为了加速数据生成,我们使用VLLM。我们准备了两种使用VLLM进行推理的方法,以提供更稳健的实现,您可以试验并选择最适合您的环境的方式。我们以LLaMA3-8B为例。对于其他模型,您需要调整eos_ids。
您可以创建一个test_gen.sh文件,并将以下内容复制到文件中并运行 bash test_gen.sh
。
# 第一种方法:初始化4个VLLM进程并将提示集合分配给4个代理
# 生成的样本将存储在output_dir + local_index + ".json"中
my_world_size=4 # 你使用的GPU数量
infer_model=meta-llama/Meta-Llama-3-8B-Instruct
prompt_dir=RLHFlow/test_generation_2k
mkdir data
output_dir=./data/gen_data
conda activate vllm
CUDA_VISIBLE_DEVICES=0 python ./generation/get_hf2.py --model_name_or_path ${infer_model} --dataset_name_or_path ${prompt_dir} --output_dir ${output_dir} --K 4 --temperature 1.0 --local_index 0 --my_world_size ${my_world_size} --eos_ids 128009 &
CUDA_VISIBLE_DEVICES=1 python ./generation/get_hf2.py --model_name_or_path ${infer_model} --dataset_name_or_path ${prompt_dir} --output_dir ${output_dir} --K 4 --temperature 1.0 --local_index 1 --my_world_size ${my_world_size} --eos_ids 128009 &
CUDA_VISIBLE_DEVICES=2 python ./generation/get_hf2.py --model_name_or_path ${infer_model} --dataset_name_or_path ${prompt_dir} --output_dir ${output_dir} --K 4 --temperature 1.0 --local_index 2 --my_world_size ${my_world_size} --eos_ids 128009 &
CUDA_VISIBLE_DEVICES=3 python ./generation/get_hf2.py --model_name_or_path ${infer_model} --dataset_name_or_path ${prompt_dir} --output_dir ${output_dir} --K 4 --temperature 1.0 --local_index 3 --my_world_size ${my_world_size} --eos_ids 128009 &
wait
python ./generation/merge_data.py --base_path ${output_dir} --output_dir ./data/gen_data.json --num_datasets ${my_world_size}
我们也可以使用API服务器生成新的回复。
# 第一种方法:初始化4个VLLM进程并将提示集合分配给4个代理
# 生成的样本将存储在output_dir + local_index + ".json"中
my_world_size=4
infer_model=meta-llama/Meta-Llama-3-8B-Instruct
prompt_dir=RLHFlow/test_generation_2k
mkdir data
output_dir=./data/gen_data.json
conda activate vllm
# 注册API服务器
bash ./generation/run_8gpu.sh $infer_model
python ./generation/gen_hf.py --ports 8000 8001 8002 8003 8004 8005 8006 8007 --eos_ids 128009 --tokenizer $infer_model --dataset_name_or_path $prompt_dir --output_dir $output_dir --K 4 --temperature 1.0
第三步 数据标注
然后,我们调用在第二步中训练的奖励/偏好模型对生成的回复进行排名。
accelerate launch ./annotate_data/get_rewards.py --dataset_name_or_path ./data/gen_data.json --output_dir ./data/data_with_rewards.json --K 4
如果遇到TypeError: Got unsupported ScalarType BFloat16
错误,考虑pip install transformers==4.38.2
注意:遵循LLaMA2项目,当前实现假设RM与要调整的模型共享相同的聊天模板。然而,在很多情况下,RM可能有自己的聊天模板。你可以在get_rewards.py中更新change_of_format函数并启用
# 大约在第123行
test_texts = [change_of_format(sample['prompt'], tmp_output) for tmp_output in sample['responses']]
第三步 训练
conda activate rlhflow
model_path=meta-llama/Meta-Llama-3-8B-Instruct
initial_model=meta-llama/Meta-Llama-3-8B-Instruct
mkdir models
accelerate launch --config_file ./configs/zero2.yaml ./dpo_iteration/run_dpo.py --run_name rlhflow_iter1 --output_dir ./models/rlhflow_iter1 --model_name_or_path $model_path --ref_model $initial_model --learning_rate 2e-7 --max_steps 1200 --choose_type max_min --train_dir ./data/data_with_rewards.json --eval_dir ./data/data_with_rewards.json --loss_type sigmoid --lr_scheduler_type cosine
如果遇到RuntimeError: CUDA error: invalid device ordinal, CUDA kernel errors might be asynchronously reported at some other API call
错误,你需要根据你的GPU数量调整配置文件中的num_of_process。
整合一切
我们将所有步骤整合在一起,使迭代训练可以自动运行。注意我们设置了sleep 1m以等待注册推理的API。你可能需要根据你的环境调整此参数。
bash run_loop.sh
致谢
作者感谢伟大的开源社区,包括Huggingface TRL团队、Huggingface H4团队、Allen Institute AI RewardBench团队、Meta LLaMA团队和Axolotl团队,感谢他们分享模型、代码和训练集。
引用
如果你觉得该仓库的内容有帮助,请考虑如下引用:
@misc{dong2024rlhf,
title={RLHF Workflow: From Reward Modeling to Online RLHF},
author={Hanze Dong and Wei Xiong and Bo Pang and Haoxiang Wang and Han Zhao and Yingbo Zhou and Nan Jiang and Doyen Sahoo and Caiming Xiong and Tong Zhang},
year={2024},
eprint={2405.07863},
archivePrefix={arXiv},
primaryClass={cs.LG}
}
@inproceedings{xiong2023iterative,
title={Iterative preference learning from human feedback: Bridging theory and practice for RLHF under KL-constraint},
author={Xiong, Wei and Dong, Hanze and Ye, Chenlu and Wang, Ziqi and Zhong, Han and Ji, Heng and Jiang, Nan and Zhang, Tong},
booktitle={ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models}
}