作者/机构: Jian Hu, Xibin Wu, Wei Shen, Jason Klein Liu, Zilin Zhu Weixun Wang, Songlin Jiang, Haoran Wang, Hao Chen, Bin Chen, Weikai Fang, Xianyu, Yu Cao, Haotian Xu, Yiming Liu* (Team janhu9527@gmail.com)

A1 主要贡献

OpenRLHF 旨在解决现有强化学习从人类反馈(RLHF)和从可验证奖励中强化学习(RLVR)框架中普遍存在的推理瓶颈和复杂性障碍问题。研究目标是提供一个兼具高性能、可扩展性和易用性的框架,使研究人员和实践者能够轻松入门。为此,OpenRLHF 提出了一系列创新点:

  • 首个基于 Ray 的开源 RLHF 和 RLVR 架构:利用 Ray 灵活的分布式计算原语,OpenRLHF 能够简化 RLHF 和 RLVR 工作流的编排和资源管理,显著降低了分布式操作和部署的复杂性,同时提升了可用性和灵活性。
  • 结合 DeepSpeed-ZeRO 和环形注意力的 3D 并行:为实现大模型无缝高效扩展,OpenRLHF 集成了 DeepSpeed-ZeRO 提供的自动张量并行(AutoTP),并通过环形注意力实现序列并行。这种集成方式实现了一种高效的“3D”并行策略(张量、数据和序列并行),无需复杂的工程或大量的用户配置。
  • 首次使用 vLLM 加速 CoT 推理:针对长思维链(long-CoT)RLHF 和 RLVR 工作负载中的关键推理瓶颈,OpenRLHF 整合了先进的 vLLM 推理引擎。通过词元级并行解码、高级缓存和优化的动态批处理,该集成显著加速了推理过程,从而大幅提升了整体训练运行时的效率。
  • 异步数据流与远程引擎交互:OpenRLHF 支持异步数据流和远程引擎通信,以最大化分布式 RLHF 训练期间的系统吞吐量和资源利用率。在该架构中,rollout 引擎、actor 引擎和远程引擎独立运行并通过消息传递进行通信,使得数据一旦可用就能立即处理。这种设计减少了空闲时间,提高了流水线效率,并增强了在大型分布式 GPU 集群上的可扩展性和灵活性。利用异步远程引擎交互,OpenRLHF 可以轻松扩展以支持可扩展的智能体强化学习训练。

A3 相关工作

强化学习从人类反馈(RLHF)与从可验证奖励中强化学习(RLVR)。RLHF 已成为使大型语言模型与人类偏好对齐的关键范式【4, Deep reinforcement learning from human preferences, 2017, Advances in neural information processing systems】【24, Learning to summarize with human feedback, 2020, Advances in Neural Information Processing Systems】。其基础框架通过人类偏好数据训练一个奖励模型,并使用如 PPO 【18, Proximal policy optimization algorithms, 2017, arXiv preprint arXiv:1707.06347】等强化学习算法优化语言模型。该方法已成功应用于 InstructGPT 【15, Training language models to follow instructions with human feedback, 2022, Advances in neural information processing systems】、ChatGPT 和 GPT-4 【1, Gpt-4 technical report, 2023, arXiv preprint arXiv:2303.08774】等知名模型。在 RLHF 的基础上,RLVR 利用了来自数学验证、代码执行或其他自动评估机制的可自动验证信号【7, Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning, 2025, arXiv preprint arXiv:2501.12948】【22, Exploring data scaling trends and effects in reinforcement learning from human feedback, 2025, arXiv preprint arXiv:2503.22230】【5, Training verifiers to solve math word problems, 2021, arXiv preprint arXiv:2110.14168】。尽管 PPO 因其稳定性仍占主导地位【18, Proximal policy optimization algorithms, 2017, arXiv preprint arXiv:1707.06347】【30, Fine-tuning language models from human preferences, 2019, arXiv preprint arXiv:1909.08593】,但直接偏好优化(DPO)【16, Direct preference optimization: Your language model is secretly a reward model, 2023, Advances in Neural Information Processing Systems】等替代方法因其计算效率而受到关注。然而,RLHF 和 RLVR 方法都需要大量的计算资源和复杂的分布式训练策略。这些计算挑战催生了各种优化技术,包括高效的推理引擎【10, Efficient memory management for large language model serving with pagedattention, 2023, Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles】、内存高效的训练策略【17, Zero: Memory optimizations toward training trillion parameter models, 2020, SC20: International Conference for High Performance Computing, Networking, Storage and Analysis】以及分布式编排框架【27, Deepspeed-chat: Easy, fast and affordable rlhf training of chatgpt-like models at all scales, 2023, arXiv preprint arXiv:2308.01320】。尽管如此,现有解决方案通常为了简化而牺牲性能,或为了优化而牺牲可访问性,这正是 OpenRLHF 旨在解决的差距。

