本文提出 StreamRL 框架,通过分离式流生成架构优化 RL 训练,解决了流水线和偏斜气泡问题,提高了 LLMs RL 训练的吞吐量和成本效率。
Reinforcement Learning, Large Language Model, Efficiency, Heterogeneous Systems, Resource Allocation, Scalability
Yinmin Zhong, Zili Zhang, Xiaoniu Song, Hanpeng Hu, Chao Jin, Bingyang Wu, Nuo Chen, Yukun Chen, Yu Zhou, Changyi Wan, Hongyu Zhou, Yimin Jiang, Yibo Zhu, Daxin Jiang
Peking University, StepFun, Unaffiliated
Generated by grok-3-mini-latest
Background Problem
强化学习(RL)已成为大型语言模型(LLMs)的核心后训练技术,但现有的共置架构存在资源耦合问题,导致可扩展性和成本效率低下。具体来说,共置架构要求生成和训练阶段共享相同的资源和硬件类型,这与两阶段工作负载的根本差异相冲突:生成阶段是内存带宽绑定的,而训练阶段是计算绑定的。这种耦合在大规模训练中导致资源利用率低下,并且难以利用异构硬件或跨数据中心资源。本文重新审视分离式架构,以解决这些问题,但分离式架构也面临流水线气泡(由于阶段依赖导致资源空闲)和偏斜气泡(由于输出长度长尾分布导致GPU利用率低下)的挑战。
Method
- 核心思想: StreamRL 基于分离式架构设计,旨在通过流生成服务(SGS)和训练器(Trainer)解耦生成和训练阶段,最大化资源利用率。
- 工作原理: StreamRL 将 RL 训练分为生成和训练两个阶段,使用流生成技术实时返回完成样本,实现阶段重叠。针对同步 RL,通过动态批处理流水线减少流水线气泡;针对异步 RL,实现完全重叠。使用基于分析器的资源分配算法平衡阶段执行时间,包括静态配置和动态调整机制。针对偏斜气泡,引入输出长度排名模型预测长尾样本,并采用偏斜感知调度和分派策略优化生成过程。具体步骤包括:
- 资源分配: 在单数据中心或跨数据中心场景下,使用算法枚举配置,优化 GPU 分配,确保生成和训练时间平衡。
- 流水线优化: 通过流生成,允许训练阶段在生成过程中处理样本,实现异步执行。
- 偏斜处理: 训练一个小型排名模型预测输出长度,基于预测结果进行样本分派和调度,使用公式 Sample\text{ }Latency = PTL(BS) \times L \tag{1} 和 Latency = PTL(BS) \times L_{\text{aug}} \times \lceil \frac{M}{BS} \rceil \tag{2} 最小化生成延迟。
- 关键实现: 不修改原始 RL 算法,仅通过系统级优化实现高效执行,支持异构硬件和弹性调整。
Experiment
- 实验设置: 使用 Qwen2.5 模型(7B 到 72B),数据集为内部 CodeMath 提示,输出长度设置为 5K、10K 和 20K 以模拟训练不同阶段。比较基线包括 verl(共置架构的开源框架)和 ColocationRL(内部共置框架),评估指标为样本吞吐量。实验覆盖单数据中心和跨数据中心场景,进行消融研究验证各组件效果。
- 为什么设置合理: 实验设计全面,考虑了不同模型大小、序列长度和部署场景,确保结果的泛化性。使用真实数据集和硬件(NVIDIA H800 和 H20 GPU),模拟长尾效应,并通过标准化输出长度确保公平比较。
- 结果分析: StreamRL 同步版本比基线提高 1.06×–1.41× 吞吐量,异步版本进一步提高至 1.30×–2.66×。结果符合预期,证明了流生成和偏斜感知调度的有效性:例如,输出长度排名模型对长尾样本的召回率高达 87%,资源分配算法显著平衡阶段延迟。跨数据中心实验显示成本效益提高 1.23×–1.31×,通信开销小(不足 2%)。整体实验验证了方法的鲁棒性和实际部署潜力。
Further Thoughts
这项工作突出了分离式架构在 RL 训练中的潜力,未来可探索与其他领域结合,如与联邦学习集成以处理分布式数据源,或与弹性并行策略(如 Alpa 或 MegaScale)协同优化大规模模型训练。同时,输出长度排名模型的预测机制可能扩展到其他 AI 任务中,例如优化推理服务或预测任务难度,以提升资源利用率;此外,结合最近的 LLM 推理优化(如 vLLM 或 DistServe),可能进一步减少生成阶段的延迟,拓宽 RL 训练的应用场景。