OpenRLHF: 一个易使用的大规模RLHF训练框架
官方 GitHub | 官方文档 初七 (hijkzzz) OpenRLHF Llamafia 编辑推荐理由：RLHF(基于人类反馈的强化学习)是当前大语言模型对齐的主流方法,但其训练过程涉及多个模型的协调,在大规模模型上实现高效训练具有挑战性。OpenRLHF提供了一个开源、易用且高性能的RLHF训练框架,支持70B以上参数规模的模型训练。该框架不仅优化了训练性能,还实现了多种对齐算法,为大模型研究人员提供了强大的工具。 随着大语言模型(LLM)规模的不断扩大,基于人类反馈的强化学习(RLHF)因其出色的性能而备受关注。然而,与预训练或微调单个模型不同,RLHF训练涉及四个模型的协调,在大规模模型上实现高效训练具有挑战性。本文介绍了OpenRLHF,这是一个开源框架,能够实现高效的RLHF扩展训练。 OpenRLHF的主要特点 创新的调度设计: 不同于TRL等框架将四个模型放置在同一GPU上,OpenRLHF利用Ray、vLLM和DeepSpeed重新设计了70B以上参数模型的调度方案,提高了资源利用率。 性能优化: OpenRLHF在生成和训练阶段都进行了优化,如使用vLLM加速样本生成,使用Flash Attention 2加速Transformer模型训练等。 算法实现: 实现了SFT-packing, PPO、DPO、KTO、Iterative DPO、拒绝采样、CSFT等多种对齐技术。 易用性: 与Hugging Face模型库无缝集成,提供开箱即用的解决方案, 包括优化的算法和启动脚本的文档。 灵活性: 支持流行技术如Mixture of Experts (MoE)和LoRA/QLoRA。 设计细节 调度优化 OpenRLHF利用Ray进行模型放置和精细编排,管理vLLM等推理优化库和DeepSpeed等训练优化库。它将四个模型分布在多个GPU上,而不是将它们放在同一个GPU上。这种设计自然支持RLHF训练过程中的多奖励模型。 性能优化 性能优化主要体现在两个阶段: 生成阶段: 我们在 prompt 最大长度 1k 生成最大长度 1k 的设定下测试 RLHF 训练的性能瓶颈，发现 PPO 主要耗时在生成阶段占比约 80% 其次我们通过Ray分布式的放置四个模型(Actor,Critic,Reference,Reward), 并支持自由模型节点合并复用, 使得有更多的GPU内存来增大推理 Batch Size 避免生成阶段的 Memory Bound，极大提升效率 并且使用vLLM的张量并行和其他高级技术加速进一步样本生成，并支持 70B 模型的训练 学习阶段: 将Adam优化器状态卸载到CPU, 其节省GPU内存的意义和生成阶段一致 使用 Flash Attention 2 加速Transformer模型训练 使用 PyTorch 张量切片移除训练样本中的冗余 padding 总结：牺牲训练效率提升推理效率，提升整体性能收益...