RLHF(或 RLVR)框架。RLHF 和 RLVR 训练的计算复杂性推动了专门框架的开发。为小规模网络设计的通用 RL 框架【12, Rllib: Abstractions for distributed reinforcement learning, 2018, International conference on machine learning】【6, Openai baselines, 2017, https://github.com/openai/baselines】【20, Auxiliary-task based deep reinforcement learning for participant selection problem in mobile crowdsourcing, 2020, Proceedings of the 29th ACM International Conference on Information & Knowledge Management】无法解决 LLM 的特定挑战,并且通常采用具有复杂进程间通信的多控制器架构,导致学习曲线陡峭。针对 RLHF 的特定框架则面临着显著的性能与可访问性之间的权衡。像 TRL 【25, Trl: Transformer reinforcement learning, 2020, https://github.com/huggingface/trl】、DeepSpeed-Chat 【27, Deepspeed-chat: Easy, fast and affordable rlhf training of chatgpt-like models at all scales, 2023, arXiv preprint arXiv:2308.01320】和 ColossalChat 【11, Colossal-ai: A unified deep learning system for large-scale parallel training, 2023, Proceedings of the 52nd International Conference on Parallel Processing】等开源解决方案提供了易于访问的实现,但通常缺乏复杂的编排能力,并且在推理优化方面存在困难。而像 Nemoaligner 【19, Nemo-aligner: Scalable toolkit for efficient model alignment, 2024a, Preprint, arXiv:2405.01481】、ChatLearn 【2, chatlearn, 2017, https://github.com/alibaba/ChatLearn】和 Verl 【23, Hybridflow: A flexible and efficient rlhf framework, 2024, arXiv preprint arXiv: 2409.19256】等工业级框架则提供了包括 3D 并行和内存管理在内的先进优化。然而,它们的架构紧密耦合,需要大量的工程专业知识和广泛的基础设施设置,为学术研究人员设置了可访问性障碍。这些系统采用静态资源分配范式,导致利用率不佳和适应性有限。大多数框架未能同时充分解决推理瓶颈和可用性挑战,迫使用户在高性能和易用性之间做出选择。

A2 方法细节

3.1 概览:基于 Ray 的 RLHF 架构

OpenRLHF 的整体架构。OpenRLHF 是首个基于 Ray 的开源 RLHF 架构,它将一批 GPU 分配给不同的角色,并利用 Ray 的调度能力管理这些角色之间的数据流和工作流。如图 1 所示,它定义了两个主要角色:负责为给定提示生成响应的 rollout 引擎,以及负责计算 logprobs、参考策略 logprobs 和处理模型训练的 ZeRO 引擎。(关于 PPO 工作流的详细设计,请参阅附录 B;关于深入的实现技巧,请参阅博客【21, Advanced tricks for training large language models with proximal policy optimization, 2024b, Notion Blog】)。此外,我们利用 VLLM 实现 rollout 引擎,以最小的 GPU 内存使用实现高效的响应生成。对于模型训练,我们采用 DeepSpeed,实现了 3D 并行——包括自动张量并行、ZeRO/数据并行和序列并行——以高效地训练 actor 和 value 模型。

