RailWave improves communication speed for large expert-parallel AI models
RailWave: Adaptive Spatial and Temporal Scheduling for Expert-Parallel Communication
Distributed, Parallel, and Cluster Computing
Summary
Communication between parts of very large AI models can get stuck because some connections get overloaded. The authors introduce RailWave, a new system that smartly spreads out and times data traffic to prevent these slowdowns. It adapts to different phases of work to keep data flowing efficiently. Tests show RailWave speeds up communication by up to 5.8 times compared to standard methods.
What this means in practice
- •For machine learning engineers: Improve training speed of large-scale expert-parallel models by reducing communication bottlenecks during expert routing phases.
- •For data center network operators: Optimize network traffic scheduling for AI workloads to better utilize network resources and reduce peak congestion.
Authors
Chutian Wang, Wenhao He, Jingmin Zhu, Qingyu Yin, Heng Xu, Xiuyu Li
Abstract
Irregular All-to-All communication is a major bottleneck in expert-parallel Mixture-of-Experts (MoE) models. Even with fixed expert routing and placement, uneven utilization of parallel network Rails and incast can limit communication performance. We present RailWave, a phase-adaptive communication layer built on DeepEP that addresses these bottlenecks below the routing layer through spatial and temporal traffic shaping. RailBalance redistributes source traffic across eligible Rails using source-local information, while a reusable, topology-derived permutation schedule limits concurrent senders per receiver without rebuilding demand-dependent schedules for each communication phase. A lightweight calibrated selector chooses an execution path according to each phase's traffic characteristics and offline profiling results. On training-derived communication workloads from the 106B GLM-4.5-Air model, RailWave delivers up to 5.84x speedup on H800 and 4.36x on H20 over Native. Code is available at https://github.com/CyberSecurityErial/RailWave-EP.