TLDR: MindSpeed RL is a new system designed for scalable and efficient reinforcement learning (RL) training on Ascend NPU clusters. It tackles key bottlenecks in RL dataflow, specifically dispatch overhead in sample flow using a distributed transfer dock strategy, and redundant memory usage in resharding flow via an allgather-swap technique. By integrating these innovations with various parallelization and acceleration methods, MindSpeed RL achieves significantly higher training throughput (1.42 to 3.97 times) for large language models compared to existing systems, demonstrating robust performance and scalability.
Reinforcement learning (RL) is a powerful method increasingly used to enhance the capabilities of large language models (LLMs), helping them achieve impressive results in complex tasks like mathematics and coding. However, training these RL systems, especially at a large scale, presents significant challenges. Mainstream RL algorithms often involve multiple components, such as actor and reference workers, and frequently switch between training and generation states. This complexity leads to heavy dependencies and data movement, often resulting in poor cluster scalability and inefficient memory use.
Introducing MindSpeed RL
A new system called MindSpeed RL has been developed to tackle these challenges, offering an effective and efficient solution for large-scale RL training. Unlike traditional centralized approaches, MindSpeed RL adopts a distributed perspective to manage the crucial data dependencies within RL training, specifically focusing on what are called ‘sample flow’ and ‘resharding flow’.
Addressing Dataflow Bottlenecks
The researchers behind MindSpeed RL identified two primary bottlenecks in existing RL systems: dispatch overhead in the sample flow and redundant memory usage in the resharding flow.
The **sample flow** describes how data, such as prompts and generated responses, moves between different stages and workers in the RL algorithm. Existing systems often use a centralized ‘replay buffer’ for this, which can become a bottleneck, causing significant delays, especially as the training scale increases. To overcome this, MindSpeed RL introduces a **distributed transfer dock strategy**. This strategy breaks down the centralized replay buffer into multiple ‘warehouses’ distributed across the cluster, each handling a portion of the data. Additionally, dedicated ‘controllers’ are assigned to each worker state, managing the metadata for data requests. This distributed setup significantly reduces the communication burden and dispatch overhead, allowing for more efficient data transfer.
The **resharding flow** involves transferring updated model weights from the training stage to the generation stage. This process often leads to redundant memory consumption, where memory is unnecessarily held even when not actively used. MindSpeed RL addresses this with an **allgather-swap technique**. This technique involves temporarily buffering weights for the allgather operation, then selectively copying only the necessary weight slices for generation. Crucially, it then swaps the original, no-longer-needed update weights from the device memory to the host memory. This frees up valuable device memory for other operations, like generating new data. Before the next update stage, the weights are efficiently swapped back from the host to the device, often overlapped with other processes to maintain efficiency.
Also Read:
- Virne: A New Benchmark for Deep Reinforcement Learning in Network Resource Allocation
- Optimizing Network Routing with AI: A Simulation-Driven Approach
Comprehensive Optimization and Performance
Beyond these innovative dataflow mechanisms, MindSpeed RL integrates a wide array of parallelization strategies and acceleration techniques for systematic optimization. These include various forms of parallelism like data parallelism, tensor parallelism, pipeline parallelism, expert parallelism (especially for Mixture-of-Experts models), and context parallelism for long sequences. It also incorporates optimized computational operations through ‘fused kernels’ such as FlashAttention, which combine multiple steps into a single, more efficient operation.
Comprehensive experiments were conducted on a super pod of Ascend NPUs, a powerful neural processing unit cluster. MindSpeed RL was tested with popular large language models like Qwen2.5-Dense-7B/32B, Qwen3-MoE-30B, and DeepSeek-R1-MoE-671B. The results showed that MindSpeed RL significantly boosts training throughput, achieving 1.42 to 3.97 times higher throughput compared to existing state-of-the-art systems. The system also demonstrated stable training processes and reliable performance on various benchmarks, even for extremely large models.
MindSpeed RL is an open-source framework, aiming to foster further development of RL algorithms on Ascend NPUs. For more technical details, you can refer to the original research paper.