Figure 1: OpenRLHF 总体架构。系统将 GPU 分配给两个主要角色:专门用于响应生成的 rollout 引擎和负责计算对数概率及模型训练的 actor 引擎。OpenRLHF 利用 Ray 进行分布式调度,集成 VLLM 以实现低 GPU 内存占用的高效响应 rollout,并采用 DeepSpeed-Zero 进行 3D 并行(包括张量、数据和序列并行)以实现高效模型训练。底层模型通过灵活的 Transformer 架构实例化,使系统易于扩展和适应不同的研究场景。
Figure 1: OpenRLHF 总体架构。系统将 GPU 分配给两个主要角色:专门用于响应生成的 rollout 引擎和负责计算对数概率及模型训练的 actor 引擎。OpenRLHF 利用 Ray 进行分布式调度,集成 VLLM 以实现低 GPU 内存占用的高效响应 rollout,并采用 DeepSpeed-Zero 进行 3D 并行(包括张量、数据和序列并行)以实现高效模型训练。底层模型通过灵活的 Transformer 架构实例化,使系统易于扩展和适应不同的研究场景。

易用性的核心原因。易用性的核心在于简化的模型切分、无缝集成和灵活调度。rollout 引擎和训练引擎之间的模型权重交换是通过一个灵活的切分和分区流水线实现的。Hugging Face Transformer 【9, Hugging face, 2022, Introduction to transformers for NLP: With the hugging face library and models to solve problems】模型使用 DeepSpeed ZeRO、AutoTP 和 Ring-Attention 【14, Ring attention with blockwise transformers for nearinfinite context, 2023, arXiv preprint arXiv:2310.01889】模型并行进行实例化和训练。然后,这些模型切片通过 AutoTP 和 AutoPP 高效地传输到 VLLM,后者将模型动态地分区为子模块。基于 Ray 的调度机制使得在不同模型并行模式(如混合引擎和异步训练)之间无缝切换成为可能。这个简化的工作流显著降低了复杂性,使得系统非常用户友好且易于扩展。与 DeepSpeed-Chat、Transformer Reinforcement Learning (TRL) 或其他主流框架的架构相比,OpenRLHF 支持异步数据流和远程引擎交互,显著提高了训练过程和智能体工作流的整体效率。

3.2 分布式高效系统设计

结合 DeepSpeed-ZeRO 和环形注意力的 3D 并行。为了实现大模型无缝且高效的可扩展性,OpenRLHF 集成了 DeepSpeed-ZeRO 最新的自动张量并行(AutoTP)功能。在许多工业级 RLHF 架构中,用户之前需要为每个 transformer 模型手动指定一个注入策略,以识别哪些线性层和注意力输出需要在数据并行排名之间进行通信。相比之下,OpenRLHF 利用 DeepSpeed-ZeRO 的新功能,默认支持 HuggingFace 模型的自动张量并行:当未启用内核注入且未提供注入策略时,DeepSpeed 会在运行时自动确定并应用必要的策略。这极大地简化了用户体验,并将强大的张量并行支持扩展到更广泛的模型范围,无需复杂的工程或手动配置。此外,OpenRLHF 通过环形注意力实现序列并行。环形注意力采用基于环的通信拓扑,有效地将长序列的注意力计算分布到多个 GPU 上,同时最小化内存使用和通信开销。这对于涉及长思维链(CoT)推理的现代 RLHF 和 RLVR 工作负载尤为关键,因为在这些场景中,传统的注意力计算可能成为一个主要的可扩展性瓶颈。通过结合 AutoTP、数据并行和基于环形注意力的序列并行,OpenRLHF 使得在灵活的 GPU 集群上进行大规模、高效且高度可用的 RLHF 模型训练成为可能。

