OpenRLHF: 可以直接训练LLM、多模态大模型的RLHF框架，支持RLHF、DPO等主流算法

代码： https://github.com/OpenRLHF/OpenRLHF

这个仓库实现了各种RLHF算法，能够快速上手微调自己打LLM或多模态大模型，甚至支持MoE模型；RL算法也支持了主流的PPO、DPO。另外可以将奖励模型、策略模型分布在不同的设备上，最少4块RTX4090就能把RLHF玩起来。再配合这篇论文讲解，算是代码仓库的说明文档，非常适合入手RLHF。

https://wx.zsxq.com/mweb/views/joingroup/join_group.html?group_id=51111854414284 (二维码自动识别)

摘要

随着大型语言模型（LLMs）按照扩展规律持续增长，通过人类反馈的强化学习（RLHF）因其卓越的性能而受到广泛关注。然而，与对单个模型进行预训练或微调不同，将RLHF扩展用于训练大型语言模型面临协调四个模型的挑战。我们提出了OpenRLHF，这是一个支持高效扩展RLHF的开源框架。与现有RLHF框架将四个模型部署在同一GPU上的方式不同，OpenRLHF重新设计了针对70B以上参数模型的调度策略，结合了Ray、vLLM和DeepSpeed，从而提高了资源利用率并支持多样化的训练方法。

OpenRLHF与Hugging Face无缝集成，提供了优化的算法和启动脚本，实现了开箱即用的用户友好体验。OpenRLHF支持RLHF、DPO、拒绝采样以及其他对齐技术，为开发最先进的LLM提供了强大助力。

OpenRLHF 的代码可在以下网址获取：https://github.com/OpenRLHF/OpenRLHF。

1 引言

尽管大型语言模型（LLM）遵循扩展规律展现了显著的性能提升，但随着模型规模的增大，将这些模型与人类价值和意图对齐成为一项日益严峻的挑战。通过人类反馈的强化学习（RLHF） [19] 已成为解决这一长期难题的有力技术。然而，随着模型规模的扩大，传统的RLHF通常需要维护多个模型并依赖更复杂的学习流程，从而对内存和计算资源提出了更高的需求。例如，近端策略优化（Proximal Policy Optimization, PPO） [23, 19] 是RLHF中常用的一种算法，在训练过程中需要同时维护四个模型。因此，当语言模型的参数规模超过70亿时，训练和协调多个模型所需的计算资源和调度复杂性显著增加，对现有框架设计提出了新的需求和挑战。

现有开源的RLHF框架（如Transformer Reinforcement Learning (TRL)、ColossalChat (CAIChat) 和 DeepSpeed-Chat (DSChat)）依赖于诸如零冗余优化（ZeRO） [15, 30, 21] 等并行化方法，将RLHF训练中涉及的四个模型共同部署在同一块GPU上。然而，随着模型规模持续突破70亿参数，这种调度方式在GPU内存受限的情况下变得效率低下。为了应对共置方法的局限性，一些框架（如TRL）通过合并actor和critic模型或采用低秩适配（Low-Rank Adaptation, LoRA）[11]等技术来减少内存占用。然而，这些方法可能会降低模型性能，且合并后的actor-critic架构与推荐的基于奖励模型权重初始化critic模型的方法 [19] 不兼容。

另一种针对大规模模型的解决方案是利用NVIDIA Megatron [26] 提供的张量并行和流水线并行技术。然而，Megatron与流行的Hugging Face 库 [29] 不兼容，且适配新模型需要对源代码进行大量修改，从而影响了其可用性。

为实现大规模的便捷RLHF训练，OpenRLHF 采用了Ray [18]、vLLM [14] 和 DeepSpeed [22] 重新设计模型调度方案，使得训练规模突破70亿参数的模型成为可能。OpenRLHF与Hugging Face Transformers [29] 无缝集成，并支持多种流行技术，如专家混合模型（Mixture of Experts, MoE）[13]、Jamba [17] 和QLoRA [4]。此外，OpenRLHF还实现了多种对齐算法，包括直接偏好优化（Direct Preference Optimization, DPO）[20]、Kahneman-Tversky优化（KTO）[10]、条件SFT（Conditional SFT）[12] 和拒绝采样（Rejection Sampling）[28]，为用户提供了全面且易用的RLHF训练框架。表1 对比了几种流行的RLHF框架。

2 背景

2.1 人类反馈强化学习（RLHF）

经典的大型语言模型训练方法基于预训练的生成式预训练Transformer（Generative Pre-trained Transformer, GPT） [19]，通常包括以下三个步骤：监督微调（SFT）、奖励模型（RM）训练 和 PPO训练：

监督微调（Supervised Fine-tuning, SFT）
开发者使用标注者提供的人类演示数据，通过监督学习的损失函数对GPT模型进行微调，其公式如（式1）所示：

