RailS: Load Balancing for All-to-All Communication in Distributed Mixture-of-Experts Training

TL;DR

RailS is a novel load-balancing framework for distributed Mixture-of-Experts training that exploits Rail topology symmetry and local scheduling to significantly reduce communication bottlenecks and improve training efficiency.

Contribution

RailS introduces a topology-aware, local scheduling load-balancing method that leverages Rail architecture symmetry, achieving near-optimal load balance and substantial performance gains.

Findings

Improves bus bandwidth by 20%–78%.

Reduces iteration time by 17%–78%.

Achieves near-optimal load balance in MoE workloads.

Abstract

Training Mixture-of-Experts (MoE) models introduces sparse and highly imbalanced all-to-all communication that dominates iteration time. Conventional load-balancing methods fail to exploit the deterministic topology of Rail architectures, leaving multi-NIC bandwidth underutilized. We present RailS, a distributed load-balancing framework that minimizes all-to-all completion time in MoE training. RailS leverages the Rail topology's symmetry to prove that uniform sending ensures uniform receiving, transforming global coordination into local scheduling. Each node independently executes a Longest Processing Time First (LPT) spraying scheduler to proactively balance traffic using local information. RailS activates N parallel rails for fine-grained, topology-aware multipath transmission. Across synthetic and real-world MoE workloads, RailS improves bus bandwidth by 20%--78% and reduces completion time by 17%--78%. For Mixtral workloads, it shortens iteration time by 18%--40% and achieves near-optimal load balance, fully exploiting architectural parallelism in distributed training.