使用 vLLM 加速 CoT 推理。随着 LLM 在推理能力上的进步,RLHF 和 RLVR 流水线越来越面临推理瓶颈,尤其是在处理长思维链(CoT)输出时。对于像 OpenAI-o1 和 DeepSeekR1 这样的模型,CoT 生成可能占据训练时间的主导地位,使得高效的长格式推理对可扩展性至关重要。为应对这一挑战,OpenRLHF 集成了 vLLM 推理框架,该框架专为高吞吐量和内存高效的 LLM 服务而设计。vLLM 为生成 RLHF 样本和支持频繁的模型权重更新提供了一个简化的接口。vLLM 的核心创新是使用 PagedAttention 【10, Efficient memory management for large language model serving with pagedattention, 2023, Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles】高效管理注意力键和值内存。这项技术将内存浪费显著减少到 4% 以下,从而能够批处理更多序列,并提高 GPU 利用率和吞吐量。PagedAttention 还支持对高级采样方法(如并行采样和波束搜索)进行高效的内存共享,将内存使用量减少高达 55%,进一步提升推理效率。除了 PagedAttention,vLLM 还有其他几个优势,包括对传入请求的连续批处理、使用 CUDA Graph 的快速模型执行,以及使用 FlashAttention 和 FlashInference 优化的 CUDA 内核,这些都实现了快速且内存高效的注意力计算。它还具有用于更快推理的推测解码和用于减少长序列延迟的分块预填充功能。总的来说,这些增强功能使 vLLM 在大规模、长序列推理方面非常有效,特别适用于效率和可扩展性至关重要的 RLHF 和 RLVR 流水线。

异步数据流与远程引擎交互。OpenRLHF 支持异步数据流和远程引擎通信,以在分布式 RLHF 训练期间最大化系统吞吐量和资源利用率。在这种架构中,rollout 引擎、actor 引擎和远程引擎独立运行,并通过消息传递进行通信,使得数据一旦可用就能立即得到处理。完全异步执行在思维链(CoT)时代尤为重要,因为推理涉及生成长度和计算成本差异巨大的多步推理。在同步框架中,最慢的 CoT 生成会成为整个流水线的瓶颈并浪费资源。相比之下,OpenRLHF 的异步设计允许每个引擎按其自身节奏运行,确保即使对于长或可变的 CoT 任务,硬件也能得到利用。这不仅加速了训练和评估,还实现了高效扩展,并支持动态的智能体 RL 工作流。利用异步远程引擎交互,OpenRLHF 可轻松扩展,以支持在现代以 CoT 为中心的环境中进行可扩展的智能体 RL 训练。

A4 实验

实验环境

  • 实验场景:长思维链(CoT)生成场景下的 RLVR 设置,以及通用的 RLVR 和 RLHF 设置。
  • 数据集
    • 长 CoT 实验:使用能产生足够长上下文输出的基准模型进行压力测试,未指定具体数据集。
    • 通用 RLVR 实验:GSM8K 数据集【5, Training verifiers to solve math word problems, 2021, arXiv preprint arXiv:2110.14168】。
    • 通用 RLHF 实验:微调 1024 个提示。
  • 模型架构
    • 长 CoT 实验:模型参数规模为 1.5B、7B 和 14B。基座模型采用 DeepSeek 开源的蒸馏版 Qwen 系列。
  • 硬件配置
    • 8块 NVIDIA H200 140GB GPU。
  • 软件配置
    • 框架版本:OpenRLHF (v0.8.5) 对比 Verl (v0.4.0)、优化的 TRL 框架和优化的 DeepSpeed-Chat (DSChat)。
    • 依赖库:PyTorch 2.7 【8, Pytorch, 2021, Programming with TensorFlow: solution for edge computing applications】。
    • 算法与优化器:
      • 长 CoT 实验:使用解耦裁剪和动态采样策略优化(DAPO)算法【28, Dapo: An open-source llm reinforcement learning system at scale, 2025, arXiv preprint arXiv:2503.14476】,优化器为 ZeRO Stage 3 或全分片数据并行(FSDP)【29, Pytorch fsdp: experiences on scaling fully sharded data parallel, 2023, arXiv preprint arXiv:2304.11277】。
      • 通用 RLVR 实验:使用 GRPO 算法。
      • 通用 RLHF 实验:使用 PPO 算法。
    • 超参数:所有对比实验均在相同的超参数设置下进行,本地批量大小设为1,最大输入上下文长度为1024个词元。

