Skip to content

Ray与LLM强化学习框架设计

大型语言模型(LLM)强化学习框架发展迅速,Ray作为分布式计算框架在LLM训练各阶段中,特别是在RL(强化学习)阶段的应用日益广泛。本文旨在梳理Ray架构在LLM强化学习中的设计理念与演进脉络,并探讨当前主流框架的实现方式与技术挑战。

从Google Pathways看分布式RL系统设计

在讨论Ray与LLM强化学习框架之前,我们有必要从Google的Pathways系统开始,这一系统于2021年提出,作为下一代AI架构和分布式机器学习平台的代表 。Pathways系统的核心设计包括单控制器架构(Single-Controller)和多程序多数据范式(MPMD) 。

单控制器架构是指用一个中央协调器管理整个分布式计算流程,负责任务分发、资源调度和状态监控。这种设计模式在Ray中通过Driver Process实现,而在PyTorch DDP中则采用多控制器模式,每个节点独立执行程序 。

多程序多数据范式(MPMD)与传统的单程序多数据范式(SPMD)形成对比。SPMD模式下,所有节点运行相同程序处理不同数据,如PyTorch DDP;而MPMD模式则允许不同节点运行不同程序,处理不同数据 。在大模型训练中,尤其是流水线并行等复杂场景,MPMD模式更为适用 。

MPMD系统的复杂性主要在于各组件间的协调与同步。单控制器架构通过中心化的控制器简化了这一过程,统一管理分布式计算流程的关键环节 。这一设计理念与LLM强化学习场景高度契合,因为LLM强化学习本质上是一个涉及多阶段、多模型的复杂分布式任务。

典型的LLM强化学习(RLHF)流程包括三个关键阶段:

  1. 生成阶段:当前策略模型(LLM)对一批输入提示生成响应文本
  2. 评估阶段:这些响应由奖励模型评分,或通过人类/自动化偏好模型进行比较评估
  3. 训练阶段:基于获得的奖励信号更新策略模型权重,可能包括价值函数或评论家网络的更新

在早期实现中,这些阶段往往串行执行,导致大量上下文切换和资源利用率低下。因此,Pathways架构的启发在于:在保证正确性的前提下,通过良好的系统设计尽可能重叠和并行化这些工作阶段,以最大化计算资源的利用效率

下表展示了两种RLHF系统架构的对比:DeepSpeed-Chat采用SPMD串行方式,而OpenRLHF则是典型的MPMD系统 。

架构特性DeepSpeed-Chat (SPMD)OpenRLHF (MPMD)
计算范式单程序多数据,各节点执行相同代码多程序多数据,各节点可执行不同代码
资源管理静态分配,资源利用率较低动态调度,资源利用率较高
计算重叠有限,阶段间存在空闲较高,阶段间可部分重叠
扩展性受限于单阶段资源需求更灵活,支持不同阶段独立扩展

Ray与LLM强化学习框架的架构设计

Ray框架的核心设计使其非常适合开发基于单控制器+MPMD的程序,自然地适应了LLM强化学习的场景需求。事实上,社区已经基于Ray开发了大量强化学习框架,主要包括两种架构模式:共置式(Task-Collocated)架构分离式(Task-Separated)架构

img

共置式架构意味着将生成阶段和训练阶段部署在同一节点上,形成时分复用系统;而分离式架构则允许部分或全部计算任务部署在不同设备上,形成空分复用系统 。

在共置式架构中,Ray通过资源组(Placement Group)实现资源分割,例如在OpenRLHF框架中,可能将每个GPU的0.75分配给训练Actor,0.25分配给生成Actor,从而实现资源共享 。这种架构的优势在于减少GPU空闲时间,减少模型offload频率,同时尽量并行化不同节点的执行,提高资源利用效率 。

然而,随着模型规模和集群规模的增长,共置式架构也显示出局限性:

  • 资源耦合问题:虽然共置式架构比SPMD程序提升了计算任务的并行性,但生成和训练共享相同设备仍导致无法独立扩展或为每个阶段定制资源。训练任务(计算密集)和生成任务(IO密集)的瓶颈不同,不利于GPU资源的高效利用 。
  • 长尾任务问题:LLM生成的文本长度不固定,尤其在思考模型流行的情况下,生成任务中不同组的模型生成结果时间可能差异很大。例如,使用32块GPU时,若每4块GPU为一组进行生成,其中一组生成任务过长,会导致其他28块GPU空闲造成资源浪费 。

因此,共置式框架虽然成熟稳定,但随着模型规模扩大,其可扩展性受限。这为下一代RL框架的发展指明了方向:能否通过打破严格的串行约束,让生成和训练阶段真正独立并行执行?

Online Policy与Offline Policy的权衡

