DistFlow: A Fully Distributed RL Framework for Scalable and Efficient LLM Post-Training
文章标题:DistFlow:一个用于可扩展和高效的LLM后训练的完全分布式强化学习框架
作者/机构:
Zhixin Wang (上海创新研究院, 浙江大学),
Tianyi Zhou (上海创新研究院, 复旦大学),
Liming Liu (上海创新研究院),
Ao Li (上海创新研究院),
Jiarui Hu (上海创新研究院),
Dian Yang (上海创新研究院),
Jinlong Hou (上海创新研究院),
Siyuan Feng (上海创新研究院),
Yuan Cheng (上海创新研究院, 复旦大学, 上海人工智能科学研究院),
Yuan Qi (上海创新研究院, 复旦大学, 上海人工智能科学研究院)
A1 主要贡献
本文旨在解决现有大规模强化学习(RL)框架中的核心性能瓶颈问题。主流框架通常采用混合控制器架构,其中单控制器负责分派整体执行逻辑和管理数据传输,而多控制器执行分布式计算。这种架构在扩展时,微小的负载不均衡也会导致显著瓶颈,最终限制了系统的可扩展性。特别是,中心化的数据流控制器成为扩展性和效率的关键制约因素。
为了克服这些限制,本文提出了DistFlow,一个新颖的、完全分布式的RL训练框架,旨在打破扩展障碍。该框架的核心创新点和贡献如下:
-
完全分布式架构:DistFlow采用多控制器范式,将数据传输和执行任务分派给所有工作节点(worker),从而消除了中心化节点。这使得每个工作节点都能独立运行,移除了单点瓶颈,从而在扩展至数千个GPU时能实现近乎线性的可扩展性和显著的效率提升。
-
解耦的模块化流水线:DistFlow通过用户输入的有向无环图(DAG)来定义算法流水线。这种设计将算法逻辑与物理资源管理完全解耦。研究人员只需在DAG中定义整个RL工作流,专注于算法设计,框架会自动将此逻辑图映射到底层硬件,从而最大化资源利用率,并使研究人员能够快速、低成本地开发和验证新算法。
-
卓越的性能与可扩展性:通过广泛的实验评估,DistFlow在从单节点到千卡规模的各种集群配置中展现出卓越的性能和线性可扩展性。与当前最先进的同步框架相比,DistFlow在不同场景下实现了高达7倍的端到端训练吞吐量提升。
-
瓶颈分析与解决方案:本文分析了现有RL框架的核心性能瓶颈,指出了中心化数据流控制器是限制可扩展性和效率的关键。基于此分析,设计并实现了DistFlow,其多控制器范式和去中心化数据管理有效解决了这一问题。
流行的RL算法(如PPO和GRPO)的工作流可以被建模为一个有向无环图(DAG),如下图所示。图中的节点代表计算操作(如Actor生成或Actor训练),边代表决定执行顺序的数据依赖关系。
A3 背景知识与动机
2.1 针对LLM的强化学习
模型与工作流。人工智能的最新进展表明,强化学习(RL)提供了一个强大的框架,可以增强语言模型超越其预训练能力【1, Training language models to follow instructions with human feedback, 2022, Advances in neural information processing systems】【2, Deep reinforcement learning from human preferences, 2017, Advances in neural information processing systems】,使其更好地与人类偏好对齐并解决日益复杂的任务。在通过大规模RL训练的模型的单步优化迭代中,四个核心模型协同工作【12, Training a helpful and harmless assistant with reinforcement learning from human feedback, 2022, arXiv preprint arXiv:2204.05862】【13, Llama 2: Open foundation and finetuned chat models, 2023, arXiv preprint arXiv:2307.09288】:作为生成响应的目标模型的Actor模型;负责为特定任务提供基于偏好的奖励的奖励模型【14, Learning to summarize with human feedback, 2020, Advances in neural information processing systems】;作为衡量知识准确性基准的参考模型;以及评估价值的Critic模型。训练过程遵循一个精心设计的三步迭代工作流:生成(Generation)、评估(Evaluation)和训练(Training)。在生成阶段,Actor模型接收一批提示作为输入,并自回归地为每个提示生成相应的文本响应。响应生成完成后,进入评估阶段,奖励模型、参考模型和Critic模型分别为当前响应提供分数。最后的训练阶段是RL训练过程的核心,其中这三个分数被整合以计算优势函数,并通过反向传播仅更新Actor和Critic的参数,而奖励和参考模型在整个过程中保持冻结【15, Reinforcement learning optimization for large-scale learning: An efficient and user-friendly scaling library, 2025, arXiv preprint arXiv:2506.06122】。
算法。策略优化是强化学习中的关键机制,它通过系统地根据反馈信号改进策略来更新模型的决策策略,以最大化预期奖励,如图1所示。PPO通过实现一个裁剪的目标函数来增强训练过程的稳定性,该函数在保持样本效率的同时防止破坏性的策略更新【16, Scaling laws for reward model overoptimization, 2023, International Conference on Machine Learning】。GRPO通过消除对独立Critic模型的需要,并从群体分数中估计基线来降低计算成本,从而使大规模语言模型的强化学习更加高效【3, Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning, 2025, arXiv preprint arXiv:2501.12948】。
范式转变。RL不仅仅是预测下一个词元,而是将语言模型转变为能够根据奖励函数定义的目标进行优化的智能体【17, Scalable agent alignment via reward modeling: a research direction, 2018, arXiv preprint arXiv:1811.07871】【18, Improving alignment of dialogue agents via targeted human judgements, 2022, arXiv preprint arXiv:2209.14375】【19, Big bird: Transformers for longer sequences, 2020, Advances in neural information processing systems】。通过这种方法,RL框架使模型能够通过交互和反馈进行改进,而不仅仅局限于其初始训练数据【20, Finetuning language models from human preferences, 2019, arXiv preprint arXiv:1909.08593】【21, Webgpt: Browser-assisted question-answering with human feedback, 2021, arXiv preprint arXiv:2112.09332】。以前,人们普遍认为预训练的扩展能力决定了模型训练和推理的上限,但最近的研究挑战了这一观点,特别是Grok 4的性能突破【7, Grok 4, 2025, x.ai/news/grok-4】【22, Exploring the limits of transfer learning with a unified text-to-text transformer, 2020, Journal of machine learning research】【23, Electra: Pre-training text encoders as discriminators rather than generators, 2020, arXiv preprint arXiv:2003.10555】【24, Language models are unsupervised multitask learners, 2019, OpenAI blog】。事实上,尽管预训练的改进显示出边际效益递减,但RL为模型能力的显著提升提供了一条新途径【25, Multiagent communication meets natural language: Synergies between functional and structural language learning, 2020, arXiv preprint arXiv:2005.07064】【26, Retrieval-augmented generation for knowledgeintensive nlp tasks, 2020, Advances in neural information processing systems】【27, Composable deep reinforcement learning for robotic manipulation, 2018, 2018 IEEE international conference on robotics and automation (ICRA)】。
2.2 RL中的可扩展性挑战
RL中的扩展法则。扩展在先进的人工智能系统中带来了模型的涌现能力。此外,我们可以观察到,在扩大计算资源后,Grok 4在Humanity's Last Exam任务中表现出涌现能力,为未来研究指明了一条有希望的道路。计算资源的扩展显著影响模型处理复杂任务的能力,同时对高效的并行化策略提出了更高的要求,以充分利用这些资源【28, Seed rl: Scalable and efficient deep-rl with accelerated central inference, 2019, arXiv preprint arXiv:1910.06591】。
并行策略。常见的并行策略包括数据并行(DP)、张量并行(TP)、流水线并行(PP)等。DP通常使用像ZeRO【29, Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters, 2020, Proceedings of the 26th ACM SIGKDD international conference on knowledge discovery & data mining】这样的参数分片方法,在多个设备上并行化模型训练,并通过在不同设备上处理不同数据批次并同步梯度来提供样本级并行。TP,如Megatron-LM【30, Megatron-lm: Training multibillion parameter language models using model parallelism, 2019, arXiv preprint arXiv:1909.08053】中所实现的,分解矩阵乘法并将其分布在多个GPU上,提供操作级并行,有效解决了大型Transformer模型中的内存限制。PP则简单地将神经网络划分为跨设备的顺序阶段,从而能够有效地通过不同模型层对微批次进行流水线处理,同时减少峰值内存需求。
RL扩展的瓶颈。虽然扩展计算资源和数据通常能提高模型性能,但大规模RL训练面临着独特的瓶颈,阻碍了线性扩展。与监督学习中更多数据几乎总是有益不同,RL系统实际上可能会因批次大小增加而性能下降,因为连续更新之间的相关性变弱,使得学习变得不稳定【28, Seed rl: Scalable and efficient deep-rl with accelerated central inference, 2019, arXiv preprint arXiv:1910.06591】。简单地增加更多计算资源并不能保证RL系统中性能的成比例提升。随着RL训练的进行,模型倾向于越来越过度优化奖励模型,导致收益递减【16, Scaling laws for reward model overoptimization, 2023, International Conference on Machine Learning】。这种过度优化为性能增益设置了一个基本的天花板,无论增加多少计算能力。例如,生成阶段的数据偏斜问题导致严重的GPU利用率不足,因为只有少数长尾样本会使昂贵的计算资源处于等待状态【31, Optimizing rlhf training for large language models with stage fusion, 2024, arXiv preprint arXiv:2409.13221】。
2.3 RL的分布式系统架构
控制器范式。分布式机器学习系统,特别是为强化学习设计的系统,采用不同的控制器范式来管理跨硬件资源的计算。这些范式从根本上塑造了算法在实践中如何实现、扩展和优化【28, Seed rl: Scalable and efficient deep-rl with accelerated central inference, 2019, arXiv preprint arXiv:1910.06591】。
单控制器(Single-Controller)。单控制器范式采用一个中心化控制器来管理分布式程序的整体执行流程。通过中心化的控制逻辑,用户可以将数据流的核心功能构建为单个进程,而控制器自动生成分布式工作节点来执行计算。这种方法提供了硬件和数据流图的全局视图,允许对数据流任务进行灵活和优化的资源映射和执行顺序协调。然而,协调消息从控制器传递给所有工作节点,在大型集群上执行庞大的数据流图时会产生显著的调度开销【28, Seed rl: Scalable and efficient deep-rl with accelerated central inference, 2019, arXiv preprint arXiv:1910.06591】。
多控制器(Multi-Controller)。相比之下,多控制器范式通过为每个设备(或工作节点)提供自己的控制器来分配控制逻辑。这种方法在RL框架中很常见【28, Seed rl: Scalable and efficient deep-rl with accelerated central inference, 2019, arXiv preprint arXiv:1910.06591】,类似于最近的RL训练系统【15, Reinforcement learning optimization for large-scale learning: An efficient and user-friendly scaling library, 2025, arXiv preprint arXiv:2506.06122】,其中多个长期运行的分布式程序通过硬编码的数据同步来协调执行顺序。虽然这减少了中央协调瓶颈,但通常会导致复杂的实现和维护挑战,尤其是在扩展到大型模型架构或集群规模时。
2.4 现有RL系统的局限性
高要求。为了解决早期设计的效率问题,研究人员采用了依赖单控制器进行高级编排的架构,一些框架如verl【11, Hybridflow: A flexible and efficient rlhf framework, 2025, Proceedings of the Twentieth European Conference on Computer Systems (EuroSys ’25)】引入了混合范式,将其与多控制器的计算效率相结合。虽然这种方法在有限规模的场景中可以提高系统吞吐量,但其依赖单控制器管理整个数据流引入了一个根本性的关键限制:无法扩展。在当前AI研究时代,在数千个GPU上训练前沿模型已成为核心需求,高效扩展的能力至关重要。基于单控制器数据流的系统,无论是完全中心化还是混合式,都天生无法满足这一需求。这种架构选择,即由一个中心化节点协调所有与数据相关的任务,随着计算资源的增加而成为一个严重的瓶颈。这种设计在大规模下会被海量数据淹没,导致不稳定和故障,从根本上限制了研究和开发。这个关键的可扩展性问题,以及不灵活的流水线设计,表明这类架构不适合现代大规模AI研究的需求。
单控制器瓶颈。如图2所示,单控制器节点的瓶颈尤为严重。这种设计采用单控制器范式,由一个中心化控制器协调整个数据流。其职责涵盖所有与数据相关的任务,从最初的数据集加载到跨不同计算阶段的大规模中间数据。这种中心化设计迫使所有中间数据流经控制器节点,引入了昂贵的“一对多”和“多对一”通信操作。这些操作带来了巨大的I/O和通信成本,严重影响系统效率,特别是对于多模态或长文本生成等数据密集型任务。这种架构不仅效率低下,而且从根本上限制了系统的可扩展性。在大型分布式训练中,数据峰值量很容易超出控制器节点的内存,导致内存不足(OOM)错误,从而限制了系统的峰值数据处理能力。因此,单控制器成为双重瓶颈,既限制了整个系统的效率,也限制了其可扩展性。
僵化的算法流水线。此外,僵化的算法流水线是另一个限制。RL系统中的数据流和控制流本身就很复杂,而这类框架中的计算工作流被设计成一个高度集成、固定的逻辑,缺乏足够的灵活性。这种预定义的架构设计迫使用户直接修改源代码来进行任何流水线调整。这种方法不仅带来了工程挑战,还延长了创新实验的迭代周期,严重限制了框架的科学探索潜力。
2.5 设计考量
这些在可扩展性和灵活性方面的根本限制表明,依赖中心化数据流控制器的现有框架不适合日益增长的大规模AI研究需求。相比之下,多控制器架构天生就非常适合应对这一挑战,因为它能够实现一个完全分布式的系统,其中数据和计算都可以在没有中央瓶颈的情况下进行管理。因此,这项工作的核心动机是设计一个基于这一原则的新型框架。我们提出了一个完全分布式的架构,旨在解决传统系统的局限性。其核心特性包括一个去中心化、多控制器的架构,消除了中央节点,从而提供高吞吐量和线性可扩展性。此外,DistFlow利用用户输入的DAG定义的模块化流水线,将算法逻辑与物理资源管理解耦。
A2 方法细节
3 DistFlow 概述
因此,我们引入了DistFlow,一个专为大规模集群可扩展性设计的完全分布式RL框架。如图3所示,DistFlow采用多控制器范式,将所有计算和数据流均匀地分派到每个GPU上。DistFlow由三个主要组件构成:一个DAG规划器(DAG Planner)(§4),多个DAG工作节点(DAG Workers)(§5),以及一个数据协调器(Data Coordinator)(§6)。DAG规划器接收用户自定义的DAG文件,并将整个DAG分割成若干个独立的DAG任务,数量与DAG工作节点相等,每个任务不包含并行节点。DAG工作节点是核心的逻辑执行者,每个工作节点绑定到一个GPU上;它从DAG规划器接收并执行单个DAG任务。每个DAG工作节点都有一个分布式数据加载器(Distributed Dataloader),在任务执行前加载相应的数据。数据协调器管理所有DAG工作节点之间的数据协调。在每个阶段结束时,它将数据收集到一系列分布式数据缓冲区(Distributed Databuffers)中,然后在下一阶段开始时分派这些数据。
DistFlow 的工作流程。在训练开始前,用户需要提供三个基本配置:模型配置(Model Config)、训练配置(Training Config)和算法配置(Algorithm Config)。当使用自定义执行流水线时,还必须指定一个DAG配置(DAG Config);然而,在使用GRPO或PPO等内置算法时,此配置是可选的。模型配置指定了数据流中每个模型的架构和相应的并行策略。训练配置和算法配置分别提供必要的训练设置和超参数。
任务分派与执行。收到这些配置后,DAG规划器处理全局DAG,并将其分解为最小的可执行单元,称为DAG任务。DAG任务可以在不同的DAG工作节点之间复制。这些DAG任务随后被分派给一组DAG工作节点,其中每个DAG工作节点绑定到一个GPU,构成基本的计算单元。每个DAG工作节点通过加载其分配的模型和数据进行初始化,然后根据配置的并行策略设置所需的训练和推理引擎。初始化后,所有DAG工作节点开始并行执行其分配的DAG任务。
数据协调。在执行过程中,由于不同阶段之间并行策略的变化,数据需要在DAG工作节点之间重新分配。这由数据协调器管理,它能自动感知每个阶段的并行配置。它通过一系列分布式数据缓冲区来协调数据的收集和分派,确保整个系统的数据流无缝衔接。系统为每个节点分配一个数据缓冲区。
实现细节。我们的系统基于PyTorch【32, Pytorch: An imperative style, highperformance deep learning library, 2019, Advances in neural information processing systems】实现。对于GPU和CPU资源的管理,我们使用Ray【33, Ray: A distributed framework for emerging AI applications, 2018, 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18)】,这是一个用于轻松构建和扩展机器学习及Python应用程序的开源框架。我们的系统架构集成了用于不同操作的专用引擎。我们使用PyTorch的完全分片数据并行(FSDP)【32, Pytorch: An imperative style, highperformance deep learning library, 2019, Advances in neural information processing systems】作为训练引擎来执行推理和训练操作。对于生成操作,我们利用vLLM【34, Efficient memory management for large language model serving with pagedattention, 2023, Proceedings of the 29th Symposium on Operating Systems Principles】和SGLang【35, Sglang: Efficient execution of structured language model programs, 2024, Advances in neural information processing systems】等推理引擎,这些引擎专为高效的自回归生成而设计。为了管理这些组件,并受到verl的分层API设计的启发,我们的系统使用3DParallelWorker
基类。
4 DAG 规划器
DAG定义的执行模型。为了解决实现多样化RL工作流的复杂性,我们的框架以一个由DAG定义的执行模型为中心。核心设计原则是将算法工作流的逻辑表示与其物理计算资源解耦。这种关注点分离通过两个主要组件实现:一个供用户使用的声明式DAG接口和一个将逻辑图转换为可执行任务链的后端DAG规划器。
4.1 输入DAG定义
通过配置文件定义工作流。该框架允许用户通过一个配置文件定义完整的RL工作流。该文件指定了一个DAG,其中每个节点代表一个原始计算步骤。节点抽象由四个关键属性定义:用于识别的唯一节点ID;指定其功能目的的角色(例如,ACTOR, CRITIC, REWARD, REFERENCE);阐明计算性质的类型(例如,MODEL_INFERENCE, MODEL_TRAIN, COMPUTE);以及建立节点间执行顺序和数据流的依赖关系。通过使用这种高级抽象,用户可以专注于算法逻辑,而不是分布式调度的底层复杂性。
4.2 DAG 分解
逻辑图到线性化执行流水线的转换。执行用户定义的DAG的一个主要挑战是确保其安全高效地适应资源受限的、多个大型模型共享同一资源池的同地部署(colocated)架构。我们的框架通过DAG规划器来应对这一挑战。其根本职责是将逻辑图转换为一个具体的、线性的执行流水线,以避免资源争用和潜在的OOM错误。为实现这一点,规划器通过分析每个节点的逻辑深度来自动序列化工作流。如果多个节点存在于同一深度,这意味着并行执行,规划器会有系统地引入依赖关系以强制执行顺序。例如,如图4所示,如果一个输入DAG包含两个并行节点,Inference I和Inference II,规划器会通过使一个节点成为另一个节点的前置条件来转换该图。
解耦与自动化。这种设计将高层算法逻辑与底层执行计划解耦,允许研究人员定义复杂的工作流而无需手动管理序列化。这个自动化过程保证了在任何给定时间只有一个节点处于活动状态。
5 DAG 工作节点
核心执行单元。在定义了逻辑工作流之后,中心挑战是将其转换为一个具体且可扩展的执行模型。我们的系统通过引入DAG工作节点来解决这个问题,这是一个核心组件,旨在单个GPU上执行序列化的任务链,同时为算法实验提供最大的灵活性。
设计抽象。DAG工作节点是框架的基本执行单元。其设计由两个关键抽象支配:一个结构化的生命周期和一个动态函数分派机制。生命周期由一个初始化阶段(工作节点准备其计算环境)和一个迭代的执行阶段(处理任务链)组成。动态函数分派机制将节点的逻辑定义(角色和类型)与其实现解耦,从而实现模块化和可插拔的架构。
具体操作流程。这种抽象设计通过一个具体的操作流程得以实现。在初始化阶段,工作节点根据抽象DAG实例化其环境。它使用分布式数据加载器获取数据,加载指定的模型,并初始化后端引擎,如vLLM、SGLang、PyTorch FSDP或Megatron。然后,它将任务链具体化为一个执行队列,并将适当的函数绑定到每个节点,如图5所示。
执行与扩展性。在随后的执行阶段,工作节点为每个RL迭代进入一个循环。它请求一个数据批次,并按顺序执行链中的每个节点。一个数据缓冲区(databuffer)充当中间状态管理器,为每个函数提供必要的输入并存储其输出。完成链后,通过高效的多对一通信将指标聚合到全局rank 0。该模型的主要好处是其固有的可扩展性。工作流结构(在DAG中定义)与其操作逻辑(由函数映射解析)的解耦,使研究人员能够快速创新。例如,引入一种新的奖励计算方法或不同的策略损失函数,无需改变核心数据流。研究人员只需在自定义函数中实现新逻辑,并将其映射到DAG中的一个节点,即可无缝地将其集成到执行流水线中,而无需修改周围的框架。
6 数据协调器
数据管理挑战。在大型分布式RL中,数据管理面临双重挑战。首先,初始数据加载的中心化方法从根本上与可扩展的分布式架构相冲突。强制单个节点加载然后分发海量数据集会产生固有瓶颈,限制了可扩展性和效率。其次,RL工作流的动态性,其中像生成和训练这样的计算阶段可能采用不同的并行策略,需要一种高效且正确的机制来在工作节点之间重新分配中间数据,以防止系统停顿或无声的训练错误。
统一的数据协调器。为了应对这些挑战,我们的框架引入了一个统一的数据协调器。该协调器是整个数据生命周期的高级抽象,由两个专门的分布式组件组成:分布式数据加载器(Distributed Dataloader)和分布式数据缓冲区(Distributed Databuffer)。数据加载器负责初始数据集的静态、一次性加载,确保数据在源头被正确分区。相比之下,数据缓冲区管理计算阶段之间中间数据的动态、瞬时流动,确保数据的正确循环。这两个组件共同保证了从初始数据加载到整个数据流在DAG工作节点之间的负载均衡。
分布式数据加载器工作原理。在初始化阶段,分布式数据加载器获取其关联的DAG工作节点的并行策略信息。它根据DP大小将数据集划分为相应的区域,并根据DP组的rank加载适当的数据。图6以一个包含512个样本的数据集为例,说明了分布式数据加载器的执行逻辑,其中rollout阶段采用DP为2、TP为2的并行策略。如图6所示,分布式数据加载器根据DP大小将数据集划分为两个区域:样本0-255设为第一个区域,样本256-511设为第二个区域。在加载期间,DP组rank为0的两个DAG工作节点将从第一个区域加载数据,而DP组rank为1的两个工作节点将从第二个区域加载。数据加载是并行执行的,这带来了加载效率的提升。
分布式数据缓冲区工作原理。分布式数据缓冲区在具有不同数据并行(DP)大小的阶段之间重新划分数据。每个源数据缓冲区对其本地数据进行切片,通过all-to-all模式将切片分发给所有其他数据缓冲区,然后每个目标数据缓冲区连接接收到的数据,为下一阶段形成一个大小正确的批次。
6.1 分布式数据加载器
去中心化数据加载。在大型场景中,一个节点加载整个数据集的中心化方法从根本上是低效和不可扩展的。因此,为了保持架构一致性和性能,我们的框架实现了一个分布式数据加载器。分布式数据加载器的数量等于DAG工作节点的数量,即GPU的数量。每个数据加载器只加载其对应的DAG工作节点在rollout阶段所需的数据,避免任何冗余数据。这种方法天生避免了单节点内存瓶颈,并通过并行化实现了更高的数据加载效率。
6.2 分布式数据缓冲区
数据流的核心组件。分布式数据缓冲区是数据流的核心组件,负责RL阶段之间的数据重新分配。每个节点分配一个实例,并由本地工作节点共享。其主要功能是充当一个感知并行的中介,确保在阶段转换期间数据流的正确性和效率,此时连续阶段的DP大小可能不同。
操作逻辑。操作逻辑在一个计算阶段结束后开始。为避免来自多个模型副本的数据冗余,只有TP rank为0的DAG工作节点将其生成的数据放入其本地数据缓冲区。然后,数据缓冲区会预测下一阶段的需求,并根据DP大小是否改变处理两种不同的情况:
根据并行需求调整逻辑。数据缓冲区的操作逻辑根据后续阶段的并行需求进行调整。当DP大小保持不变时,数据缓冲区执行一个快速路径操作,将接收到的数据直接放入节点的共享内存中。这使得下一阶段的工作节点可以直接检索数据进行计算,开销最小,效率更高。相反,如果DP大小不同,数据缓冲区会启动一个更复杂的多步重新分配过程。它首先将其本地持有的数据切片成等于集群中数据缓冲区数量的分区,并通过all-to-all通信模式分发它们,如图7所示。之后,每个数据缓冲区连接从其对等节点接收到的多个分区,并存储新整合的数据块。最后,当下一阶段的DAG工作节点请求数据时,数据缓冲区根据其新的DP rank将适当的切片分发给其节点内的工作节点。
示例。例如,如图8所示,假设一个生成阶段(DP=2, TP=2)转换到一个训练阶段(DP=4),全局批次大小为64。在生成阶段,TP rank为0的工作节点会将其32个样本放入相应的本地数据缓冲区。数据缓冲区知道后续阶段的DP大小为4,然后会执行重新分配过程,并随后为四个DAG工作节点各分配16个样本进行训练。这种自动化处理确保了任何阶段转换的正确数据流和负载均衡。
A4 实验
实验环境
- 硬件配置 (Testbed): 实验在一个包含128个节点的集群上进行。每个节点配备8块NVIDIA Hopper GPU,通过NVLink互连。节点之间通过支持RoCE v2的RDMA网络连接。
- 软件配置: PyTorch 2.6.0, CUDA 12.6, vLLM 0.8.5.post1, NCCL 2.21.5。
- 模型与算法:
- 算法: PPO 和 GRPO。对于PPO实验,使用函数奖励代替奖励模型,Critic模型大小与Actor模型相同。
- 模型: 语言模型使用Qwen-2.5-Instruct系列,视觉语言模型(VLM)使用Qwen-2.5-VL-Instruct系列,模型规模包括7B、32B和72B。
- 数据集:
- 语言模型: DeepScaleRPreview-Dataset【36, Deepscaler: Surpassing o1-preview with a 1.5b model by scaling rl, 2025, Notion Blog】,包含约4万个独特的数学问题。
- 视觉语言模型: MMEureka-Dataset【37, Mm-eureka: Exploring the frontiers of multimodal reasoning with rule-based reinforcement learning, 2025, arXiv preprint arXiv:2503.07365】。
- 配置: 所有实验默认最大提示长度为2048,最大响应长度为4096,对较短的响应进行填充。
- 基线系统 (Baseline): 与verl【11, Hybridflow: A flexible and efficient rlhf framework, 2025, Proceedings of the Twentieth European Conference on Computer Systems (EuroSys ’25)】 v0.4.0进行比较,这是一个最先进的RL训练系统。DistFlow和verl都使用vLLM作为推理引擎,PyTorch FSDP作为训练后端。
- 评估指标: 主要性能指标是吞吐量(tokens/second),通过全局批次中的总词元数除以一次迭代的时间计算得出。结果是经过预热期后多次迭代的平均值。
实验结果
端到端评估
DistFlow在所有测试配置中始终优于基线verl(图9和图10)。这主要归功于其完全分布式的数据流架构,该架构消除了基线系统中的单控制器瓶颈。
* PPO算法: DistFlow实现了1.09倍至1.64倍的加速比。随着GPU数量增加,加速效果更明显。值得注意的是,基线在32个GPU上训练72B模型时出现OOM错误,而DistFlow成功完成(图9)。
* GRPO算法: GRPO涉及更大数据量,DistFlow的优势更加突出,速度比基线快高达2.62倍(图10)。这进一步证明了DistFlow分布式数据系统在数据密集型场景下的高效性。
可扩展性评估
DistFlow在从32个GPU扩展到1024个GPU时表现出近乎线性的扩展能力(图11),这得益于其完全分布式设计,能有效平衡计算和数据通信负载。
* 线性扩展: 对于32B模型,从64个GPU扩展到512个GPU时,平均每个GPU的吞吐量保持了初始性能的80.5%。
* 与基线对比: 基线系统由于OOM错误无法完成相同的线性度测试。为了进行直接比较,实验在每个集群规模下找到基线能支持的最大全局批次大小(表1),并在此约束下进行测试。结果显示,DistFlow的速度比基线快高达7倍,尤其是在VLM设置和大规模集群上(图12)。
长上下文评估
在处理长上下文这一数据密集型任务时,DistFlow的分布式数据流优势更加明显。
* 性能趋势: 随着上下文长度增加,DistFlow的性能优势也随之增长。对于7B模型,其吞吐量加速比从8k上下文长度的1.48倍增长到64k上下文长度的2.03倍(图13)。
* 鲁棒性: 基线系统在处理32k上下文长度的72B模型时遭遇OOM错误,而DistFlow能够顺利完成,展示了其架构在处理极端数据密集型任务时的鲁棒性和可扩展性。
收敛性测试
为了验证性能提升是否影响模型准确性,实验对比了DistFlow和verl的收敛情况。
* 结果: 在使用相同超参数训练32B模型20个周期后,DistFlow与基线的奖励和熵曲线轨迹完全相同(图14),表明训练准确性没有受到影响。
* 效率: 在达到相同结果的同时,DistFlow将总执行时间减少了21%,证明其效率和可扩展性的提升并未以牺牲训练准确性为代价。
A5 结论
本文介绍了DistFlow,一个旨在解决大规模强化学习(RL)训练中固有的可扩展性和灵活性挑战的新型框架。为了应对现有方法中常见的单控制器数据流瓶颈,我们提出了一个完全分布式的架构。DistFlow的核心采用了多控制器范式,将数据加载、计算和中间数据传输等任务均匀地分派到所有工作节点,从而彻底消除了中心瓶颈,实现了近乎线性的可扩展性。此外,为了赋予研究人员更大的灵活性,DistFlow引入了由用户定义的DAG驱动的模块化流水线。这种设计将算法逻辑与物理资源管理解耦,极大地加速了新算法的实验和迭代周期。实验证明,与最先进的框架相比,DistFlow实现了高达7倍的端到端吞吐量提升。我们相信,这项工作通过提供一个更高效、更灵活、真正可扩展的解决方案,为大规模RL研究铺平了道路。
💬 评论讨论
欢迎在这里分享您的想法和见解!