奖励模型训练（Reward Model training）

以去掉最终解码层的SFT模型为基础，训练奖励模型以接收提示和回复，并输出一个标量奖励。奖励模型的损失函数如（式2）所示：

PPO训练（Proximal Policy Optimization, PPO Training）

在一个“上下文-回复”带宽环境中，开发者使用PPO算法对语言模型进行微调。环境根据提示和回复对生成的结果给出奖励，并结束回合。为防止对奖励模型的过度优化，在每个标记上添加了基于监督微调模型的每标记KL散度惩罚（Kullback-Leibler divergence）。通过奖励模型权重初始化价值函数，为强化学习（RL）微调提供稳定的起点。PPO的损失函数如（式3）所示：

2.2 Ray

Ray [18] 是一个分布式执行框架，为并行和分布式计算任务提供了强大的调度和扩展能力。它使用内置的分布式调度器，高效地将任务分配到集群中的可用资源上，从而支持从单机到数千个节点的大规模部署无缝扩展。Ray 的调度机制智能地管理任务并行性，将计算分解为可并发执行的小任务，分布在多个核心和机器上执行。凭借其可扩展架构和高效的调度能力，Ray 非常适合加速机器学习、科学计算以及高性能数据处理管道等数据密集型工作负载。Ray 为并行处理提供计算层支持，使用户无需成为分布式系统的专家。

2.3 vLLM

vLLM [14] 是一个快速且易于使用的库，用于LLM推理和服务。通过高效管理注意力键和值的内存（PagedAttention）、对传入请求的连续批处理，以及使用CUDA图实现快速模型执行，vLLM 提供了业界领先的服务吞吐量。vLLM 的灵活性和易用性体现在其与流行的 Hugging Face 模型的无缝集成、高吞吐量服务支持多种解码算法、分布式推理的张量并行支持以及流式输出等方面。它还支持实验性功能，如前缀缓存和多-LoRA支持。vLLM 可无缝支持 Hugging Face 上最受欢迎的开源模型，包括类似Transformer的LLM（如Llama）和专家混合LLM（如Mixtral [13]）。

2.4 DeepSpeed

DeepSpeed [22] 是一个优化库，旨在提高大规模深度学习模型的训练效率。其 零冗余优化器（Zero Redundancy Optimizer, ZeRO） [21] 通过在数据并行进程之间划分模型状态、梯度和优化器状态，大幅减少了内存消耗，从而支持训练拥有数万亿参数的模型。此外，DeepSpeed 的数据卸载功能可在CPU和GPU内存之间无缝传输数据，从而进一步优化资源利用效率，使硬件GPU内存受限的情况下也能高效训练大规模模型。DeepSpeed 同样可以无缝支持 Hugging Face 上最受欢迎的开源模型。

3 OpenRLHF 的设计

3.1 调度优化

将RLHF训练扩展到更大模型需要高效地将至少四个组件模型（actor、critic、reward 和 reference）分配到多个GPU上，因为每个加速器的内存限制（例如，NVIDIA A100 的显存限制为80GB）。OpenRLHF 在模型调度上进行了创新，利用 Ray [18] 实现模型的分布式部署和精细的任务编排。同时，基于Ray的调度器管理针对推理优化的库（如vLLM [14]）以及针对训练优化的库（如DeepSpeed）。OpenRLHF 将四个模型分布到多个GPU上，而不是将它们共置于同一GPU上（如图1所示）。

这种设计在RLHF训练过程中自然支持多奖励模型（如图2所示），以实现不同算法的实现选择。例如，算法工程师可以快速构建多种对齐策略，如将“有用性”和“有害性”分离，而无需关心底层数据流的具体细节。

通过结合Ray和DeepSpeed，调度器支持灵活的模型合并或卸载策略。例如，可以合并actor-reference模型或critic-reward模型以节省GPU资源。除了高度可定制的算法实现优势外，调度器还通过优化GPU的编排显著提升整体训练性能。更多细节将在后续章节中讨论，但调度优化是实现进一步效率提升的基石。

3.2 性能优化

RLHF算法的性能依赖于训练和推理的效率。通过使用LLaMA2 7B和NVIDIA A100进行性能分析（如图4a所示），发现主要的瓶颈位于PPO样本生成阶段，其占据了总训练时间的80%。这是因为在生成阶段，自回归解码的复杂度为，且受到内存带宽的限制。图4b显示，增大推理批量大小可以显著提高生成吞吐量。然而，像DeepSpeedChat和TRL这样的RLHF框架，由于将所有模型共享同一GPU，导致生成阶段可用内存不足，无法增大批量大小，从而加剧了低内存访问效率的问题。

