LlamaRL: A Distributed Asynchronous Reinforcement Learning Framework for Efficient Large-scale LLM Training

作者/机构: Meta GenAI

A1 主要贡献

本文介绍了一个名为 LlamaRL 的全分布式、异步强化学习（RL）框架，该框架专为在拥有少量至数千台 H100 GPU 的集群上高效训练各种规模（8B、70B 和 405B 参数）的大型语言模型（LLM）而优化。

核心问题与挑战: 随着强化学习（RL）在提升大型语言模型（LLM）能力方面变得日益重要，开发一个能够可靠管理拥有数千亿参数的策略模型的高效 RL 框架面临着巨大的工程挑战，主要体现在以下三个方面：
1. 灵活性以支持各种 RL 算法: RL 训练过程复杂，涉及多种模型（如 actor, critic, reward model）和它们之间多样的数据通信。随着 PPO、RAFT、ReMax、CGPO 等算法的不断涌现，需要一个灵活的训练框架来轻松扩展以支持不同数量的模型和相应的数据/参数流。
2. 扩展至更大模型和更多 GPU: 训练更大的模型需要复杂的并行策略来处理内存和计算限制。例如，使用 PPO 训练 Llama 3.1 405B 模型需要至少 512 个 H100 GPU，并结合张量并行（TP）和全分片数据并行（FSDP）等技术。因此，需要一个高度可扩展的 RL 框架来高效支持不同模型规模、并行策略和 GPU 配置。
3. 因空闲气泡导致的 GPU 使用率不足: RL 算法中，响应生成和权重更新是顺序执行的。由于不同工作节点上响应的长度不同，处理时间也不同，这会导致 GPU 工作节点出现空闲（即“气泡”）。模型越大，生成时间的差异越大，气泡问题就越严重，从而降低了 GPU 利用率。

研究目标与创新点: 为应对上述挑战，本文提出了 LlamaRL 框架，其核心贡献如下：
* 简单且模块化的架构: LlamaRL 采用了一个基于原生 PyTorch 构建的精简单控制器架构。这种设计能够无缝扩展至数千个 GPU，支持对 405B 参数等大规模 LLM 的高效训练。其模块化结构和直观的控制逻辑使用户能够轻松地调整和扩展框架，以支持多样化的 RL 算法。
* 高效且可扩展的最佳实践: LlamaRL 引入了一套专为大规模、持续扩展的 LLM 进行 RL 训练而定制的最佳实践，与 DeepSpeed-Chat 等先进系统相比，在 405B 参数模型上实现了高达 10.7 倍的效率提升，且效率优势随模型规模增大而增长。这些核心实践包括：
* 协同定位的模型卸载与细粒度并行及量化: 认识到生成过程是内存密集型且耗时的，LlamaRL 将其完全从训练集群中卸载。这种完全卸载提升了可扩展性和灵活性，允许进行细粒度的并行和量化优化，从而显著降低计算和内存需求。
* 异步离策略（Off-policy）RL: 为了最小化计算气泡，训练和生成过程以异步方式并行运行（如图2所示），显著提高了吞吐量和资源利用率。这种方式不可避免地引入了 1 到 n 步延迟的离策略性。为此，本文提出了一种名为 AIPO（异步重要性加权策略优化）的离策略学习算法，有效缓解了大规模训练中的不稳定性。
* 通过直接 GPU 内存访问实现全分布式权重同步: 开发了一种全分布式、GPU 原生的同步方法，用于高效同步卸载的生成器和策略模型之间的模型权重。该方法利用 NVLink 实现 GPU 间的直接零拷贝传输，实现了线性可扩展性，能够在约 2 秒内完成跨数千 GPU 的 TB 级模型的权重更新（如图4所示）。

此外，本文还提供了 LlamaRL 效率的理论分析，包括一个形式化证明，表明其异步设计能带来严格的 RL 加速。通过广泛的实验，验证了 LlamaRL 在训练时间、模型质量以及离策略校正在异步训练中的重要性方面的优势。

A3 背景知识

大型语言模型 (LLM)
LLM 是一类强大的生成模型。经过大规模预训练和相对轻量的微调后，这些模型能够执行摘要、编码和翻译等多种任务。因此，LLM 可以成为非常有帮助的人类助手，有时在专业能力上甚至能超越人类用户或专家。

规模法则 (Scaling laws)
预训练研究表明，LLM 的性能可以通过模型大小、数据大小和计算量的幂律关系进行合理预测，这促进了工业级大模型的规模化。最近的研究发现，通过在可验证任务上进行 RL 来扩展推理时计算，是赋予 LLM 强大推理能力的有效方法，从而提升其在数学问题解决、代码生成和安全对齐方面的表现。这些进展凸显了扩大 LLM 计算规模的潜力，并强调了开发更高效强化学习系统的必要性。

用于 LLM 的强化学习
强化学习已成为微调顶尖 LLM（如 GPT-4, Claude, Gemini, Llama 2-Chat 等）的关键组成部分。RL 使得 LLM 能够超越监督训练数据的分布建模，调整采样分布以获得奖励函数的更高评价。对于来自人类反馈的强化学习 (RLHF)，其目标是使模型更具创造性、更安全、更有帮助，奖励通常由人类评估员提供。更广泛地说，奖励也可以来自基于规则的评分器，例如用于代码应用的单元测试执行或用于数学推理的与真实答案的匹配。

RL 的训练算法和系统
常用于训练 LLM 的 RL 算法，如近端策略优化 (PPO) 【索引29，Proximal policy optimization algorithms，2017，arXiv】，通常包含四个模型：策略模型、评论家模型（或价值模型）、参考策略模型和奖励模型。PPO 以迭代方式进行，每次迭代包括三个阶段：(1) 使用策略模型和一批提示生成响应；(2) 通过对评论家、参考策略和奖励模型进行单次前向传播来评分生成的响应，从而准备训练数据；(3) 通过前向和后向计算更新 actor 和 critic 模型。许多 RL 流水线遵循类似的阶段，但涉及不同数量的模型和数据依赖关系。例如，评论家模型是可选的，根据应用，奖励模型可以被基于规则的评分器取代。

图1展示了一个仅使用策略模型和几个基于规则的评分器的在线 RL 训练流程示例。为了简化，该流程中省略了用于正则化的参考模型。该示例也没有使用学习的评论家模型，而是从组得分中估计基线来计算用于策略更新的优势函数。这个例子利用了基于规则的评分器，这在代码和推理应用中很常见。策略模型有两个实例，分别基于全分片数据并行（FSDP）和 CUDA Graph 实现，以分别进行训练和推理优化。

图 1 在线 RL 训练的一个示例流程。为了简化，此流程中省略了用于正则化的参考模型。该示例流程也没有使用学习到的评论家模型，而是从组分数中估计基线来计算策略更新的优势函数。该示例使用了基于规则的评分器，这在代码和推理应用中很常见。策略模型有两个实例，分别基于全分片数据并行（FSDP）和CUDA图进行训练和推理优化。

A2 方法细节

4 高层策略

