TAOT: Topology-Aware Optimal Transport for Dynamic Expert Replica Placement in MoE Training

2026-08-04

AI Summary

This work addresses severe load imbalance in Mixture-of-Experts (MoE) training caused by dynamic expert routing and the high communication overhead of existing dynamic replication strategies, which neglect cross-node topology when migrating expert weights. To tackle this, we introduce— for the first time—topology-aware entropy-regularized optimal transport theory to model expert replica placement as a scheduling problem between overloaded and underutilized resources. We solve the resulting transport plan via Sinkhorn-Knopp iterations and integrate it with integer replica matching and token assignment to produce an executable schedule. By overlapping communication and computation to hide latency, our method jointly optimizes load balancing and inter-node communication costs. Experiments show it achieves up to 1.43× end-to-end training speedup, matches or exceeds state-of-the-art load balance, and reduces weighted communication cost by up to 74%.

Abstract

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.

Problem

• Mixture-of-Experts
• load imbalance
• expert replica placement
• communication cost
• multi-node topology

Innovation

• optimal transport
• topology-aware
• expert replica placement
• Mixture-of-Experts
• load balancing