Search papers, labs, and topics across Lattice.
This paper introduces TAOT, a topology-aware optimal transport method designed to optimize dynamic expert-replica placement in Mixture-of-Experts (MoE) training. By framing the load balancing and communication costs as an entropy-regularized optimal transport problem, TAOT effectively reduces the communication overhead associated with expert weight transfers while maintaining competitive load balance. Experimental results demonstrate a 1.43x speedup in MoE training and a significant 74% reduction in weighted expert-communication costs compared to existing methods.
Achieving a 1.43x speedup in MoE training while slashing communication costs by up to 74% could redefine efficiency benchmarks in large-scale LLM training.
Mixture-of-Experts (MoE) has become a key architecture for scaling large language models (LLMs), yet its dynamic routing causes severe load imbalance in expert-parallel training. Existing dynamic-replica methods copy hot experts onto idle ranks to share computation, but they optimize load balance alone and ignore the cost of moving expert weights across a multi-node topology, so the resulting cross-node communication can outweigh the balancing gain and inflate training cost. We present TAOT, a topology-aware optimal transport method for dynamic expert-replica placement. TAOT models the overload on hot ranks and the spare capacity on lightly loaded ranks as a balanced entropy-regularized optimal transport problem with a communication-cost matrix, solves it with Sinkhorn-Knopp iterations to produce rank-level flow hints, and combines integer replica matching with token assignment into an executable schedule. At the system level, it overlaps guest-weight transfer with home-expert computation to hide the communication overhead. Experiments show TAOT achieves a 1.43x end-to-end MoE training speedup, reaches balance quality competitive with or better than existing state-of-the-art methods, and attains the lowest weighted expert-communication cost across all configurations, with up to a 74% reduction.