在LLM强化学习框架中,Online PolicyOffline Policy是两种关键策略,它们直接影响系统效率与算法收敛性 。

Online Policy严格要求使用当前最新策略生成的数据进行训练。在RLHF中,这意味着每次训练迭代必须等待当前Actor模型完成新一轮文本生成,然后立即使用这些"新鲜"样本来更新模型权重 。这种方式保证了训练数据与当前策略的一致性,但代价是强制同步等待,导致计算空泡(pipeline bubbles)。

Offline Policy则允许使用稍旧的策略生成的数据进行训练。在实际系统中,这意味着训练阶段可以使用来自之前几个迭代版本的Actor模型生成的样本,而无需严格等待当前版本生成完成 。虽然引入了策略陈旧性(policy staleness),但这种设计允许生成和训练阶段真正并行化,大幅提升系统吞吐量 。

从理论角度看,Online Policy的样本效率确实优于Offline Policy,这主要源于其保证训练数据与当前策略分布的一致性。在LLM强化学习场景中,策略陈旧性会引入分布偏移问题——生成数据的策略分布与当前训练策略产生偏差,这种不匹配可能降低训练样本的有效性,影响模型收敛的稳定性和效率 。

然而,在工业实践中,这种理论差异往往被系统吞吐量的提升所抵消。现代LLM RL系统通过以下方式缓解策略陈旧性的影响:

  1. 增加经验缓冲区容量:存储更多历史样本,减少策略陈旧性的影响 。
  2. 优化批采样策略:平衡新旧样本的使用比例 。
  3. 采用更频繁的模型同步机制:减少陈旧策略的使用时间 。

实际部署中,有团队发现通过合理调整超参数(如学习率衰减、梯度裁剪),Offline Policy系统在保持训练稳定性的同时,能够获得成倍的训练吞吐量提升。这种工程权衡推动了RLHF向多轮迭代和DPO向Iterative DPO等更适合并行化的算法变种演进。

Online Policy的限制在共置式框架中尤为明显。当使用最新策略生成的数据进行训练时,训练步骤必须等待生成步骤完成,导致在迭代之间形成大量空泡 。而分离式框架则允许生成和训练阶段真正并行,从而消除这一问题 。

Disaggregated架构:从Offline Policy到Streaming RL

近年来,分离式架构逐渐成为主流 。根据统计,2025年发布的主流框架如Kimi、SeedThinking-1.5、StreamRL、AReaL等均采用分离式架构,而早期的DeepSpeed-Chat、NeMo-Aligner、RLHFuse、verl等则多采用共置式架构 。

框架发布时间架构
DeepSpeed-Chat2023Task-Collocated
NeMo-Aligner2024Task-Collocated
RLHFuse2024Task-Collocated
verl2024Task-Collocated
OpenRLHF2024Task-Separated
K1.52025Task-Collocated
SeedThinking-1.52025Task-Separated
StreamRL2025Task-Separated
AReal2025Task-Separated
slime2025Task-Separated/Collocated
siiRL2025Task-Separated

分离式架构成为主流的原因可归结为两点:提高效率降低成本

首先,训练与推理的计算负载特性明显不同。训练是计算密集型任务,而推理解码阶段则是访存密集型任务。两者在集群规模扩大和序列长度增加时,资源需求差距进一步扩大。业界开始探索在LLM强化学习中采用Offline Policy算法的可能性,通过引入策略陈旧性来换取更高的系统吞吐量。

ASYNCHRONOUS RLHF(arXiv:2410.18252)的研究表明,适度的策略陈旧性是可以接受的,并且不会显著影响训练效果 。这一发现随后在StreamRL(arXiv:2504.15930)的工作中得到验证和扩展 。Meta在LlamaRL(arXiv:2505.24034)中提出了新的计算流,形成了流式的分离式架构 。

分离式架构的核心优势在于异构资源的独立扩展能力。例如,若发现生成阶段成为瓶颈(如复杂推理任务需要更长生成时间),可单独为生成服务增加更多GPU,而不影响训练集群配置;反之,若奖励模型计算或PPO更新成为瓶颈,可针对性扩展训练侧资源。这种弹性伸缩能力在云环境和多租户场景中尤为宝贵 。

通过接受一定程度的陈旧策略数据,流式框架显著改善了资源利用率。不需要所有GPU在迭代间同步暂停;相反,生成和训练可达到稳定吞吐量。StreamRL明确解决了同步设计的流水线"空泡"和落后任务的长尾问题,实现了完全异步流水线,使权重更新、生成和训练尽可能重叠 。

更进一步,为协调训练集群和推理集群,通常引入数据缓冲区或队列作为交互点。这种设计可实现生成和训练的近乎完美重叠,消除大部分空闲时间。AsyncFlow(arXiv:2507.01663)框架采用TransferQueue传输数据并控制训练执行 。