本节详细介绍 LlamaRL 实现大规模训练的几个关键高层策略。为简化起见，我们不考虑参考模型，并讨论类似于图2所示的算法设置，其中评论家模型被移除，并应用基于规则的评分器。

4.1 协同定位模型卸载

分布式模型放置以克服内存与并行限制。如前所述，协同定位模型放置策略限制了每个模型的可用内存使用，并强制所有模型共享相同的并行设置。因此，我们的初步尝试是将具有高昂内存或计算开销的模块卸载到专用的组中，即采用分布式模型放置。我们卸载了生成阶段，因为它是一个内存密集型过程，并且占据了主要的执行时间。因此，推理策略模型和训练策略模型被置于不同的处理组中，分别通过 Meta 内部的推理库和开源的 FSDP 实现。基于规则的评分器与训练策略模型分配在一起，并使用轻量级的 Python 程序进行计算。

通过 DDMA 实现高效的完全分布式模型放置。卸载后，处理组之间的通信至关重要，主要有两种类型：一是来自推理策略模型、奖励模型等的生成内容的通信；二是模型权重或 logits 的通信。我们没有采用像混合引擎（Hybrid Engine）那样部分协同定位推理和训练策略模型以节省权重通信开销的方式，而是开发了一种分布式直接内存访问（DDMA）方法，以实现秒级的权重通信。使用完全分布式模型放置的主要原因是它为采用异步 RL 训练提供了灵活性。

4.2 异步离策略 RL

采用异步学习以消除计算气泡。在切换到使用分布式组来容纳模型后，由于 RL 训练中模型的顺序执行，框架仍然存在 GPU 内存利用率低和可扩展性差的问题。为了克服这些问题，我们为 LlamaRL 框架应用了一种异步离策略学习算法，从而解除了 RL 工作流程中每个组件的阻塞。作为对比，图2展示了 (a) 同步在线策略 RL 和 (b) 异步离策略 RL 的过程。

异步并行执行提升训练速度。对于异步 RL，每个处理组并行运行，并在每个训练步骤结束时进行通信。这种设计使得训练器和生成器组件能够同时进行计算，从而减少了总训练时间。相比之下，同步在线策略 RL 在训练器和生成器步骤之间交替进行，一个组件运行时会阻塞另一个，这减慢了整体训练速度。

通过部分 rollout 缓解掉队者效应。为了进一步缓解数据并行组内部以及训练器和生成器组件之间的掉队者效应（straggler effects），我们采用了一种受部分 rollout 启发的策略【索引22，Kimi k1.5: Scaling reinforcement learning with llms，2025，Kimi Team】。具体来说，我们分解长的响应生成，缓存未完成的提示，并在后续迭代中恢复它们。

通过离策略校正处理异步性。然而，异步学习不可避免地会引入离策略性：策略模型是在由其前一个迭代版本生成的样本上进行训练的。尽管存在离策略性，我们的结果表明，在离策略校正的帮助下，离策略训练的影响与完全在线策略训练相比可以忽略不计。我们将在第6节进一步讨论算法细节。

a) 同步在线策略 RL

b) 异步离策略 RL

图 2 (a) 同步在线策略 RL 和 (b) 异步离策略 RL 的过程演示。对于异步 RL，生成器和训练器并行运行，互不阻塞，这与同步 RL 相反。这种设计显著加快了整体训练过程，且不影响模型质量。

4.3 细粒度并行与量化

支持灵活并行与量化以优化效率。从图 2(b) 的示例中可以看出，由于所有处理组的执行时间并未完美匹配，处理组1可能会出现 GPU 空闲气泡。为了进一步优化训练效率，必须支持灵活且细粒度的并行和量化。在 LlamaRL 中，不同处理组中的模型可以使用不同的并行策略和数据精度。我们提供以下几个例子：

数据并行（Data Parallelism）：回到协同定位策略，所有模型堆叠在一起，这意味着模型的数据并行（dp）大小必须相同。在 LlamaRL 中，模型的数据并行大小是解耦的，这为更好地对齐不同处理组的吞吐量和提高训练效率提供了灵活性。
模型并行（Model Parallelism）：模型并行在不同处理组之间是解耦的。根据经验，推理侧较小的模型并行（mp）大小（尤其当 mp > 8 时）可以显著减少节点间通信，从而减少生成时间。
流水线并行（Pipeline Parallelism）：流水线并行将模型计算分布在不同的流水线阶段。通过解耦处理组间的流水线并行（pp）大小，它为微调每个处理组的执行时间以实现更好的时间对齐提供了灵活性。
量化（Quantization）：在推理侧进行量化（fp8 或 fp4）允许模型以更小的 mp 大小进行生成，从而提供更快的生成速度。

功能实现。这些功能主要在通信通道中实现。更多细节可以在下一节找到。

5 系统设计

本节我们介绍 LlamaRL 的高层架构和接口。与依赖第三方编排系统（如 Ray【索引25，Ray: A Distributed Framework for Emerging AI Applications，2018，USENIX Association】）的 OpenRLHF 和 HybridFlow 等先前工作不同，LlamaRL 完全基于原生 PyTorch 实现。这一设计选择降低了复杂性，提升了可靠性，并简化了用户工作流。下面，我们将解释关键的架构组件、用户构建 RL 算法的接口，以及一个在 LlamaRL 中组装典型 RL 流水线的示例。

5.1 高层架构

图 3 LlamaRL 架构示例，展示了两个执行器之间通过通信通道进行交互，并由一个单一的控制器管理。

架构概览。如图 3 所示，LlamaRL 编排了多个执行器（executors），每个执行器负责 RL 流水线的一个或多个阶段，如策略推理、奖励计算和策略训练。这些执行器附加于分布式处理组，各自拥有自己的 GPU 数量和并行配置。由于其分布式特性，执行器依赖通信通道（communication channels）来在它们之间传递数据（如提示、生成的轨迹、奖励）和更新的模型权重。LlamaRL 并未依赖重量级的外部调度器，而是由一个单一的控制器（controller）负责根据需要启动、协调和同步执行器及数据流。

5.1.1 执行器（Executor）

分布式模型放置策略。像 PPO 这样的 RL 算法利用多个模型。对于这些模型的分配，一个直接的策略是将所有模型以相同的并行方式堆叠在所有 GPU 上，即协同定位策略。由于第 2 节中讨论的缺点，分布式模型放置策略更受青睐。

执行器作为自包含单元。LlamaRL 采用分布式策略，并进一步将分布式模型封装在执行器中。一个执行器是一个自包含的单元，运行在预定义数量的 GPU 上，并具有特定的并行配置。一个执行器的基类包含以下函数：
1. Init: 加载或构建模型，配置并行环境（例如，FSDP、流水线并行或张量并行组），并准备任何需要的数据加载器或分词器。
2. step: 执行一个 RL 步骤（例如，策略训练的一次前向/后向传播，或生成样本的一次前向传播）。
3. save_checkpoint: 定期保存模型状态、优化器状态和其他进度标记。
4. get_output / get_model: 暴露输出或内部状态（例如，最近生成的序列），以便其他执行器可以通过通信通道接收它们。

