OpenRLHF: 一个易使用的大规模RLHF训练框架

官方 GitHub | 官方文档

初七 (hijkzzz) OpenRLHF

Llamafia 编辑推荐理由：RLHF(基于人类反馈的强化学习)是当前大语言模型对齐的主流方法,但其训练过程涉及多个模型的协调,在大规模模型上实现高效训练具有挑战性。OpenRLHF提供了一个开源、易用且高性能的RLHF训练框架,支持70B以上参数规模的模型训练。该框架不仅优化了训练性能,还实现了多种对齐算法,为大模型研究人员提供了强大的工具。

随着大语言模型(LLM)规模的不断扩大,基于人类反馈的强化学习(RLHF)因其出色的性能而备受关注。然而,与预训练或微调单个模型不同,RLHF训练涉及四个模型的协调,在大规模模型上实现高效训练具有挑战性。本文介绍了OpenRLHF,这是一个开源框架,能够实现高效的RLHF扩展训练。

OpenRLHF的主要特点

1. 创新的调度设计: 不同于TRL等框架将四个模型放置在同一GPU上,OpenRLHF利用Ray、vLLM和DeepSpeed重新设计了70B以上参数模型的调度方案,提高了资源利用率。

2. 性能优化: OpenRLHF在生成和训练阶段都进行了优化,如使用vLLM加速样本生成,使用Flash Attention 2加速Transformer模型训练等。

3. 算法实现: 实现了SFT-packing, PPO、DPO、KTO、Iterative DPO、拒绝采样、CSFT等多种对齐技术。

4. 易用性: 与Hugging Face模型库无缝集成,提供开箱即用的解决方案, 包括优化的算法和启动脚本的文档。

5. 灵活性: 支持流行技术如Mixture of Experts (MoE)和LoRA/QLoRA。

设计细节

调度优化

OpenRLHF利用Ray进行模型放置和精细编排,管理vLLM等推理优化库和DeepSpeed等训练优化库。它将四个模型分布在多个GPU上,而不是将它们放在同一个GPU上。这种设计自然支持RLHF训练过程中的多奖励模型。

性能优化

性能优化主要体现在两个阶段:

1. 生成阶段:

   • 我们在 prompt 最大长度 1k 生成最大长度 1k 的设定下测试 RLHF 训练的性能瓶颈，发现 PPO 主要耗时在生成阶段占比约 80%
   • 其次我们通过Ray分布式的放置四个模型(Actor,Critic,Reference,Reward), 并支持自由模型节点合并复用, 使得有更多的GPU内存来增大推理 Batch Size 避免生成阶段的 Memory Bound，极大提升效率
   • 并且使用vLLM的张量并行和其他高级技术加速进一步样本生成，并支持 70B 模型的训练

2. 学习阶段:

   • 将Adam优化器状态卸载到CPU, 其节省GPU内存的意义和生成阶段一致
   • 使用 Flash Attention 2 加速Transformer模型训练
   • 使用 PyTorch 张量切片移除训练样本中的冗余 padding

总结：牺牲训练效率提升推理效率，提升整体性能收益

PPO实现技巧

为了稳定大语言模型的PPO训练, OpenRLHF采用了多种技巧,如:

• 仅在序列的结束标记上预测奖励
• 使用基于token-level的强化学习
• 在PPO中使用KL散度损失项
• 应用奖励归一化以提高训练稳定性
• 使用线性预热余弦退火学习率调度器
• 用奖励模型的权重初始化 Critic
• 较低的 Actor 学习速率，而 Critic 学习速率更高
• 学习初始阶段冻结 Actor 的权重，更好的初始化 Critic
• 使用 GAE + Advantage Normalization
• 使用分开的 Actor 和 Critic 模型

更多的细节在 Zhihu Blog 和 Notion Blog

实验结果

性能对比

OpenRLHF与优化后的DSChat在不同规模模型上的性能对比:

表格展示了训练1024个提示词,进行1个PPO epoch所需的平均时间。结果显示,OpenRLHF在各种规模的模型上都实现了显著的性能提升。

训练稳定性和收敛性

实验基于LLaMA3-8B, 结果显示:

• PPO训练曲线稳定,奖励和回报值稳步上升,KL散度和损失值保持稳定。
• 在Chat-Arena-Hard评估中，RLHF模型显著优于SFT模型

Chat-Arena-Hard
-------------------------------------------
llama-3-8b-sft                 | score: 5.6   
llama-3-8b-rlhf-100k           | score: 20.5

实验细节请参考 https://huggingface.co/OpenRLHF/Llama-3-8b-rlhf-100k

Ray RLHF 快速上手

# Launch the docker container (可选的)
docker run --runtime=nvidia -it --rm --shm-size="10g" --cap-add=SYS_ADMIN -v $PWD:/openrlhf nvcr.io/nvidia/pytorch:24.02-py3 bash

# pip install
pip install openrlhf[vllm]

# launch the master node of ray
ray start --head --node-ip-address 0.0.0.0 --num-gpus 8

# launch ray job
ray job submit --address="http://127.0.0.1:8265" \
   --runtime-env-json='{"working_dir": "/openrlhf"}' \
   -- python3 -m openrlhf.cli.train_ppo_ray \
   --ref_num_nodes 1 \
   --ref_num_gpus_per_node 2 \
   --reward_num_nodes 1 \
   --reward_num_gpus_per_node 2 \
   --critic_num_nodes 1 \
   --critic_num_gpus_per_node 2 \
   --actor_num_nodes 1 \
   --actor_num_gpus_per_node 2 \
   --vllm_num_engines 2 \
   --vllm_tensor_parallel_size 2 \
   --colocate_critic_reward \
   --colocate_actor_ref \
   --ref_reward_offload \
   --pretrain OpenRLHF/Llama-3-8b-sft-mixture \
   --reward_pretrain OpenRLHF/Llama-3-8b-rm-mixture \
   --save_path /openrlhf/examples/checkpoint/llama3-8b-rlhf \
   --micro_train_batch_size 8 \
   --train_batch_size 128 \
   --micro_rollout_batch_size 16 \
   --rollout_batch_size 1024 \
   --max_samples 100000 \
   --max_epochs 1 \
   --prompt_max_len 1024 \
   --generate_max_len 1024 \
   --zero_stage 3 \
   --bf16 \
   --actor_learning_rate 5e-7 \
   --critic_learning_rate 9e-6 \
   --init_kl_coef 0.01 \
   --prompt_data OpenRLHF/prompt-collection-v0.1 \
   --input_key context_messages \
   --apply_chat_template \
   --normalize_reward \
   --adam_offload \
   --flash_attn \
   --gradient_checkpointing \
   --use_wandb {wandb_token}