PipelineRL

我们很高兴开源 PipelineRL，这是一个实验性的强化学习（RL）实现，它解决了大型语言模型（LLM）强化学习中的一个基本挑战：推理吞吐量与在线数据收集之间的权衡。PipelineRL 的关键创新在于 RL 训练期间的**飞行中权重更新**（见下图 1）。这使得 PipelineRL 能够实现持续高推理吞吐量，并最大限度地减少用于 rollout 的权重与最新更新的模型权重之间的滞后。结果是：大型语言模型训练快速且稳定。

在这篇博客文章中，我们展示了 1) 飞行中权重更新不会损害训练过程，以及 2) PipelineRL 在使用更简单的 RL 算法的情况下，取得了与 Open-Reasoner-Zero 相当的结果。我们还介绍了模块化的 PipelineRL 架构，该架构有助于尝试新的推理/训练器组合。

传统强化学习 vs PipelineRL

在传统的强化学习方法中（图 1a），高吞吐量推理和在线数据收集之间存在权衡。为了解释这种权衡，我们首先以算法形式定义传统强化学习：

current_policy = initial_policy
opt_state = init_optimizer(current_policy)

while True:
    # RL step starts
    # inference
    inference_policy = current_policy
    list_of_prompts = [sample_prompts(training_batch_size) \
        for _ in range(num_grad_steps)]
    list_of_rollouts = [sample_rollouts(prompts, inference_policy) \
        for prompts in list_of_prompts]
    # training
    lag = 0 # lag between the inference and current policies
    for rollouts in list_of_rollouts:
        current_policy, opt_state = policy_update(current_policy, opt_state, rollouts)
        lag += 1
    # RL step ends

为了实现高吞吐量，推理服务器必须使用大批量处理，因此需要为多个策略优化步骤生成数据。然而，每个优化步骤都会增加当前策略与使用推理策略收集的数据之间的滞后，从而逐渐使收集到的数据更偏离策略且对训练效果更差。在线学习需要单次优化步骤的数据。但是，使用多个 GPU 生成少量数据效率低下，因为这意味着每个 GPU 的批量大小很小。此外，随着推理服务器完成短序列，仅处理剩余少量最长序列时，批量大小会下降。

PipelineRL (图 1b) 通过飞行中权重更新解决了这一权衡。我们在每个优化器步骤之后更新推理服务器中的权重，**且从不停止推理**。我们只在所有推理服务器暂停，用于接收新权重所需的时间。飞行中权重更新允许推理服务器持续保持最佳批处理大小，同时确保数据保持在线策略或接近在线策略，这分别带来了更好的 GPU 利用率和更有效的学习。

PipelineRL 有效！

为了展示 PipelineRL 的有效性和飞行中权重更新的优势，我们在 Open-Reasoner-Zero 数据集上训练了一个 7B 模型和一个 32B 模型。从学习曲线可以看出，PipelineRL 在流行的推理测试基准：AIME 2024 和 MATH 500 上，与 Open-Reasoner 相比，性能不相上下或更优（见图 2）。

值得注意的是，我们的强化学习实现比 Open-Reasoner-Zero 简单得多。Open-Reasoner-Zero 使用值函数，而我们的实现是 GRPO 的简化版本。特别是，我们发现不需要信任区域重要性权重钳制即可实现稳定训练。DAPO 论文中的过长序列过滤或奖励整形也未被使用。为了归一化损失，我们仅将批次中的序列数量作为分母，为所有 token 赋予相等权重。我们没有使用 KL 惩罚，也没有使用熵奖励（尽管我们的实现确实支持参考模型 KL）。尽管我们的实现很简单，或者可能正因如此，训练非常稳定，您可以在此 wandb 报告中看到。

人们可能会认为，飞行中权重更新会导致训练过程不稳定，因为**序列生成是在 KV 缓存中使用先前模型版本计算的陈旧键和值进行的**。然而，我们的实验表明这并不会对稳定性产生不利影响。

PipelineRL 架构

PipelineRL 采用模块化设计，旨在利用高度专业化的推理和训练软件（SGLang、vLLM、Nvidia Dynamo、DeepSpeed、FSDP、TorchTitan、FastLLM 等）的快速改进。我们为推理和训练组件之间提供了清晰的契约，以便在新的推理和训练解决方案可用时轻松集成它们。