参数延迟更新的一步异步Off-Policy算法

一步异步算法指的是训练和推理模型最多只差一个版本。具体而言,训练最新模型时的输入数据,是由上一个版本的推理模型所生成 。实现这一异步的关键在于调整批大小和参数传输机制。

通过令推理阶段的全局批大小为训练的K倍,可实现K-1步的Off-Policy算法。然而,由于算法收敛性限制,K不能无限增大,导致空泡占比无法进一步缩小。因此,AReaL框架提出了一种创新的异步参数传输机制。

当训练集群进行参数更新后,最新模型参数会异步传输至推理集群的CPU侧内存,并做好Resharding准备。推理进程正常进行,待全局批推理任务完成后,触发一次H2D操作(主机到设备内存传输),将新模型加载至显卡上 。由于机内H2D带宽通常高于网络传输带宽,这种方式可有效掩盖训推参数同步的开销,发挥分离式架构的优势 。

半步异步算法:优化模型收敛与系统吞吐

"一步异步"算法看似已将Off-Policy对收敛性的影响降至最低,但仍有改进空间。在实际训练中,为保证流水线正常运转,需满足推理总吞吐≥训练总吞吐 。这意味着有机会让部分推理实例轮流更新最新权重,而其他实例继续使用旧权重推理,维持训练流水不中断

这种"半步异步"算法有两个主要优势:

  1. 降低策略陈旧性影响:在一个训练全局批中,样本部分来自最新模型权重,部分来自旧权重,相比一步异步算法中整个全局批都由旧权重生成,这种设计有助于提高模型效果。
  2. 优化参数加载开销:不同推理实例的参数更新时机不同,使得参数加载的H2D时间可被掩盖,从而提高训练效率 。

受限于时间精力,这一策略尚未被广泛实现,但仍是一个值得探讨的方向 。

Ray技术挑战:从快速原型到生产部署

随着LLM强化学习框架的演进,Ray对框架的依赖正在逐渐减少。Ray最初是这些框架快速原型化的理想选择,因为它能够轻松处理集群配置、进程启动,并提供简洁的远程函数调用和Actor类API 。例如,OpenRLHF、VeRL、Slime和AsyncFlow等框架都将Ray作为关键"胶水层"来协调复杂的训练循环 。

然而,随着这些系统在生产环境中的大规模部署,Ray的固有限制开始显现:

调试复杂性是首要挑战。当远程Worker深层出现异常时,通常只能收到模糊的错误信息,难以追踪问题根源。这在大规模分布式训练中尤为棘手,因为问题可能出现在数百个节点中的任意一个 。

通信开销是另一个关键瓶颈。Ray的核心设计依赖Python对象序列化和gRPC通信,这在RLHF场景下面临巨大压力。每个经验批次可能包含数万个token的生成文本、完整的logits分布以及各种奖励和价值函数输出,数据体积可观。当这些数据需要在生成阶段和训练阶段之间频繁传输时,序列化开销变得不可忽视 。

此外,资源管理机制也面临挑战。Ray通过资源组(Placement Group)实现资源分割,但显存分割需依赖NVIDIA的MIG(硬件级)或MPS(用户态)技术,而非Ray原生功能。例如,OpenRLHF可能通过CUDA Hook实现显存共享,但这增加了系统复杂性 。

分离式架构的调度优化:LPT算法与动态批处理

StreamRL框架采用最长处理时间优先(LPT)算法结合动态批处理解决分离式架构的调度问题 。LPT算法是一种近似算法,用于解决作业调度问题,其调度长度是最优调度长度的4/3-1/(3m) 。

StreamRL的LPT调度算法实现如下:

  1. 按输出长度降序排序:利用输出长度排序器模型(通过监督微调训练)对输入提示的输出长度进行估计和排序。
  2. 偏态感知调度:将提示按估计输出长度排序,标记长尾样本,为其分配专门计算资源和较小批量大小,常规样本组成大批次,充分利用GPU资源。
  3. 动态批处理流水线:摒弃传统批量生成方式,样本一完成就立即发送给训练阶段,训练阶段根据生成速度进行动态批处理,减少训练阶段空闲时间 。

这一设计解决了分离式架构中的两个主要问题:流水线气泡偏态气泡。流水线气泡源于两阶段串行执行,而偏态气泡则因LLM推理工作负载中长尾输出长度分布导致。通过LPT调度和动态批处理,StreamRL在异构、跨数据中心场景下实现了最高2.66倍的吞吐量提升 。

AReaL的Decoupled PPO Objective

AReaL框架提出了一种解耦的PPO目标函数(Decoupled PPO Objective),从算法层面解决训推模型参数不同导致的收敛性问题 。这一创新直接针对策略陈旧性问题,允许使用不同版本策略生成的数据而不显著影响收敛性 。

