PipelineRL：通过飞行中权重更新优化大语言模型强化学习

PipelineRL 是一种实验性的强化学习（RL）实现，旨在解决大规模 LLM 训练中高推理吞吐量与在策略数据收集之间的根本权衡。通过实现 飞行中权重更新，PipelineRL 使推理服务器能够在不停止推理过程的情况下接收更新的模型权重，确保 GPU 高利用率，同时保持训练数据接近在策略。

解决推理吞吐量与在策略权衡

在传统的 RL 工作流中，效率与数据新鲜度之间存在冲突。为了实现高吞吐量，推理服务器通常使用大批量大小，这会为多个策略优化步骤生成数据。然而，每一次后续的优化步骤都会增加用于收集数据的权重与当前策略权重之间的“滞后”，导致数据越来越偏离在策略，训练效果下降。

PipelineRL 通过在每次优化器步骤后更新推理服务器中的权重而不停止推理来解决此问题。系统仅在接收新权重所需的短暂时间内暂停推理服务器。此方法使推理服务器能够保持最佳批量大小，同时确保数据保持在策略或接近在策略，从而实现更稳定和更有效的学习。

性能与稳定性结果

在 Open-Reasoner-Zero 数据集上使用 7B 和 32B 模型进行的实验表明，PipelineRL 在 AIME 2024 和 MATH 500 推理基准上匹配或超越了 Open-Reasoner 的性能。

简化的 RL 算法

尽管性能竞争力强，PipelineRL 使用的 RL 实现比 Open-Reasoner-Zero 简单得多。关键简化包括：

  • 简化的 GRPO：它使用了 Group Relative Policy Optimization（GRPO）的简化版本，且不包含价值函数。
  • 无复杂过滤：实现省略了信任域重要性权重限制、过长序列过滤和奖励整形。
  • 基本损失归一化：损失使用批次中序列数量作为分母进行归一化，对所有 token 赋予相等权重。
  • 无惩罚：系统不使用 KL 惩罚或熵奖励（尽管支持参考模型 KL）。

KV 缓存陈旧数据的影响

飞行中权重更新的主要担忧是序列生成仍在使用 KV 缓存中陈旧的键和值，因为这些是用模型的先前版本计算的。然而，实验结果表明这并未对训练稳定性产生不利影响。

模块化架构与技术契约

PipelineRL 旨在实现模块化，以便与专用推理（例如 SGLang、vLLM）和训练（例如 DeepSpeed、FSDP、TorchTitan）软件集成。这通过两个主要契约实现：

推理契约

要与 PipelineRL 集成，推理软件必须提供三个特定的 API：

  1. 进程组初始化：一个 HTTP POST /init_process_group 请求，用于初始化用于权重更新的进程组。
  2. 权重更新触发：一个 HTTP POST /request_weight_update 请求，通知推理服务器暂停并通过 NCCL 接收权重广播。
  3. 聊天完成：标准的 HTTP POST /v1/chat/completion 请求，用于 actor 交互。

训练器契约

训练软件必须提供以下操作的 Python API：

  • 工作节点初始化：加载并分片训练权重和优化器状态。
  • 前向传播：生成 token 的对数似然。
  • 反向步骤：计算并累积 RL 目标的梯度。
  • 优化器步骤：执行优化器步骤。
  • 权重收集与广播：逐层收集更新后的权重，以广播到推理服务器。

实验配置

PipelineRL 在 7B 和 32B 模型上使用以下超参数进行测试：

  • 批量大小：4096
  • 学习率：1e-6
  • 最大生成 token 数：8192

训练计算需求约为 7B 模型在 2 台节点上 3.5 天，32B 模型在 4 台节点上 6 天。

Sources