Unleashing Efficient Asynchronous RL Post-Training via Staleness-Constrained Rollout Coordination

TL;DR

StaleFlow is a system that improves asynchronous reinforcement learning post-training by jointly controlling data staleness and skewness, leading to higher throughput without sacrificing convergence.

Contribution

It introduces a unified approach with a global staleness protocol and redesigned architecture to address data staleness and skewness simultaneously in asynchronous RL.

Findings

Achieves 1.42-2.68× higher throughput than existing systems.

Maintains convergence while reducing data staleness and skewness.

Demonstrates effectiveness across various RL post-training scenarios.

Abstract

Reinforcement learning (RL) post-training has become pivotal for enhancing the capabilities of modern large models. A recent trend is to develop RL systems with a fully disaggregated architecture, which decouples the three RL phases (rollout, reward, and training) onto separate resources and executes them asynchronously. However, two critical data-level concerns arise: (1) asynchronous execution leads to data staleness in trajectories (the data generated by rollout) as the model parameters used in rollout may not be up to date, which impairs RL convergence; and (2) the length variation of trajectories introduces severe data skewness, leading to workload imbalance and degraded system performance. Existing systems fail to address these two concerns in a unified manner. Techniques that tightly control data staleness often constrain effective data skewness mitigation, while aggressive data skewness mitigation tends to exacerbate data staleness. As a result, systems are forced to trade off convergence for performance, or vice versa. To address this, we propose StaleFlow, an RL post-training system that jointly tackles data staleness and skewness. First, to control staleness, StaleFlow introduces a global consistency protocol that tracks the full lifecycle of each trajectory and constrains staleness. Second, to mitigate skewness, StaleFlow re-designs the RL system architecture by constructing data servers for trajectories and parameters to achieve flexible rollout coordination. Subsequently, we develop a suite of staleness-aware, throughput-oriented strategies to enhance system performance. Evaluations show that StaleFlow achieves up to 1.42-2.68$\times$ (1.17-2.01$\times$ on average) higher throughput than state-of-the-art systems, without compromising convergence.