模块化设计。通过继承基类执行器，实现了不同的执行器类。例如，图 3 展示了用于 RLOO 算法的两个执行器：生成器（generator）和训练器（trainer），分别负责推理和训练阶段。基于规则的奖励模型包含在训练器中。通过将 RL 的各个阶段隔离到不同的执行器中，LlamaRL 使得对单个组件进行扩展或替换变得更加容易——例如，更改模型并行度或从 PPO 切换到另一种 RL 变体——而无需进行整体代码的修改。

5.1.2 通信通道（Communication Channel）

定义执行器间的通信。通信通道定义了执行器之间的一种有向通信范式。它包括：
* 一个名称，用于引用正在传输的数据或权重。
* 一个outbound_executor（发送方）和一个inbound_executor（接收方）。
* 一个communication_type，指示数据如何分发：
* BROADCAST: 出站数据被相同地发送到每个入站进程。
* SCATTER: 出站数据被分区并以块的形式分发到不同的入站进程。
* GATHER: 来自多个出站进程的数据被一个入站执行器聚合。
因此，执行器之间复杂的交互可以通过一个通信通道列表来定义。至关重要的是，它们还处理模型权重的更新，这是大规模 RL 中的一个核心挑战。

5.1.3 控制器（Controller）

编排训练流程。最后，控制器将执行器和通道组合成一个单一的训练过程，如算法1所示。在底层，它初始化分布式环境，在正确的 GPU ranks 上启动每个执行器，并通过以下方式编排每个训练“tick”：

为每个执行器设置迭代步骤（这样每个执行器都知道要处理哪个 mini-batch 或提示）。
按正确的顺序执行通信通道。例如，生成器执行器的输出被传递给奖励执行器，然后奖励执行器的输出流向策略训练器。
触发每个执行器的 step。在单次迭代中，多个执行器异步运行，各自执行自己的计算，无需等待其他执行器完成。
保存检查点。每个执行器定期独立地（或根据用户定义的触发器）写出模型或状态数据。

简洁的事件循环。因为每个执行器都是一个自主的 SPMD（单程序多数据）进程，所以控制器保持简洁且易于理解——本质上只是一个事件循环。

算法 1: ExecutorController 训练循环概述
输入: executor_group: List[Executor], communication_channels: List[CommunicationChannel], max_steps: int, init_communication_channels: Optional[List[CommunicationChannel]]
输出: 通过协调执行器来控制训练

1 Class: ExecutorController ;
2   executor_context: ExecutorContext ← 初始化一个上下文对象，存储不同执行器中不同进程间的所有 torch 分布式组。这些信息将被通信通道使用 ;
3   local_executor: Executor ← 当前进程本地的执行器实例。每个进程独立运行相同的控制器逻辑（SPMD 模型）。;
4
5 Function: init() ;
6   将 torch.distributed.group.WORLD 设置为本地执行器组;
7   调用 local_executor.init();
8
9 Function: run() ;
10  while local_executor.curr_step < max_steps do
11    调用 local_executor.set_step();
12    将 torch.distributed.group.WORLD 设置为全局组;
13    foreach channel in communication_channels do
14      调用 communicate(channel);
15    end
16    将 torch.distributed.group.WORLD 设置为本地组;
17    调用 local_executor.step();
18    调用 local_executor.save_checkpoint();
19  end
20  调用 executor_context.post_training_step() ;
21  调用 executor_context.shutdown();
22
23 Function: communicate(channel) ;
24  if local_executor == channel.outbound_executor then
25    调用 SEND_OPS[channel.communication_type](local_executor, channel);
26  end
27  if local_executor == channel.inbound_executor then
28    调用 RECV_OPS[channel.communication_type](local_executor, channel);
29  end

5.2 分布式直接内存访问（DDMA）权重更新

传统参数服务器的瓶颈。大规模 RL 的一个瓶颈是执行器之间模型权重的逐步同步，即将更新后的策略参数从训练执行器发送到推理执行器。传统方法通常在参数服务器（PS）中执行【索引23，Scaling distributed machine learning with the parameter server，2014，USENIX Association】，该服务器充当当前模型参数的键值存储。然而，传统的参数服务器架构存在两个主要瓶颈：（1）网络带宽不足——尤其是当 PS 与工作节点共同部署时【索引7，Project Adam: Building an efficient and scalable deep learning training system，2014，USENIX Association】；（2）无法有效扩展的低效软件栈【索引27，GPU Direct RDMA for infiniband on NVIDIA GPUs: A case study，2013，IEEE】。特别地，PS 软件栈中的数据复制、聚合和优化开销会阻碍吞吐量，即使在高带宽网络上也是如此。

DDMA 的优化。我们提出了一种分布式直接内存访问（DDMA）方法，专门用于大规模 RL 训练中语言模型的高效权重同步。如图 4 所示，DDMA 引入了针对上述瓶颈的关键优化：
1. 完全分布式: 每个 GPU 只存储或更新其分配的分片，利用与训练期间相同的张量和并行组。这消除了任何单个节点上的内存瓶颈，并实现了在内存和网络吞吐量方面的线性可扩展性，可扩展至数千个 GPU 卡和数 TB 的模型权重，权重更新延迟约为 2 秒。
2. 直接内存访问进行权重更新: 通过利用 NVIDIA 的节点内 NVLink 和节点间 Infiniband 互连，DDMA 支持零拷贝通信，允许设备间的 CUDA 内存区域直接进行内存传输——完全绕过 CPU 内存。这极大地减少了权重更新路径中的延迟，并缓解了因数据移动而导致的大规模场景下步长时间膨胀数百秒的瓶颈【索引17，Openrlhf: An easy-to-use, scalable and high-performance rlhf framework，2024，】。

图 4 通过分布式直接内存访问（DDMA）进行模型权重同步。

5.3 示例：使用卸载的奖励和生成器进行 RL

组装 LlamaRL 作业。算法 2 展示了如何将策略训练器、奖励计算器和生成器组装成一个 LlamaRL 作业。关键点包括：
* 执行器初始化: 每个执行器被分配一个不同的 GPU 进程组（伪代码第 4-7 行）。
* 通信通道: 我们定义了用于将更新的策略权重传递给生成器、将生成的序列传递给奖励模型，以及将带有奖励评分的序列传回给训练器的通道（第 10-16 行）。
* 卸载: 奖励模型和生成器在不同的 GPU 集上运行，以便每个模型可以专门化其并行策略。
* 控制器: 设置完成后，ExecutorController.run() 方法处理异步调度和重复迭代（第 18-19 行）。

模块化设计的优势。这种模块化设计使您能够换入替代的奖励模型，添加额外的模型（例如，评论家、参考策略），或实现不同的 RL 优化算法，而无需改变流水线的其余部分。

总结。通过将 RL 流水线的每个组件分离到其自己的执行器中，并使用通信通道进行数据交换，LlamaRL 提供了一个高度可扩展且灵活的框架。分布式方法与 DDMA 相结合，极大地缓解了 GPU 内存压力，并减少了 RL 流水线中的空闲时间。在接下来的章节中，我们将评估这些设计选择如何转化为实际的速度提升和扩展性能。

