A training-infrastructure method for the load imbalance MoE's dynamic routing causes in expert-parallel training. Prior dynamic-replica schemes copy hot experts onto idle ranks but optimize load balance alone, ignoring the cost of moving expert weights across a multi-node topology — so the extra cross-node communication can outweigh the balancing gain. TAOT casts overload-vs-spare-capacity as an entropy-regularized optimal-transport problem with a communication-cost matrix, solves it with Sinkhorn-Knopp to produce rank-level flow hints, and overlaps guest-weight transfer with home-expert compute to hide the overhead. 1.43× end-to-end MoE training speedup, balance quality competitive with SOTA, and the lowest weighted expert-communication cost across configs (up to 74% reduction). Baidu, Inc. (senior author Dou Shen).

Paper

moetraininginfrastructureresearch