OpenRLHF 使用Ray将四个模型分布在多个GPU上，有效缓解了这一问题。此外，为加速样本生成并支持无法在单个GPU上加载的更大LLM（如70B模型），OpenRLHF 借助vLLM的张量并行和其他先进技术（如连续批处理和分页注意力 [14]），如图1所示。

在RLHF学习阶段，OpenRLHF还应用了以下技术进行改进（见图3）：

将Adam优化器状态卸载到CPU，释放GPU内存，使得在生成（无vLLM）和训练期间支持更大的批量大小。这种方法提高了计算效率并减少了ZeRO通信开销。为缓解梯度聚合期间GPU与CPU通信的开销，采用了固定内存和梯度累积技术。
使用Flash Attention 2 [3] 加速Transformer模型的训练。
通过PyTorch张量切片移除训练样本中的冗余填充。

如图2所示的其他三个模型使用ZeRO阶段3（对模型、梯度和优化器进行分片）。OpenRLHF通过NVIDIA NCCL和vLLM权重加载器同步ZeRO与vLLM引擎的权重，确保了快速且简便的集成。在第4.1节中，我们将OpenRLHF与精心调优的DSChat性能进行对比。

3.3 PPO 实现技巧

对于大型语言模型（LLM）的训练，像PPO这样的强化学习算法容易出现不稳定问题。我们尽最大努力验证了实现细节，其推理和学习的通用流程分别在图2和图3中展示。此外，OpenRLHF 在PPO实现中采用了多种技巧以稳定训练 [8]：

仅在序列的文本结束标记（end-of-text token）上预测奖励。
对语言模型使用标记级别的强化学习。
在PPO中使用KL散度损失项。
在PPO中使用基于策略损失相对比例调整的预训练损失项。
应用奖励归一化以稳定训练。
使用具有全局统计的分布式优势归一化。
使用线性预热余弦退火学习率调度器。
用奖励模型的权重初始化Critic。
Actor使用较低的学习率，而Critic使用较高的学习率。
在初始学习阶段冻结Actor的权重，以更好地初始化Critic。
使用GAE（广义优势估计）。

更多细节请参考[博客][25]。

3.4 易用性

为提升用户友好性，OpenRLHF 提供了一键可训练脚本，支持的算法完全兼容Hugging Face库，用户可以指定模型和数据集的名称或路径。以下是一个在16个A100 GPU上对70B模型进行RLHF训练的配置示例：

{
  "model_name": "LLaMA-70B",
  "dataset_path": "/path/to/dataset",
  "num_gpus": 16,
  "batch_size": 64,
  "learning_rate": 3e-5,
  "training_steps": 10000
}

我们为所有支持的算法提供了使用脚本（见附录A），详细信息可在 OpenRLHF 使用文档中找到。

4 实验

4.1 性能基准测试

表3和表4展示了RLHF性能实验中使用的配置。我们使用NVIDIA A800 GPU和NVLINK进行数据传输，对LLaMA2模型进行了训练。为了尽可能优化DSChat的性能，我们采用了一些技术，如启用Adam优化器卸载以及对奖励模型（RM）和参考模型（Ref）进行卸载，以增加推理阶段的微批量大小并避免内存不足问题。此外，我们修复了DSChat中的一些Bug，使其能够在LLaMA2上启用Hybrid Engine（HE）。

表5和表6分别展示了OpenRLHF和优化后的DSChat在PPO各阶段的详细耗时。尽管OpenRLHF的PPO性能尚未完全优化（例如，将Actor与Reference模型节点合并，以及将Critic与Reward模型节点合并以减少GPU资源的使用），实验结果仍表明，OpenRLHF在表2中的性能具有显著优势。

性能优势分析

OpenRLHF相较于DSChat的性能优势主要来源于vLLM和Ray的应用：

vLLM的生成加速性能优于Hybrid Engine。
Ray的分布式模型部署将模型分布到不同的节点上，这看似会降低GPU利用率，但却避免了过度的模型分片和模型权重卸载，从而节省了GPU内存并减少了通信开销。这种方式使得每个GPU上的微批量大小和张量并行中的矩阵乘法规模得以增大，从而提升了整体性能。

尽管OpenRLHF的某些方面（如PPO性能）未完全优化，实验结果仍然证明了其性能相较于DSChat的显著优势。

4.2 训练稳定性与收敛性

我们基于LLaMA2 7B模型评估了OpenRLHF框架的训练收敛性。在监督微调（SFT）阶段，我们使用了来自OpenOrca的50k数据集。奖励模型的训练偏好数据集则由Anthropic HH [3]、LMSys Arena [4] 和 Open Assistant [5] 数据集中的约200k样本混合构成。**近端策略优化（PPO）**训练随机从上述数据集中采样了80k条提示，**直接策略优化（DPO）**训练使用与奖励模型训练相同的数据集。