实验结果

性能比较

长 CoT 实验性能分析。如表 1 所示的实验结果表明,OpenRLHF 在所有配置下均取得了优于 Verl 的训练速度。OpenRLHF 的加速比从 1.5B 模型的 1.22 倍到 14B 模型的 1.68 倍不等,且随着模型规模和上下文长度的增加,性能优势愈发明显。例如,在 14B-8K 的设置中,OpenRLHF 实现了 1.56 倍的加速(328.6 秒 vs. 511.1 秒),而在 7B-2K 配置下,加速比为 1.56 倍(30.3 秒 vs. 47.3 秒)。这些速度提升归功于 OpenRLHF 的算法设计,包括 DAPO 优化策略,它有效缓解了长上下文场景下的内存开销和计算瓶颈。跨不同规模的一致性加速突显了 OpenRLHF 在当代 RLHF 流水线中的效率优势。

Table 1: OpenRLHF 与 Verl 在不同模型大小和上下文窗口长度下的平均每步训练时间(秒)对比。加速比是 Verl 时间 / OpenRLHF 时间在所有序列长度上的几何平均值。
Table 1: OpenRLHF 与 Verl 在不同模型大小和上下文窗口长度下的平均每步训练时间(秒)对比。加速比是 Verl 时间 / OpenRLHF 时间在所有序列长度上的几何平均值。

通用 RLVR 实验。为确保对强化学习微调中的训练效率进行公平可控的评估,我们在 GSM8K 数据集上使用 GRPO 算法对 OpenRLHF 和优化的 TRL 框架进行了一个训练周期的比较。两个系统均配置相同的超参数并在相同的硬件上运行,以隔离框架设计对性能的影响。选择 GRPO 算法是因为其在基于奖励的微调场景中的相关性,而 GSM8K 是算术推理任务的代表性基准。在训练效率方面,OpenRLHF 显示出显著优势,完成一个周期耗时 1,657 秒,而 TRL 需要 5,189 秒,加速比约为 3.1 倍。这一结果凸显了 OpenRLHF 实现的卓越效率和可维护性,这得益于其简化的设计和有针对性的系统级优化。

通用 RLHF 实验。为了评估现代 RLHF 框架的训练效率,我们将 OpenRLHF 与优化的 DeepSpeed-Chat (DSChat) 实现进行了比较。实验涉及使用 PPO 算法对 1,024 个提示进行一个周期的微调,硬件和超参数配置完全相同。这种设置确保了观察到的性能差异仅归因于两个框架在系统设计和优化策略上的不同。在训练时间方面,OpenRLHF 以 236.8 秒完成任务,显著优于需要 855.09 秒的 DSChat,实现了 3.6 倍的加速。这一性能提升主要由 OpenRLHF 的两个系统级创新驱动:使用 vLLM 加速词元生成和使用 Ray 进行高效的分布式执行。总的来说,这些设计选择共同造就了一个更具可扩展性和卓越性能的 RLHF 框架。

易用性比较

