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

TL;DR

RailS targets the bottlenecked all-to-all traffic in distributed MoE training by exploiting the Rail topology's deterministic, multi-NIC parallelism. It converts the global optimization into a distributed, local scheduling problem using Flow Splitting, LPT-based Scheduling, and Multipath Spraying, underpinned by a min-max load-balancing framework and a linear program that bounds the optimal time. The key theoretical insights show that Rail supports $N$ parallel rails with capacity $NR_2$, that uniform sending achieves global optimality due to topology symmetry, and that an LPT-based discrete flow allocator approaches this optimum with provable bounds. Practically, RailS implements a fully distributed LPT scheduler and a flow-splitting spraying pipeline, delivering substantial gains in bus bandwidth and iteration time across uniform, sparse, skewed, and real MoE workloads, and demonstrating full architectural parallelism exploitation. The work advances topology-aware transport co-design for large-scale MoE systems and offers a scalable path to higher training throughput on Rail-based data-center networks.

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.

Figures (36)

• Figure 1: Hybrid parallelism for a MoE model (data parallelism, pipeline parallelism, expert parallelism, and tensor parallelism).
• Figure 2: Mixtral-MoE
• Figure 3: LLaMA-MoE
• Figure 4: Qwen-MoE
• Figure 5: Non-uniform all-to-all traffic

Theorems & Definitions (8)

• Theorem 1
• proof
• Theorem 2
• proof
• Theorem 3
• proof
• Theorem 4
• proof