传统PPO的actor loss定义为: $$ L_{\text{PPO} }(\theta) = \mathbb{E}\left[ \frac{\pi_\theta(a|s)}{\pi_{\text{old} }(a|s)} A(s,a) \right] - \beta \cdot \text{KL}(\pi_\theta \parallel \pi_{\text{old} }) $$ 其中,$\pi_{\text{old} }$是模型更新前的策略,$A(s,a)$是优势函数,$\beta$是KL散度权重 。

AReaL的Decoupled PPO通过引入行为策略($\pi_{\text{behav} }$)和近端策略($\pi_{\text{prox} }$)进行重构: $$ L_{\text{decoupled} }(\theta) = \mathbb{E}\left[ \frac{\pi_\theta(a|s)}{\pi_{\text{behav} }(a|s)} A(s,a) \right] - \beta \cdot \text{KL}(\pi_\theta \parallel \pi_{\text{prox} }) $$ 其中,$\pi_{\text{behav} }$代表实际生成策略,$\pi_{\text{prox} }$作为近端策略,用于规范$\pi_\theta$的更新 。

这一重构的关键在于:即使一条序列的响应是由不同权重生成的,也可视为来自某个统一的行为策略 。AReaL团队证明,一个轨迹中不一致的策略版本等价于单一行为策略$\pi_{\text{behav} }$ 。

在实际实现中,AReaL将$\pi_{\text{behav} }$相关的logP计算放在rollout后端,其余部分放在trainer后端。由于不同框架间存在精度差异,这一分离可能导致精度问题,但通过实验验证,即使策略陈旧性达8步,模型性能仍保持稳定 。

分离式架构的未来趋势

分离式架构+Offline Policy的实现与MPMD范式一致,代表了LLM强化学习系统的发展趋势。这一趋势体现在以下几个方面:

系统复杂性与性能的平衡:分离式架构虽然增加了系统复杂性,但通过精心设计(如AsyncFlow、StreamRL),这些问题变得可控。例如,StreamRL通过C++推理引擎、CUDA内核和RL-RPC通信框架(支持GPU-Direct RDMA零拷贝张量传输)优化通信效率,同时具备TCP fallback机制确保兼容性 。

跨数据中心训练:分离式架构允许生成和训练服务部署在不同数据中心,通过点到点链路连接,实现灵活的资源分配和异构硬件选择。在异构、跨数据中心设置下,StreamRL的成本效益最高提升1.33倍 。

动态工作负载管理:LLM生成任务的不固定长度和思考模型的流行,使得动态工作负载管理变得至关重要。分离式架构通过独立扩展生成和训练资源,更好地应对这一挑战 。

算法与系统协同优化:如AReaL的解耦PPO目标函数与系统级优化(解耦CPU与GPU操作、并发请求处理、无填充的序列打包策略)相结合,实现了高达2.57倍的训练吞吐量提升,同时在512 GPU上展现出近似线性扩展效率 。

总结与展望

Ray框架凭借其单控制器+MPMD的设计理念,为LLM强化学习提供了理想的分布式计算基础。随着模型规模扩大和训练需求增长,分离式架构逐渐取代共置式架构成为主流,通过接受一定程度的策略陈旧性,显著提高了系统吞吐量和资源利用率 。

从算法角度看,解耦的PPO目标函数(如AReaL框架)和半步异步算法(如Kimi的Partial Rollout技术)为处理策略陈旧性提供了有效解决方案 。这些算法创新与系统设计的结合,使得LLM强化学习训练在保持模型性能的同时,实现了显著的加速。

然而,Ray框架在调试复杂性通信开销方面仍面临挑战 。随着LLM强化学习框架的演进,Ray的依赖可能减少,但其设计理念和分布式计算能力将继续影响这一领域的发展。

未来发展方向可能包括:

  1. 更细粒度的资源管理:结合硬件级显存分割技术(如NVIDIA MIG)与Ray的资源组机制,实现更高效的资源利用 。
  2. 更智能的调度算法:如StreamRL的LPT调度与动态批处理结合,进一步优化生成和训练阶段的并行效率。
  3. 更稳定的异步算法:如AReaL的解耦PPO和半步异步算法,平衡策略陈旧性与模型收敛性。
  4. 跨数据中心训练优化:针对分离式架构在跨数据中心部署的场景,优化通信效率和数据同步机制 。

Ray与LLM强化学习框架的结合,不仅提高了训练效率,也为AI技术的规模化应用提供了新的可能性。随着这一领域的持续发展,我们有望看到更多创新的架构和算法,进一步推动LLM强化学习的效率与效果提升。

Maintained by Robin