代码复杂度对比。如图 2 所示,OpenRLHF 在实现的简洁性和训练性能之间取得了令人信服的平衡。尽管其代码行数仅为 8,523 行,是第二简洁的框架——远少于 TRL 的 19,071 行和 Verl 的 32,325 行——OpenRLHF 在标准 RLHF 基准测试中仍展现出明显的性能优势。这种精简的代码库不仅便于开发者理解和修改,还降低了与集成到自定义流水线相关的工程开销。除了其轻量化设计,OpenRLHF 还为各种强化学习微调范式提供了全面支持,包括监督微调(SFT)、直接偏好优化(DPO)【16, Direct preference optimization: Your language model is secretly a reward model, 2023, Advances in Neural Information Processing Systems】、奖励模型(RM)和过程奖励模型(PRM)【13, Let’s verify step by step, 2023, The Twelfth International Conference on Learning Representations】。这种广泛的功能性,结合其模块化和文档齐全的架构,显著降低了研究和生产使用的入门门槛。总体而言,OpenRLHF 的设计通过将最小的代码复杂性与广泛的功能和有竞争力的性能相结合,体现了高度的可用性。

Figure 2: RLHF 框架核心代码复杂度比较(代码行数)。数值越低表示实现越简单,可维护性越好。
Figure 2: RLHF 框架核心代码复杂度比较(代码行数)。数值越低表示实现越简单,可维护性越好。

A5 结论与局限性

结论

OpenRLHF 是一个简单而高性能的开源框架,它弥合了 RLHF 和 RLVR 训练中性能与可用性之间的鸿沟。通过整合 Ray 的分布式计算、vLLM 的推理优化、DeepSpeed ZeRO 的内存效率以及环形注意力的序列并行,OpenRLHF 实现了四大关键创新:基于 Ray 的架构简化了编排,3D 并行实现了高效扩展,加速的 CoT 推理解决了关键瓶颈,异步数据流最大化了吞吐量。该框架已被卡内基梅隆大学的学术课程和大型科技公司的生产部署等领先机构和公司广泛采用,验证了其在现实世界中的有效性。OpenRLHF 对后续框架开发的影响及其在推动 RLHF 研究大众化方面的作用,证明了其对该领域的重大贡献。通过开源此框架,旨在促进创新,加速研究进展,并使更广泛的社群能参与到对齐人工智能的开发中。

局限性

尽管 OpenRLHF 在平衡性能和可访问性方面表现出显著优势,但仍需承认几个局限性。首先,尽管进行了优化,OpenRLHF 的峰值性能可能无法与那些拥有专门工程团队和大量资源的、高度专业化的工业级框架相媲美。其次,作为一个没有专门经济支持的社区驱动的开源项目,OpenRLHF 在快速集成尖端功能方面面临资源限制,可能导致与资金充足的商业框架相比出现延迟。目前,该框架主要关注语言模型,不支持视觉-语言模型或其他多模态架构,这限制了其在多模态 AI 对齐研究中的应用。此外,OpenRLHF 的模块化设计引入了对 Ray、vLLM 和 DeepSpeed 等外部系统的依赖,这些上游系统的更新可能需要维护工作或引入兼容性问题。尽管存在这些局限性,OpenRLHF 在提升 RLHF 研究的可访问性和大众化方面的贡献为社区提供了巨大价值。

A6 附录

A. 广泛影响与应用

学术与社区认可。OpenRLHF 自发布以来,在学术界和工业界都获得了显著的应用和影响。该框架在高性能和易用性之间取得的平衡,使其成为从研究实验到生产部署等多种应用的首选。它已被正式整合到学术课程中,特别是被卡内基梅隆大学 2025 年春季高级自然语言处理课程采纳为核心教学框架,这证明了其对 RLHF 新手的教育价值和可及性。该框架还获得了主流技术社区的认可,包括受邀在 2025 年 PyTorch 专家交流会上发表演讲,突显了其对更广泛的深度学习生态系统的技术贡献。

