MindSpeed RL: Distributed Dataflow for Scalable and Efficient RL Training on Ascend NPU Cluster
作者: Laingjun Feng, Chenyi Pan, Xinjie Guo, Fei Mei, Benzhe Ning, Jianxiang Zhang, Xinyang Liu, Beirong Zhou, Zeng Shu, Chang Liu, Guang Yang, Zhenyu Han, Jiangben Wang, Bo Wang
分类: cs.LG, cs.AI
发布日期: 2025-07-25
备注: 9 pages
💡 一句话要点
MindSpeed RL:昇腾NPU集群上可扩展高效强化学习训练的分布式数据流系统
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 强化学习 分布式训练 数据流管理 昇腾NPU 大规模语言模型 集群可扩展性 内存优化
📋 核心要点
- 现有强化学习训练系统在处理大规模数据和复杂模型时,面临着集群可扩展性差和内存利用率低的挑战,严重制约了训练效率。
- MindSpeed RL通过分布式传输坞站策略和allgather-swap策略,优化了样本流和重分片流,从而降低了调度开销和内存占用。
- 实验结果表明,MindSpeed RL在多个大型语言模型的强化学习训练中,相比现有系统,吞吐量提升了1.42到3.97倍,展现了其优越的性能。
📝 摘要(中文)
强化学习(RL)越来越多地被用于对齐大型语言模型。流行的RL算法利用多个worker,并且可以建模为一个图,其中每个节点是worker的状态,每条边表示节点之间的数据流。由于繁重的跨节点依赖关系,RL训练系统通常面临较差的集群可扩展性和低内存利用率的问题。本文介绍了MindSpeed RL,一个有效且高效的大规模RL训练系统。与现有的集中式方法不同,MindSpeed RL从分布式角度组织了RL训练中的基本数据依赖关系,即样本流和重分片流。一方面,设计了一种分布式传输坞站策略,该策略在传统的回放缓冲区的基础上设置控制器和仓库,以释放样本流中的调度开销。提出了一种实用的allgather-swap策略,以消除重分片流中的冗余内存使用。此外,MindSpeed RL进一步集成了许多并行化策略和加速技术,以进行系统优化。与现有的最先进系统相比,在Qwen2.5-Dense-7B/32B、Qwen3-MoE-30B和DeepSeek-R1-MoE-671B的RL训练上的综合实验表明,MindSpeed RL的吞吐量提高了1.42~3.97倍。最后,我们开源了MindSpeed RL,并在具有384个神经处理单元(NPU)的昇腾超级集群上进行了所有实验,以证明昇腾的强大性能和可靠性。
🔬 方法详解
问题定义:现有强化学习训练系统在扩展到大规模集群时,面临着严重的通信瓶颈和内存效率问题。具体来说,集中式的数据管理方式导致调度开销大,而频繁的数据重分片则造成了大量的冗余内存占用,限制了整体的训练吞吐量。
核心思路:MindSpeed RL的核心思路是从分布式角度出发,重新设计强化学习训练过程中的数据流管理。通过将数据调度和存储分散到各个节点,并优化数据重分片策略,从而降低通信开销和内存占用,提高集群的整体利用率。
技术框架:MindSpeed RL的整体架构包含以下几个主要模块:1) 分布式传输坞站:用于分散样本的调度和存储,降低中心节点的压力;2) allgather-swap策略:用于优化数据重分片过程,减少冗余内存的使用;3) 并行化策略和加速技术:用于进一步提升训练效率。整个流程包括样本收集、数据传输、模型训练和策略更新等环节。
关键创新:MindSpeed RL的关键创新在于其分布式数据流管理策略。传统的强化学习训练系统通常采用集中式的数据管理方式,而MindSpeed RL则将数据调度和存储分散到各个节点,从而降低了通信开销和内存占用。此外,allgather-swap策略也有效地减少了数据重分片过程中的冗余内存使用。
关键设计:分布式传输坞站策略通过在传统回放缓冲区的基础上设置控制器和仓库来实现。控制器负责样本的调度,而仓库负责样本的存储。allgather-swap策略则通过在allgather操作后进行数据交换,从而避免了冗余内存的分配。具体的参数设置和网络结构取决于具体的强化学习算法和模型。
🖼️ 关键图片
📊 实验亮点
实验结果表明,MindSpeed RL在Qwen2.5-Dense-7B/32B、Qwen3-MoE-30B和DeepSeek-R1-MoE-671B等大型语言模型的强化学习训练中,相比现有最先进系统,吞吐量提升了1.42到3.97倍。这些实验在拥有384个昇腾NPU的超级集群上进行,充分验证了MindSpeed RL的性能和可扩展性。
🎯 应用场景
MindSpeed RL可广泛应用于大规模语言模型的强化学习训练,例如指令微调、奖励模型训练等。其高效的分布式数据流管理能力,能够显著提升训练效率,加速模型的迭代和优化,从而推动人工智能技术的发展和应用。
📄 摘要(原文)
Reinforcement learning (RL) is a paradigm increasingly used to align large language models. Popular RL algorithms utilize multiple workers and can be modeled as a graph, where each node is the status of a worker and each edge represents dataflow between nodes. Owing to the heavy cross-node dependencies, the RL training system usually suffers from poor cluster scalability and low memory utilization. In this article, we introduce MindSpeed RL, an effective and efficient system for large-scale RL training. Unlike existing centralized methods, MindSpeed RL organizes the essential data dependencies in RL training, i.e., sample flow and resharding flow, from a distributed view. On the one hand, a distributed transfer dock strategy, which sets controllers and warehouses on the basis of the conventional replay buffer, is designed to release the dispatch overhead in the sample flow. A practical allgather--swap strategy is presented to eliminate redundant memory usage in resharding flow. In addition, MindSpeed RL further integrates numerous parallelization strategies and acceleration techniques for systematic optimization. Compared with existing state-of-the-art systems, comprehensive experiments on the RL training of popular Qwen2.5-Dense-7B/32B, Qwen3-MoE-30B, and DeepSeek-R1-MoE-671B show that MindSpeed RL increases the throughput by 1.42 ~ 3.97 times. Finally, we open--source MindSpeed RL and perform all the experiments on a super pod of Ascend with 384 neural processing units (NPUs) to demonstrate the powerful performance and reliability of Ascend.