6 LlamaRL 离策略学习

框架与符号。在 LlamaRL 中，我们采用基于重要性采样的离策略学习框架来适应大规模异步训练过程。为了设定符号，我们令 $x \in X$ 为 RL 训练的提示集， $y \in Y$ 为可能的生成结果。生成结果 $y = y_{1:T}$ 由单个 token $(y_t)_{t=1}^T$ 组成，其中 $T$ 表示生成长度。下面，我们使用策略梯度算法举例说明。

离策略校正。学习器模型为 $\pi$，它以自回归方式定义了生成结果 $y=y_{1:T}$ 上的分布 $\pi(y|x) = \prod_{t=1}^T \pi(y_t|x, y_{1:t-1})$。生成结果由一个 actor 模型采样，该模型通过另一个网络以自回归方式进行推理 $y_t \sim \mu(\cdot | x, y_{1:t-1})$。由于异步训练，$\mu$ 与 $\pi$ 不同（例如，滞后一个更新步骤），因此这两个网络产生不同的分布。重要性采样比率 $\frac{\pi(y_t | x, y_{1:t-1})}{\mu(y_t | x, y_{1:t-1})}$ 用于在学习过程中调整两个网络之间的离策略差异。

学习器更新。在训练期间，生成结果 $y_t$ 及其概率 $\mu(y_t|x, y_{1:t-1})$ 从生成器传送到训练器。我们对学习器网络的最终更新是一个重要性采样加权的策略梯度：

$$\sum_{t=1}^{T} \min \left(\frac{\pi\left(y_{t} \mid x, y_{1: t-1}\right)}{\mu\left(y_{t} \mid x, y_{1: t-1}\right)}, \rho\right) \cdot A\left(x, y_{1: t}\right) \nabla \log \pi\left(y_{t} \mid x, y_{1: t-1}\right),$$

其中 $A(x, y_{1:t}) = r(x, y_{1:t}) - v(x, y_{1:t})$ 是逐 token 的优势估计。这里的 $r(x, y_{1:t})$ 是逐 token 的奖励，可以从奖励模型计算得出【索引8, Deep reinforcement learning from human preferences, 2017, Advances in Neural Information Processing Systems (NeurIPS); 索引26, Training language models to follow instructions with human feedback, 2022, arXiv】。对于许多应用，奖励可以通过基于规则的评分器计算【索引13, Rlef: Grounding code llms in execution feedback with reinforcement learning, 2024, arXiv; 索引24, Let’s verify step by step, 2023, arXiv】，从而为每个生成结果 $r(x, y)$ 得到一个单一的奖励。在接下来的部分中，我们将重点讨论这种情况。奖励通常会加上 KL 正则化项 $\lambda_{KL} \cdot D_{KL}(\pi(y|x), \pi_{base}(y|x))$，由参数 $\lambda_{KL}$ 调整，以惩罚与参考策略 $\pi_{base}$ 的偏离。

基线计算。基线 $v(x, y_{1:t})$ 用于减少方差，有几种计算方法。例如，它可以从一个评论家模型计算得出【索引29，Proximal policy optimization algorithms，2017，arXiv】，或使用分支 rollout【索引29, Trust region policy optimization, 2015, PMLR; 索引21, Vineppo: Unlocking rl potential for llm reasoning through refined credit assignment, 2024, arXiv】。另一种简单而有效的方法，也是我们下面将重点关注且与图2中工作流兼容的方法是，从单个提示 $x$ 出发，我们采样 $n$ 个生成结果，并将基线计算为这些生成结果 $y_i, i \in [n]$ 的平均奖励，奖励为 $r(x, y_i)$。这为整个 token 序列产生一个恒定的基线 $v(x, y_{1:t}) = \frac{1}{n} \sum_{i=1}^n r_i(x, y)$，这在许多应用中被证明是相当有效的【索引1, Back to basics: Revisiting reinforce style optimization for learning from human feedback in llms, 2024, arXiv; 索引30, Deepseekmath: Pushing the limits of mathematical reasoning in open language models, 2024, arXiv】。

重要性采样比率裁剪。值得注意的是，我们使用一个固定的常数 $\rho$ 来从上方裁剪重要性采样比率 $\frac{\pi(y_t | x, y_{1:t-1})}{\mu(y_t | x, y_{1:t-1})}$。这种裁剪策略在梯度估计中引入了偏差和方差的权衡：当 $\rho$ 被赋予一个较大的值时，估计的偏差较小，但由于重要性采样权重，方差较大。我们将整个算法称为 AIPO（Asynchronous Importance weighted Policy Optimization）。

缓解训练不稳定性。在实践中，我们发现裁剪常数 $\rho \in [2, 10]$ 似乎普遍有效，并有助于缓解大规模实验中出现的大部分训练不稳定性。这种训练不稳定性通常表现为训练性能的突然或缓慢下降，并且更常由于异步训练与模型架构、数据混合、多模态、额外的训练和推理优化等因素之间的复杂相互作用而出现。

与 PPO 裁剪的比较。对于熟悉该主题的读者来说，很明显上述校正与 PPO 中的双边裁剪【索引29，Proximal policy optimization algorithms，2017，arXiv】以及最近的 GRPO【索引30，Deepseekmath: Pushing the limits of mathematical reasoning in open language models，2024，arXiv】略有不同。我们早期的消融实验表明，单边裁剪在我们的设置中效果更好。我们还认为，不同的算法设计有其充分的理由：简而言之，PPO 裁剪基于置信域策略优化公式，这与 LlamaRL 中的异步训练设置有细微差别。我们在附录A中提供了额外的讨论。

7 LlamaRL 加速的理论证明

理论视角。在本节中，我们旨在从理论角度解释 LlamaRL 实现的效率增益。我们希望这能揭示加速的来源，并帮助从业者为这类训练系统建立一个良好的心智模型。

核心定理（非正式版）。更确切地说，我们将证明以下定理，为便于理解，此处以非正式形式陈述：
定理 7.1 (LlamaRL 的理论加速，定理 7.5 的非正式版本)：在相同的硬件预算和内存限制下，LlamaRL 的配置可以使其强化学习（RL）的执行速度严格快于任何传统同步 RL 框架的可能配置。

证明核心思想。定理 7.1 的证明核心是将寻找最优效率的过程构建为一个受 GPU 内存约束的优化问题。我们将对同步和异步框架进行数学建模，并为每个框架推导出闭式解。关键的洞见在于，LlamaRL 的异步设计将同步基线的共享内存约束分解为两个独立的约束——一个用于训练器，一个用于生成器。这消除了两个组件之间的相互依赖，允许每个组件被单独优化，从而产生更大的可行域。

基本量和符号。我们首先根据典型 RL 训练框架的结构引入一些基本量。然后，在现代硬件通常满足的合理假设下，我们提出了一个形式化的证明，表明我们的异步设计比其传统的同步对应方案产生严格更高的效率。下面，我们首先引入一组将在整个推导过程中使用的通用常数和符号。