行业与机构采用。该框架已被包括谷歌、字节跳动、百度、英伟达、腾讯、中国电信、Vivo、NexusFlow、JSC、加州大学伯克利分校的 Starling 团队、美团和香港科技大学在内的领先科技公司和研究机构广泛采用。这种跨越主要科技公司、电信供应商和学术机构的广泛采用,证明了 OpenRLHF 在处理从研究原型到生产规模部署等不同规模 RLHF 工作负载时的多功能性和稳健性。

作为基础框架的贡献。OpenRLHF 的模块化设计和可扩展架构使其成为专业化框架的基础。著名例子包括用于多模态强化学习的 LMM-R1、用于高级推理任务的 MARTI,以及用于多模态应用的 MM-EUREKA。这些衍生框架展示了 OpenRLHF 在支持多样化研究方向方面的灵活性,及其作为 RLHF 生态系统中创新平台的作用。

对后续框架发展的影响。OpenRLHF 中引入的设计原则和技术创新影响了该领域后续框架的开发。包括 Verl、Alibaba ROLL、SLIME 和 Open-Reasoner-Zero 在内的几个著名框架在其文档和出版物中承认了 OpenRLHF 的贡献,并引用其分布式架构设计、基于 Ray 的编排方法和集成策略作为其自身开发的影响因素。这种承认反映了 OpenRLHF 在推动 RLHF 系统设计前沿和建立分布式 RLHF 训练最佳实践方面的作用。

推动研究大众化。除了直接使用和引用外,OpenRLHF 通过降低学术研究人员和小型组织的入门门槛,为 RLHF 研究的大众化做出了贡献。该框架在不牺牲性能的前提下强调易用性,使得更广泛的研究社区能够参与 RLHF 研究,从而在那些可能缺乏复杂分布式训练设置资源的社群中 fostering 创新。

B. PPO 工作流设计

PPO 训练流程概览。本节介绍了 OpenRLHF 全面的基于 PPO 的 RLHF 训练工作流,该工作流协调多个专用引擎,以高效处理复杂的多阶段训练过程(如图 3 所示)。OpenRLHF PPO 工作流由四个主要阶段循环执行:(1) Rollout 生成,即当前策略为提示生成响应;(2) 奖励计算,即使用训练好的奖励模型评估响应;(3) 优势估计,即使用 GAE(广义优势估计)计算优势和回报;(4) 策略优化,即使用 PPO 损失更新策略。此流水线会重复迭代多次直至收敛。

Rollout 生成阶段。训练始于 Rollout 引擎使用当前策略 $\pi_{\theta}$ 为一批提示生成响应。从训练数据集中采样一批提示 $\{x_1, x_2, ..., x_B\}$,配备了 vLLM 以实现高效推理的 Rollout 引擎使用当前策略生成响应 $\{y_1, y_2, ..., y_B\}$。在生成过程中,引擎记录动作的对数概率 $\log \pi_{\theta}(y_i|x_i)$ 和注意力掩码。生成的序列 $(x_i, y_i)$ 及其元数据被收集起来用于后续处理。Rollout 引擎利用 vLLM 的优化,包括连续批处理、KV 缓存管理和 PagedAttention,以在此推理密集型阶段最大化吞吐量。

奖励与参考策略对数概率计算。一旦生成了 rollouts,系统会计算奖励和参考策略的对数概率。训练好的奖励模型 $R_{\phi}$ 会评估每个提示-响应对,以产生标量奖励 $r_i = R_{\phi}(x_i, y_i)$。同时,冻结的参考策略 $\pi_{ref}$ 计算用于 KL 正则化的对数概率 $\log \pi_{ref}(y_i|x_i)$,而 critic 网络 $V_{\psi}$ 则估计用于优势计算的状态值 $V_{\psi}(x_i, y_i, :t)$。这些计算可以在多个 GPU 上并行化,并且通常为了效率而批量处理。

