StreamRL: Scalable, Heterogeneous, and Elastic RL for LLMs with Disaggregated Stream Generation

作者/机构: Yinmin Zhong1, Zili Zhang1, Xiaoniu Song2, Hanpeng Hu2, Chao Jin1, Bingyang Wu1, Nuo Chen2, Yukun Chen2, Yu Zhou2, Changyi Wan2, Hongyu Zhou2, Yimin Jiang3, Yibo Zhu2, Daxin Jiang2
(1 School of Computer Science, Peking University; 2 StepFun; 3 Unaffiliated)

A1 主要贡献

本文探讨了大型语言模型(LLM)的强化学习(RL)训练架构。RL训练通常包括在线生成样本(generation)和使用这些样本进行训练(training)两个阶段。

核心问题:
主流的“同地部署(colocated)”架构(即生成和训练阶段在同一组GPU资源上分时复用)虽然解决了资源闲置问题,但在大规模训练中暴露了“资源耦合”的根本缺陷。生成阶段是内存带宽密集型,而训练阶段是计算密集型,两种负载特性迥异。同地部署强制它们使用相同数量和类型的硬件,导致资源利用率低下,无法为各阶段选择最具成本效益的硬件。此外,这种架构难以扩展到跨数据中心的异构资源池,因为训练阶段的全连接通信(full-mesh communication)会产生巨大的跨数据中心通信开销。

研究目标:
本文重新审视并优化了“分离式(disaggregated)”架构,即为生成和训练阶段分配独立的专用资源。分离式架构天然具有资源分配灵活、支持异构硬件、可扩展至跨数据中心的优势。然而,现有分离式框架存在两大性能瓶颈:
1. 流水线气泡(Pipeline bubbles): 因阶段间的串行依赖导致一个阶段工作时另一阶段资源空闲。
2. 倾斜气泡(Skewness bubbles): LLM生成样本的输出长度呈长尾分布,导致在生成后期只有少数长尾样本在处理,GPU利用率严重不足。

创新点 (StreamRL):
为了解决上述问题,本文提出了一个基于分离式架构设计的RL训练框架——StreamRL。其核心创新点如下:
1. 提出流式生成(Stream Generation): 将生成和训练阶段抽象为流式生成服务(SGS)和训练器(Trainer)。SGS一旦生成一个样本,就以流的方式立即返回给Trainer,从而打破了传统的阶段边界。这使得:
* 在同步RL算法中,可以实现“动态批次流水线(dynamic-batch pipelining)”,有效减少流水线气泡。
* 在异步RL算法中,可以实现“完全异步流水线(fully asynchronous pipelining)”,几乎完全重叠两个阶段的执行,最大化资源利用率。
2. 解决长尾倾斜问题: 提出了一套倾斜感知调度机制。
* 输出长度排序器模型(Output-length ranker model): 训练一个轻量级模型来预测并排序输入提示(prompts)的生成长度,以识别长尾样本。
* 倾斜感知分发与调度(Skewness-aware dispatching and scheduling): 根据排序结果,将长尾样本和常规样本分派到不同的生成实例上,并为长尾样本分配更小的批次大小以加速其处理,从而减少整体生成延迟。
3. 实现弹性与异构部署:
* 自动资源分配与动态调整: 提供了基于性能分析的资源分配算法,用于在训练前平衡两个阶段的执行时间。同时,设计了动态资源调整机制,在训练过程中根据负载变化(如生成长度增加)自动扩展SGS的资源,维持阶段平衡。
* 支持异构和跨数据中心部署: StreamRL的设计充分利用了分离式架构的优势,实验证明在异构、跨数据中心的环境下,能将生成阶段部署在成本效益更高的推理GPU上,将训练阶段部署在高性能计算GPU上,从而将成本效益提升高达1.33倍。

实验结果表明,与现有最先进的系统相比,StreamRL的吞吐量提升高达2.66倍。


图 1. 两种代表性的RL框架架构。(a) 分离式架构,为生成和训练阶段分配专用资源,但存在资源闲置问题。(b) 同地部署架构,两阶段分时复用同一组资源,解决了闲置问题但引入了资源耦合。

A3 背景知识与关键动机

2.1 背景知识

2.2 同地部署架构的问题


图 2. 生成和训练阶段在不同资源数量(左)和序列长度(右)下的性能敏感性差异。

Table 1. NVIDIA GPU规格.

2.3 分离式架构的动机

2.4 分离式架构的挑战


图 3. 分离式架构中的资源浪费。

A2 方法细节

3 StreamRL 概览


图 4. StreamRL系统架构。

4 解决流水线气泡

4.1 重叠设计


图 5. 流式技术如何增强现有方案以更好地缓解流水线气泡。

4.2 阶段平衡

Algorithm 1 资源分配算法

5 解决倾斜气泡

本节介绍SGS用于处理倾斜气泡的技术,其目标是在给定资源下最小化生成时间。

5.1 问题与机遇


图 6. 左:倾斜感知分发相对于随机分发的优势。右:7B LLM在NVIDIA H800上使用vLLM分析的每token解码延迟随批次大小增加的趋势。

5.2 输出长度排序器

5.3 倾斜感知调度

Algorithm 2 倾斜感知分发算法

6 实现

A4 实验

实验环境

实验结果

A5 结论

本文重新审视了用于RL训练的分离式架构,并强调了其相对于广泛采用的同地部署架构的显著优势:灵活的资源分配、对异构硬件的支持以及跨数据中心的可扩展性。为了充分发挥分离式架构的潜力,我们提出了StreamRL框架,它通过流式生成和倾斜感知调度有效解决了现有分离式RL框架中存在的流水线气泡和由数据倾斜导致的效率低下问题。实验表明,与当前最先进的RL框架相比,StreamRL实现了高达2.66倍的加速。我们希望这项工作能鼓励社区重新审视分离式架构,并对其有效性有更深入的理解。