定义 7.2. (GPU 的通用常数)
令 $G_0$ 表示可用 GPU 的总数，$B_0$ 表示全局批量大小，$M_0$ 表示每个 GPU 的最大可用内存，$W_0$ 表示单个模型副本的内存占用。令 $b_t$ 表示训练期间使用的微批量大小，$b_g$ 表示生成器批量大小，它指示了最大的解码并发数。最后，令 $m_t$ 和 $m_g$ 分别表示训练器和生成器的模型并行度（即每个模型实例的 GPU 数量）。

内存组成。在表2中，我们列出了在训练器和生成器中导致每个 GPU 内存主要消耗的模型组件（模型权重、优化器状态、梯度、隐藏激活、KV 缓存）。例如，对于模型并行度为 $m_t$、大小为 $W_0$ 的单个模型，每个 GPU 的内存消耗为 $W_0/m_t$。

简化分析的概念。我们引入一些有用的概念，以简化后续的陈述和分析。

表 2 典型 RL 训练迭代中主要组件的 CUDA 内存占用。其中，$A_t$ 和 $K_g$ 是仅由模型大小决定的常数，与批量大小或模型并行大小无关。

定义 7.3 (处理时间)
对于给定的设置，令 $\tau_t(b)$ 和 $\tau_g(b)$ 分别表示处理大小为 $b$ 的批次的训练和生成时间。此外，我们定义单位样本处理时间为：

请注意，处理时间与吞吐量成反比，即处理时间越短意味着吞吐量越大。我们现在引入支撑主要结果的关键假设：

假设 7.1 (批量大小扩展)
$\eta_t$ 和 $\eta_g$ 都是关于批量大小 $b$ 的单调递减函数。

假设的实证支持。这一假设得到了现代 GPU 上大规模深度学习的经验观察的充分支持。随着批量大小的增加，由于硬件利用率的提高——例如更高的算术强度、更高效的内存访问模式和更好的内核调度——单位样本处理时间趋于减少。在像 NVIDIA A100 和 H100 这样的架构上，更大的批量能更好地分摊计算和内存开销，特别是对于 GEMM 操作，导致总计算时间的次线性增长，从而使单位样本成本呈现下降趋势。

先前工作与验证。这一现象在先前关于高效训练和推理的工作中被广泛观察和报道，例如 DeepSpeed【索引28，Zero: Memory optimizations toward training trillion parameter models，2019，SC20: International Conference for High Performance Computing, Networking, Storage and Analysis】，他们在其中利用更大的批量大小来最大化 GPU 利用率并提高可扩展性。我们在我们的设置中通过广泛的批量大小范围对这一假设进行了经验验证，如图5所示。

RL 步长时间的定义。现在，我们通过正式定义给定 RL 训练框架的 RL 步长时间，将上述规模扩展讨论与 RL 训练加速联系起来：

定义 7.4 (RL 步长时间)
给定 RL 框架 F，我们将步长时间 T 定义为微批量大小 $(b_t, b_g)$ 和模型并行大小 $(m_t, m_g)$ 的函数。我们特别考虑：
* $T_{LlamaRL}(b_t, b_g, m_t, m_g, \theta)$，它允许训练器和生成器有独立的 mp_size 以及灵活的 GPU 分配。这里 $\theta \in (0, 1)$ 表示分配给训练器的 GPU 比例，剩余的 $(1-\theta) \cdot G_0$ 由生成器使用。
* $T_{baseline}(b_t, b_g, m)$，对应于同步基线，其中约束 $m_t = m_g \equiv m$ 成立。

定理 7.5 的正式陈述。在温和的假设下——具体来说，设备间通信开销可以忽略不计（我们的经验测量支持这一点），并且表2中列出的组件是 RL 训练期间 CUDA 内存消耗的主要来源——我们现在建立 LlamaRL 相对于同步基线的正式效率增益。

定理 7.5 (LlamaRL 的理论加速)
在定义 7.2 和假设 7.1 的设置下，存在满足 GPU 内存约束的 $b_t, b_g, m_t, m_g$，使得 $T_{LlamaRL}(b_t, b_g, m_t, m_g)$ 严格小于最小可能的 $T_{baseline}$。

图 5 70B 模型上批量大小扩展假设 7.1 的经验验证。左图：每 128 个样本的训练时间随着微批量大小的增加而减少。右图：每 64 次生成的完成时间随着最大解码并发数的增加而减少。两个结果都显示总处理时间的次线性增长，支持了单位样本时间（ηt, ηg）随批量大小增加而减少的假设。

备注 7.1。值得注意的是，该定理是在最小假设下建立的。特别是，我们没有对批量处理时间 $\tau$ 或单位样本处理时间 $\eta$（见定义7.3）假设任何特定的解析形式。唯一的要求是 $\eta$ 随批量大小单调递减——这一行为既在实践中被先前工作广泛报道，也在我们的实验中得到了经验验证，正如我们之前讨论的。

证明。我们首先计算每个框架的步长时间。在同步情况下，总步长时间是生成器和训练器时间的总和。相反，对于异步框架，步长时间由两者中较慢的一个决定。利用这一观察和定义7.3，我们可以将两种步长时间表示如下：

$$\begin{aligned} \begin{aligned} T_{\text{baseline}}(b_t, b_g, m) &= \frac{B_0 m}{G_0 b_t} \cdot \tau_t(b_t) + \frac{B_0 m}{G_0 b_g} \cdot \tau_g(b_g) \\ &= \frac{B_0}{G_0} \cdot m \cdot (\eta_t + \eta_g), \end{aligned} \end{aligned}$$

另一方面，根据表2，训练器每个 GPU 的内存占用由下式给出：

$$\frac{4 W_{0}}{m_{t}}+\frac{A_{t} b_{t}}{m_{t}}.$$

同样，生成器消耗：

现在结合方程（2）、（3）、（4）和（5），$T_{baseline}$ 的可行最小化解是以下约束优化问题的解：

$$\begin{aligned} \begin{aligned} &\min_{b_t,b_g,m} \quad \frac{B_0}{G_0} \cdot m \cdot (\eta_t(b_t) + \eta_g(b_g)) \\ &\text{subject to} \quad \frac{(4W_0 + A_t b_t) + (W_0 + K_g b_g)}{m} \le M_0. \end{aligned} \end{aligned}$$

同样，使用方程（3）、（4）和（5），$T_{LlamaRL}$ 的可行最小化解解决了：

现在，考虑问题（6）的一个最小化解，同样表示为 $(b_t^, b_g^, m^)$，它达到了最小的可行步长时间 $T_{baseline}^$。根据引理B.1，$(b_t^, b_g^, m^*)$ 满足 (12)，因此以下严格不等式成立：

$$\begin{aligned} \begin{cases} \dfrac{4W_0 + A_t b_t^*}{m^*} < M_0, \\ \dfrac{W_0 + K_g b_g^*}{m^*} < M_0. \end{cases} \end{aligned}$$