Figure 3: OpenRLHF 的整体 PPO 工作流。
Figure 3: OpenRLHF 的整体 PPO 工作流。

优势与回报计算。接着,系统使用广义优势估计(GAE)计算优势。首先,计算时间差分残差:$\delta_t = r_t + \gamma V_{\psi}(s_{t+1}) - V_{\psi}(s_t)$。然后,使用 GAE 计算优势:$A_t = \sum_{l=0}^{\infty}(\gamma\lambda)^l\delta_{t+l}$,并计算折扣回报为 $R_t = A_t + V_{\psi}(s_t)$。优势计算中加入了 KL 惩罚项,以防止策略偏离参考策略太远:$r'_t = r_t - \beta KL[\pi_{\theta}||\pi_{ref}]$。

策略更新。最后阶段使用 PPO 的裁剪目标函数更新策略。ZeRO 引擎使用裁剪的替代目标计算策略损失:$L^{CLIP}(\theta) = \hat{E}[\min(r_t(\theta)\hat{A}_t, \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon)\hat{A}_t)]$,其中 $r_t(\theta) = \frac{\pi_{\theta}(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)}$ 是概率比。critic 损失计算为 $L^V(\psi) = \hat{E}[(V_{\psi}(s_t) - R_t)^2]$。总损失结合了策略和价值损失:$L_{total} = L^{CLIP} + c_1L^V + c_2S[\pi_{\theta}](s_t)$,其中 $c_1$ 和 $c_2$ 分别是价值损失和熵奖励的系数,而 $S[\pi_{\theta}](s_t)$ 是策略分布的熵,以鼓励探索。ZeRO 引擎使用 DeepSpeed ZeRO 优化进行梯度计算和模型参数更新,以实现内存效率和可扩展性。

基于 Ray 的架构协调。在整个工作流中,OpenRLHF 基于 Ray 的架构实现了不同引擎之间的无缝协调,同时保持了计算效率。Rollout 引擎和 ZeRO 引擎可以运行在为各自工作负载优化的不同 GPU 集群上——为 rollout 生成使用推理优化的硬件,为策略更新使用训练优化的硬件。Ray 的分布式调度自动处理分布式系统中的数据传输、同步和容错。这种异步执行模型允许系统在可能的情况下重叠计算阶段,与传统的同步 RLHF 实现相比,显著提高了整体训练吞吐量。

训练循环与高级技术。整个 PPO 训练循环会持续多个周期,每个周期处理多批 rollout 数据。基于 KL 散度阈值和性能指标的早停标准可防止策略崩溃并确保训练稳定。模块化设计便于集成高级技术,如优势归一化、梯度裁剪和自适应 KL 惩罚系数,使得 OpenRLHF 对于不同的研究和生产场景具有高度可定制性。

C. 完整贡献者名单

  • Ray 集成: Jian Hu, Xibin Wu, Songlin Jiang
  • vLLM 集成: Jian Hu, Xibin Wu, Songlin Jiang, Kaichao You (vLLM 团队)
  • 环形注意力: Zilin Zhu (智谱), Zhibo Zhou (Vivo), gzpan(GitHub 用户), Jian Hu
  • 监督学习: Haoran Wang, Xianyu
  • DeepSpeed 和 Transformers 集成: Jian Hu, Xibin Wu, Songlin Jiang, Bin Chen
  • 异步智能体 RL: Haotian Xu, Jian Hu
  • 单控制器架构: Jian Hu
  • PPO 实现: Jian Hu
  • GRPO 实现: Jason Klein Liu
  • 动态采样: Jian Hu
  • 动态批处理: Hao Chen
  • 文档: Jian Hu
  • 测试与 Bug 报告: Weikai Fang, Wei Shen, Weixun Wang
  • 论文撰写与演示: Wei Shen, Jason Klein Liu, Weixun Wang, Yu Cao, Jian Hu

更完整的列表可以在 OpenRLHF 的提交和发布历史中找到。