推理契约

推理软件必须向 PipelineRL[1] 开放以下 API：

1. 进程组初始化：在启动时，训练器 0（指定的协调器）向所有推理服务器发送 HTTP `POST /init_process_group` 请求。此请求初始化将用于发送权重更新的进程组。
2. 权重更新触发：一旦训练器完成学习步骤（优化器步骤和权重收集），训练器 0 就向推理端点提交 HTTP POST /request_weight_update 请求。该请求包含主训练器进程即将通过 NCCL 传输的权重的顺序和形状的详细信息。推理服务器必须暂停推理并接收权重广播。
3. 聊天完成：参与者进程使用 HTTP POST /v1/chat/completion 请求与参与者 LLM 进行交互。

如果 `init_process_group` 和 `request_weight_update` API 成为行业标准，则可以使用 PipelineRL 即插即用地尝试不同的推理实现。

训练器契约

当每个训练器工作者积累了足够数量的训练 token 后，PipelineRL 训练代码会立即将新生成的训练数据提供给它们。任何公开这些 Python API 的训练软件都可以与 PipelineRL 协同工作。

• 工作器初始化：加载并分片训练权重和优化器状态。
• 前向传播：给定输入，生成 token 对数似然。
• 反向传播：计算并累积表示所选 RL 目标的标量梯度。
• 优化器步骤：执行优化器步骤。
• 权重收集与广播：在优化器步骤之后，训练器软件必须逐层收集更新后的模型权重，以准备将其广播到推理服务器。

PipelineRL 目前使用 HuggingFace `accelerate` 库，让用户可以在 DeepSpeed 和 FSDP 之间进行选择。但我们发现 `accelerate` 的契约过于灵活，可能会令人困惑。我们将转向如上所述的更严格的契约，这将使使用其他训练器变得更容易。

PipelineRL 的下一步计划是什么？

即将推出的功能。 我们的实现仍处于实验阶段，缺少一些重要的功能。我们的首要任务包括使用协程更精确地控制推理批处理大小、多模态支持和序列并行训练。我们也欢迎贡献更多推理服务器和训练器集成。但是，我们不会试图将 `pipeline-rl` 仓库打造成一个支持所有可能算法和奖励函数的大型框架。我们认为 `pipeline-rl` 应该是一个可修改、快速的 GRPO 参考实现，并且奖励可轻松验证。如果您想使用 PipelineRL 进行研究项目，只需分叉仓库并尽情修改代码即可！

更多研究即将发布。 需要进行更多分析来理解飞行中权重更新如何影响训练动态，并仔细测量 PipelineRL 带来的加速效果。此外，关于 PipelineRL 与 LLM 异步强化学习中高度相关的先前工作之间的相似之处，还有很多值得探讨。请持续关注我们即将发布的研究论文，以获取所有这些及更多内容！

贡献者和致谢

Alexandre Piché 在开发 TapeAgents 时编写了我们 RL 代码的第一个同步版本。Dzmitry Bahdanau 重构了代码，使其变为异步和分布式，并实现了飞行中权重更新。Rafael Pardinas 实现了序列打包。Ehsan Kamaloo 协助运行实验。Xiaoyin Chen 协助调试框架。

我们感谢 TRL、OpenRLHF 和 veRL 等先前 LLM 强化学习实现，我们从它们那里借鉴了许多技巧。来自 Simple-RL、Deepscaler、DAPO 和 OpenReasoner 等其他开源推理项目的成果，对稳定 PipelineRL 至关重要。我们要感谢 Christopher Manning 和 Michael Noukhovitch 的深思熟虑的评论。最后，我们感谢更广泛的 ServiceNow 研究团队和 ServiceNow CoreLLM 团队，他们是出色的同事。

[1] 代码中当前的契约略有不同，但我们正在按上述方式进行重构。

实验细节

我们在这里报告的 7B 和 32B 实验使用了相同的超参数：

• 批量大小 4096
• 学习率 1e-6
• 最大生成 token 数量 8192
  • 请注意，在 OpenReasoner 运行中，他们允许生成 16K 个 token

我们用于报告实验的计算资源：

• 7B 模型使用 2 个节点大约 3.5 天
• 32B 模型使用 4 个节点大约 6 天