相反，考虑问题（7）的一个解，表示为 $(b_t^{}, b_g^{}, m_t^{}, m_g^{}, \theta^{**})$。根据引理B.2，它必须满足以下约束：

$$\begin{aligned} \begin{cases} m_t^{**} = \dfrac{(4W_0 + A_t b_t^{**})}{M_0}, \\ m_g^{**} = \dfrac{(W_0 + K_g b_g^{**})}{M_0}. \end{cases} \end{aligned}$$

接下来，为最优训练器和生成器时间定义变量 $T_t^{}$ 和 $T_g^{}$：

$$\begin{aligned} \left\{\begin{aligned} T_t^{**} &:= \min_{b_t} \frac{(4W_0 + A_t b_t) \eta_t}{M_0}, \\ T_g^{**} &:= \min_{b_g} \frac{(W_0 + K_g b_g) \eta_g}{M_0}. \end{aligned}\right. \end{aligned}$$

然后我们得到：

$$\begin{aligned} \begin{aligned} T_{\text {baseline }}^{*} & =\frac{B_{0}}{G_{0}} \cdot \frac{\left(4 W_{0}+A_{t} b_{t}^{*}\right)+\left(W_{0}+K_{g} b_{g}^{*}\right)}{M_{0}} \cdot\left(\eta_{t}^{*}+\eta_{g}^{*}\right) \\ & >_{(a)} \frac{B_{0}}{G_{0}} \cdot\left(\frac{\left(4 W_{0}+A_{t} b_{t}^{*}\right) \eta_{t}^{*}}{M_{0}}+\frac{\left(W_{0}+K_{g} b_{g}^{*}\right) \eta_{g}^{*}}{M_{0}}\right) \\ & >_{(b)} \frac{B_{0}}{G_{0}} \cdot\left(T_{t}^{* *}+T_{g}^{* *}\right) \\ & =\frac{B_{0}}{G_{0}} \cdot\left(\theta^{* *} \cdot \frac{T_{t}^{* *}}{\theta^{* *}}+\left(1-\theta^{* *}\right) \cdot \frac{T_{g}^{* *}}{1-\theta^{* *}}\right) \\ & ={ }_{(c)} \frac{B_{0}}{G_{0}} \cdot \max \left(\frac{T_{t}^{* *}}{\theta^{* *}}, \frac{T_{g}^{* *}}{1-\theta^{* *}}\right) \\ & ={ }_{(d)} \frac{B_{0}}{G_{0}} \cdot \max \left(\frac{\eta_{t}^{* *} m_{t}^{* *}}{\theta^{* *}}, \frac{\eta_{g}^{* *} m_{g}^{* *}}{1-\theta^{* *}}\right) \\ & =T_{\text {LLAMARL }}\left(b_{t}^{* *}, b_{g}^{* *}, m_{t}^{* *}, m_{g}^{* *}, \theta^{* *}\right) \end{aligned} \end{aligned}$$

其中不等式(a)源于直接的代数操作，不等式(b)源于定义(10)，等式(c)是引理B.3中(16)第三个恒等式的结果，等式(d)使用了引理B.3中(16)的前两个恒等式，最后一个等式由定义(3)得出。结合(11)的两边，我们完成了定理7.5的证明。 □

备注 7.2。最终推导（11）有助于识别效率增益的两个主要来源。首先，不等式（a）反映了异步设置中释放的 GPU 内存如何允许减少模型并行度，从而降低每个 GPU 的计算负载。然后，不等式（b）表明，通过解耦训练器和生成器的内存约束，异步框架可以进一步在每个组件内独立优化 mp_size 和微批量大小，以最大化整体步骤效率。

备注 7.3。为简化定理 7.5 的证明，我们有意做了一些保守的假设。例如，我们没有考虑使用量化来减小模型大小的可能性，这会进一步支持更小的 mp_size，也没有论证与减少 mp_size 相关的额外通信节省。重要的是，这些被忽略的因素中的每一个只会通过贡献更大的效率增益来加强我们的结论。

A4 实验环境

数据集：
- 训练: 使用 MATH 数据集【索引16，Measuring mathematical problem solving with the math dataset，2021，ArXiv】，这是一个包含问题与简短答案配对的数学推理数据集。
- 评估: 在 MATH 数据集的完整测试集、GSM8K 数据集【索引9，Training verifiers to solve math word problems，2021，ArXiv】以及一个为避免测试集污染而专门保留的 500 个问题的子集 MATH-500 上进行评估【索引24，Let’s verify step by step，2023，ArXiv】。
模型架构:
- 使用 LLaMA 3.1 模型家族的 8B、70B 和 405B 变体【索引11，The llama 3 herd of models，2024，ArXiv】。
- 优化器为 Adam，学习率固定为 $2 \times 10^{-7}$。
硬件配置:
- GPU: 所有实验均在 NVIDIA H100 GPU 上进行。
- GPU 数量:
  - 8B 和 70B 模型：使用 256 个 H100 GPU。
  - 405B 模型：使用 1024 个 H100 GPU。
- 具体的模型并行度（mp_size）和 GPU 分配见实验结果表3。
软件配置与基线:
- LlamaRL: 本文提出的异步 RL 框架。
- 基线 (Baseline): 采用公认的同步在线策略基线，类似于 DeepSpeed-Chat【索引41，Deepspeed-chat: Easy, fast and affordable rlhf training of chatgpt-like models at all scales，2023，】。
- 共享优化: LlamaRL 和基线都使用了相同的推理和训练优化技术。
  - 推理优化: 高效的 KV-cache 内存管理、优化的 CUDA 内核、CUDA 图【索引25，Accelerating PyTorch with CUDA graphs，2021，https://pytorch.org/blog/accelerating-pytorch-with-cuda-graphs/】和张量并行【索引32 ，Megatron-lm: Training multi-billion parameter language models using model parallelism，2019，arXiv】。
  - 训练优化: FSDP【索引42，Pytorch fsdp: Experiences on scaling fully sharded data parallel，2023，】。
- 核心区别: 基线将推理和策略模型部署在同一个 GPU 集群中，共享相同的模型并行度；而 LlamaRL 将这些组件解耦到不同的 GPU 集群，允许使用不同的并行策略和数据精度。

A4 实验结果

8.2 效率比较

实验内容: 本实验比较了 LlamaRL 和同步基线在 8B、70B 和 405B 模型上的 RL 单步耗时。为了公平比较，所有实验每个训练提示生成 n=4 个补全，每批次训练使用 512 个独立提示，从而保持全局批量大小为 2048。

实验结果 (表3): 在同等 GPU 预算下，LlamaRL 在 8B、70B 和 405B 模型上分别实现了 2.52倍、3.98倍和 10.7倍的单步加速。

分析与结论:

LlamaRL 表现出随模型规模增大的强大可扩展性。
在 405B 模型上的巨大增益，部分得益于在生成器上使用了 fp8 量化，但主要原因在于 LlamaRL 的并行灵活性。具体来说，它允许生成器和训练器使用不同的模型并行度（mp_size）。例如，可以将生成器的 mp_size 设为 8，而训练器保持 bf16 精度，mp_size 为 16，从而有效平衡生成和训练时间以最大化整体吞吐量。
相比之下，同步基线需要为生成器和训练器设置一个全局统一的 mp_size（例如 405B 模型为 64），即使对生成器进行量化，也无法改变这一限制，从而限制了其潜在的加速效果。
图 7 展示了 LlamaRL 相对于同步基线的效率增益随模型规模的变化。图中凸起的趋势表明，相对加速比随着模型规模的对数增长呈超线性增长，这预示了 LlamaRL 在未来训练更大语言模型方面的巨大潜力。
表 4 对比了 OpenRLHF 和 LlamaRL 的权重同步时间，LlamaRL 在 405B 模型上仅需 2.31 秒，远优于 OpenRLHF 在 70B 模型上的 111.65 秒，这得益于 DDMA 设计。

表 3 LlamaRL 与同步基线的 RL 速度比较。为确保公平比较，所有实验每个训练提示使用 n=4 个补全，每个训练批次使用 512 个独立提示，从而得到一个一致的全局批次大小 2048。我们为 8B 和 70B 模型使用 256 个 GPU，为 405B 模型使用 1024 个 GPU。对于基线，生成器和训练器都位于所有 GPU 上，因此未列出它们各自的 GPU 数量。对于最大解码并发数，我们通过将每个数据并行组的最大解码并发数乘以数据并行组大小来计算。

表 4 OpenRLHF 和 LlamaRL 中权重同步的时间消耗（秒）。

图 7 LlamaRL 在不同模型规模下相对于同步基线的效率增益。加速比随模型规模增大而增加，在 405B 时达到 10 倍以上。我们使用对数刻度的 x 轴以更好地可视化增长趋势。
图 8 离策略校正影响的消融实验。我们比较了在 LlamaRL 中，有和没有重要性采样校正的异步训练的稳定性。在更复杂的数据混合下，8B 和 70B 模型的训练倾向于出现零星的不稳定。相比之下，应用离策略校正稳定了所有模型规模的训练，证明了它们对于确保异步 RL 训练框架鲁棒性的必要性。

8.3 与在线策略学习的质量比较

实验内容: 本实验评估了不同 RL 训练框架下的训练质量。我们在 MATH 测试集、MATH-500 和 GSM8K 上进行评估，使用 sympy 分数作为主要指标。该分数通过解析预测和参考表达式并检查符号等价性来衡量答案的准确性。我们使用相同的分数作为训练数据的奖励。我们使用 LlamaRL 框架和 AIPO 算法训练 8B 模型，并与一个强大的同步 RL 基线进行比较。该基线可被视为 AIPO 的在线策略版本，使用相同的梯度估计器但没有重要性采样校正。所有超参数（包括全局批量大小和学习率）在两种设置中都保持一致。

实验结果 (图6): 在 MATH 测试集和 GSM8K 上，LlamaRL 的性能与传统的同步 RL 框架相当。有趣的是，LlamaRL 在整个训练过程中在 MATH-500 上显示出持续的改进。

分析与结论:
* 结果表明异步训练并未损害训练质量。
* 在 MATH-500 上的增益可能部分归因于统计变异，但也可能反映了离策略学习相对于严格在线策略方法的潜在优势。
* 本实验聚焦于 8B 模型，因为它在 MATH 训练集上能观察到更清晰的学习趋势，从而更有效地比较不同 RL 框架的训练质量。

图 6 在 MATH-500、MATH 测试集和 GSM8K 上比较 LlamaRL 与同步 RL 基线的评估结果。LlamaRL 在所有基准测试中均取得了与基线相当的性能，同时受益于更高效的异步训练框架。

离策略校正影响的消融实验

实验内容: 图 8（右侧部分）展示了在 8B 和 70B 模型上，有无重要性采样校正对异步训练稳定性的影响。

实验结果 (图8):
* 无校正: 训练过程出现零星的不稳定，表现为训练准确率的突然大幅下降。
* 有校正: 应用离策略校正后，8B 和 70B 模型的训练过程都变得稳定。

分析与结论: 该实验证明了离策略校正对于确保 LlamaRL 异步训练框架鲁棒性的必要性。

A5 结论

本文介绍了 LlamaRL，一个用于大规模训练大型语言模型（LLM）的全分布式异步强化学习框架。通过将异步训练与多样的并行策略及跨 RL 组件的混合精度技术相结合，LlamaRL 显著提升了计算和内存效率。我们提出的分布式直接内存访问（DDMA）方法进一步确保了 actor 之间的高效权重更新。因此，LlamaRL 展现出强大的可扩展性——从小型 H100 GPU 集群到数千甚至更多的 H100 GPU——为模型规模和训练需求持续增长的未来提供了一条切实可行的路径。

我们从理论和实践两方面证实了 LlamaRL 的效率增益。理论上，我们提供了形式化分析，并证明在相同资源约束下，LlamaRL 的步长时间严格低于任何同步 RL 框架。实验表明，离策略校正有效解决了异步学习的潜在缺点，保持了训练稳定性和最终性能。特别地，LlamaRL 在超大规模（405B）模型上实现了 10.7 倍的 RL 训练吞吐量提升。

总而言之，LlamaRL 标志着基于 RL 的 LLM 微调迈出了新的一步，简化了在规模上集成可能涉及多个模型的复杂算法。我们预计，我们的方法将通过在现代 AI 工作负载的 RL 流水线中提供更高的效率、增强的可靠性和更易的扩展性，使研究人员和从业者受益。未来的方向包括探索更先进的离策略学习方法，集成多任务目标，以及将异步 RL 扩展到更广泛的模态。

A6 附录

A 与 PPO 和 GRPO 的比较讨论

PPO 裁剪。从高层次看，PPO 对重要性采样比率应用了双边裁剪，并进一步根据优势函数的符号进行条件判断。其损失函数实现如下：

$$\min \left(\frac{\pi\left(y_t \mid x, y_{1: t-1}\right)}{\mu\left(y_t \mid x, y_{1: t-1}\right)} A_t, \underbrace{\operatorname{clip}\left(\frac{\pi\left(y_t \mid x, y_{1: t-1}\right)}{\mu\left(y_t \mid x, y_{1: t-1}\right)}, 1-\epsilon, 1+\epsilon\right) A_t}_{\text {double-sided clipping }}\right),$$

其中 $\epsilon > 0$ 是一个超参数。这个整体设计除了双边裁剪外，还有一些微妙的技术特性：(1) 梯度裁剪。最微妙的是，损失函数的实现方式使得当比率 $\frac{\pi(y_t | x, y_{1:t-1})}{\mu(y_t | x, y_{1:t-1})}$ 超出边界 $[1-\epsilon, 1+\epsilon]$ 时，梯度为零。(2) 优势函数符号。除了双边裁剪，PPO 还根据估计的优势函数的符号来决定裁剪方向，这体现在对两个损失项取最小值。这里的主要论点是构建策略性能的一个下界。GRPO 应用了与 PPO 相同的截断裁剪，但去除了学习评论家函数的需求。

PPO 裁剪深植于置信域策略优化。PPO 裁剪深深植根于用于稳定 RL 训练的置信域策略优化（TRPO）公式【索引20, Approximately optimal approximate reinforcement learning, 2002, Proceedings of the nineteenth international conference on machine learning; 索引29, Trust region policy optimization, 2015, PMLR】，其动机是确保参数更新足够保守以保证策略改进。在 PPO 的上下文中，行为策略 $\mu$ 具有一种特定形式：它是 $\pi$ 的前一次迭代，表示为 $\pi_{old}$。因此，裁剪确保了只有当 $\frac{\pi(y_t | x, y_{1:t-1})}{\pi_{old}(y_t | x, y_{1:t-1})}$ 在例如 $[1-\epsilon, 1+\epsilon]$ 范围内时，策略才有非零更新，即当 $\pi$ 与 $\pi_{old}$ 的偏差不大时。

与 LlamaRL 情况的关键区别。如果直接应用 TRPO 公式，我们会得到一个更同步的算法，即交替地从 $\pi_{old}$ 生成样本并更新策略到 $\pi$。在 LlamaRL 中，由于训练过程的异步性，行为策略 $\mu$ 不一定是 $\pi$ 的确切前一次迭代。例如，$\mu$ 可能与 $\pi$ 相差几次更新，此外还可能被量化或配置了与学习器网络不同的采样参数等。因此，$\mu$ 不一定是置信域正则化的理想策略。相反，一个更自然的动机可能是校正由 $\mu$ 生成的样本的离策略性，以构建对 $\pi$ 的近似在线策略更新。这通过裁剪重要性采样比率来实现，以平衡偏差和方差的权衡【索引12，Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures，2018，International conference on machine learning】。

我们早期的消融实验结果也表明，LlamaRL 中的单边裁剪与 PPO 裁剪效果相当或略好。然而，原则上，这两种裁剪策略可以结合起来，以协同 TRPO 和异步离策略学习的算法优势。

B LlamaRL 加速的理论证明引理

我们提出并证明了定理 7.5 所需的几个引理。

引理 B.1
令 $(b_t^*, b_g^*, m^*)$ 为问题 (6) 的一个最小化解。那么它必须满足：

引理 B.1 的证明。假设严格不等式成立：

$$\frac{(4W_0 + A_t b_t^*) + (W_0 + K_g b_g^*)}{m^*} < M_0,$$

这意味着仍有空间增加 $b_t$ 或 $b_g$。这允许我们进一步减小目标函数：

$$\frac{B_0}{G_0} \cdot m \cdot (\eta_t + \eta_g)$$

根据假设 7.1，$\eta_t, \eta_g$ 的单调递减性质，这与 $(b_t^*, b_g^*, m^*)$ 的最小化性质相矛盾。因此引理 B.1 得证。□

引理 B.2
令 $(b_t^*, b_g^*, m_t^*, m_g^*)$ 为问题 (7) 的一个最小化解。那么它必须满足：

引理 B.2 的证明与引理 B.1 的论证非常相似，因此省略。

引理 B.3
$T_t^{**}, T_g^{**}$ (见 (10)) 分别是以下两个优化问题的最小值：

$$\begin{aligned} \begin{aligned} & \min_{b_{t}} \quad \eta_{t}(b_{t})m_{t} \\ & subject \ to \quad \frac{4W_{0} + A_{t}b_{t}}{m_{t}} \le M_{0}, \end{aligned} \end{aligned}$$

此外，令 $(b_t^{**}, b_g^{**}, m_t^{**}, m_g^{**}, \theta^{**})$ 为 (7) 的一个解，则以下成立：

引理 B.3 的证明。第一个论断直接由引理 B.2 证明的类似论证得出。我们省略细节。首先，我们声称：

假设等式不成立，那么不失一般性，我们可以假设：

$$\frac{\eta_{t}^{**} m_{t}^{**}}{\theta^{**}} > \frac{\eta_{g}^{**} m_{g}^{**}}{1-\theta^{**}}.$$

很明显，通过稍微增加 $\theta^{**} \rightarrow \theta^{**} + \epsilon$，不等式 (18) 仍然成立，但以下成立：

$$\max(\frac{\eta_{t}^{**}m_{t}^{**}}{\theta^{**}},\frac{\eta_{g}^{**}m_{g}^{**}}{1-\theta^{**}}) = \frac{\eta_{t}^{**}m_{t}^{**}}{\theta^{**}} > \frac{\eta_{t}^{**}m_{t}^{**}}{\theta^{**}+\epsilon} = \max(\frac{\eta_{t}^{**}m_{t}^{**}}{\theta^{**}+\epsilon},\frac{\eta_{g}^{**}m_{g}^{**}}{1-\theta^{**}-\epsilon}),$$

这与 $(b_t^{**}, b_g^{**}, m_t^{**}, m_g^{**}, \theta^{**})$ 的最小化性质相矛盾。为了证明 (16) 的前两个恒等式，我们需要证明 $\eta_t m_t$ 和 $\eta_g m_g$ 都在 $(b_t^{**}, b_g^{**}, m_t^{**}, m_g^{**}, \theta^{**})$ 处达到最小。再次假设它们在不改变其他变量的情况下变化。不失一般性，让我们假设 $\eta_t^{**} m_t^{**}$ 不是最小的。这意味着存在 $b', m'$，使得：

$$\eta_{t}^{\prime} m_{t}^{\prime}<\eta_{t}^{* *} m_{t}^{* *}$$

选择足够小的 $\epsilon > 0$，使得：

此外，我们有：

$$\frac{\eta_{g}^{* *} m_{g}^{* *}}{1-\theta^{* *}-\epsilon}<\frac{\eta_{g}^{* *} m_{g}^{* *}}{1-\theta^{* *}} .$$

结合 (21) 和 (22)，可以得出：

$$\max(\frac{\eta_{t}^{\prime} m_{t}^{\prime}}{\theta^{**}+\epsilon}, \frac{\eta_{g}^{**} m_{g}^{**}}{1-\theta^{**}-\epsilon})<\max(\frac{\eta_{t}^{**} m_{t}^{**}}{\theta^{**}}, \frac{\eta_{g}^{**} m_{g}^{**}}{1-\theta^{**}}).$$

也就是说，我们找到了一个新的点 $(b_t', b_g^{**}, m_t', m_g^{**}, \theta^{**} + \epsilon)$，它为问题 (7) 产生了一个更小的值，这与

$$(b_{t}^{**}, b_{g}^{**}, m_{t}^{**}, m_{g}^{**}, \theta^{**}).$$

的最小化性质相矛盾。因此我们证明了 (16) 的前两个恒等式。(16) 的第三个恒等式现在直接由 (17) 得出。□

PaperCache

LlamaRL: A Distributed Asynchronous Reinforcement Learning Framework for Efficient Large-scale LLM Training

LlamaRL: A Distributed Asynchronous Reinforcement Learning Framework for Efficient Large-scale LLM Training

A1 主要贡献

A3 背景知识

相关工作