DigenRL:基于扩散并行与训练器辅助加速视觉生成大模型解耦强化学习
速览
针对视觉生成大模型强化学习资源耦合问题,提出DigenRL解耦框架。该框架引入生成轴流水线、时间步并行及弹性训练器辅助生成技术,有效减少执行气泡。实验显示,其吞吐量较SOTA系统提升1.56-2.10倍。
AI 深度解读
加速视觉生成大模型的解耦强化学习:基于扩散并行与训练器辅助生成
背景
强化学习(Reinforcement Learning, RL)已成为大语言模型(LLM)后训练阶段的主导范式,推动了如 veRL 等高性能 RL 系统的出现。与此同时,面向扩散模型(Diffusion)的 RL 算法(例如 DanceGRPO 和 FlowGRPO)迅速将 RL 的应用范围从语言推理扩展到了基于扩散的视觉生成和基于流(Flow)的生成领域。
然而,针对扩散生成式大模型的高效 RL 系统仍处于探索阶段。现有的实现方案(如 veRL-Omni)仍然依赖共置执行(Colocated Execution),即采样(Rollout)和训练资源部署在同一组硬件上。这种架构虽然简化了同步过程,但也带来了显著的局限性:
- 资源耦合:采样和训练资源无法独立扩展,限制了异构部署的灵活性。
- 扩展性受限:难以根据采样和训练不同的计算需求进行独立缩放。
- 效率瓶颈:在解耦架构中,执行气泡(Execution Bubbles,即资源空闲等待的时间)难以消除,导致整体吞吐量受限。
核心内容
为了解决上述问题,研究团队提出了 DigenRL,这是一个专为基于扩散的生成式大模型设计的解耦强化学习框架(Disaggregated RL Framework)。DigenRL 支持灵活的资源分配,兼容异构 GPU,并促进高效的任務调度。
为了最大限度地减少解耦架构中的执行气泡,DigenRL 引入了三项核心技术创新:
1. 生成轴流水线(GAP)与时间步并行(TSP)
在扩散架构内部,研究团队设计了生成轴流水线(Generation-axis Pipeline, GAP)和时间步并行(Time-step Parallelism, TSP)。
- 机制:通过在扩散模型内部实现更细粒度的流水线处理,使得采样(Rollout)和训练之间的协作更加紧密。
- 作用:这种细粒度的流水线机制允许在采样和训练阶段之间进行更重叠的计算,从而减少等待时间,提高硬件利用率。
2. 弹性训练器辅助生成(TAG)
提出了一种**弹性训练器辅助生成(Elastic Trainer-Assisted Generation, TAG)**方法。
- 机制:允许训练阶段的 GPU 资源动态地协助执行采样生成任务。
- 作用:当训练资源出现空闲时,可以将其转化为采样资源,反之亦然。这种弹性机制有效利用了异构集群中的闲置算力,平衡了采样和训练之间的负载不均。
3. 严格的一步约束异步策略
引入了一种严格的一步约束异步策略(Tightly One-step Constrained Asynchronous Strategy)。
- 机制:进一步利用流水线中的尾部气泡(Tail Bubble)。
- 作用:通过精确控制异步更新的步长和约束条件,确保在异步执行过程中模型更新的一致性,同时最大化利用原本会被浪费的计算周期。
实验验证
研究团队在三个包含 16-32 张 GPU 的硬件测试平台上进行了广泛实验,使用的生成模型包括:
- HunyuanVideo-13B
- Wan2.1-14B
- FLUX.1-12B
- QwenImage-20B
实验结果表明,与最先进的扩散 RL 系统(veRL-Omni 和 GenRL)相比,DigenRL 实现了 1.56 到 2.10 倍的吞吐量提升。
关键要点
- 架构革新:DigenRL 打破了传统 RL 系统中采样与训练资源共置的限制,采用解耦架构,支持异构 GPU 和灵活的资源调度。
- 细粒度并行:通过生成轴流水线(GAP)和时间步并行(TSP),在扩散模型内部实现了更细粒度的流水线,减少了采样与训练之间的同步开销。
- 资源弹性利用:训练器辅助生成(TAG)技术允许训练 GPU 动态参与采样任务,有效消除了资源闲置,提高了集群整体利用率。
- 异步优化:严格的一步约束异步策略进一步优化了流水线尾部的气泡时间,提升了整体执行效率。
- 显著性能提升:在多个主流视觉生成大模型(如 HunyuanVideo、Wan2.1、FLUX.1、QwenImage)上,DigenRL 的吞吐量比现有 SOTA 系统高出 1.56-2.10 倍。
- 适用性广泛:该技术不仅适用于扩散模型,也为基于流的生成模型提供了高效的 RL 训练解决方案。
意义与影响
DigenRL 的提出标志着视觉生成大模型强化学习训练进入了一个新的阶段。其核心贡献在于解决了大规模分布式训练中资源异构性和负载不均衡的难题。
- 降低训练成本:通过解耦架构和弹性资源调度,企业可以更灵活地利用现有异构 GPU 集群,无需为采样和训练分别购买专用硬件,显著降低了高性能视觉生成模型训练的硬件门槛和成本。
- 加速模型迭代:吞吐量 1.5-2 倍的提升意味着模型收敛速度加快,研究人员可以在更短的时间内进行更多的实验和迭代,加速视觉生成模型的进化。
- 推动 RL 在视觉领域的普及:随着高效 RL 系统的出现,基于强化学习优化视觉生成模型的质量(如一致性、细节丰富度、指令遵循能力)将成为常态,推动 AIGC 内容生成质量的进一步提升。
- 为多模态大模型提供范式参考:DigenRL 的设计思路(细粒度并行、弹性辅助、异步优化)不仅适用于扩散模型,也为其他多模态生成模型(如视频生成、3D 生成)的强化学习训练提供了可借鉴的技术范式。
