TopoEP: Topology-Aware Load Balancing for Expert-Parallel MoE Training
Organizations: University of Science and Technology of China
Abstract
Dynamic routing creates severe load imbalance in large-scale expert-parallel Mixture-of-Experts (MoE) training, turning GPUs that host hot experts into stragglers. As each MoE layer waits for its slowest rank, these stragglers prolong the expert-parallel stage and reduce overall training efficiency. Existing expert-parallelism load-balancing (EPLB) systems commonly compute load-balancing plans on the CPU, incurring device--host data transfers and cross-rank synchronization that make scheduling at every layer and microbatch expensive. Their planning formulations also overlook the hierarchical communication costs of modern scale-up and scale-out GPU clusters. We present \textit{TopoEP}, a GPU-native, topology-aware load-balancing system for large-scale MoE training. At each MoE layer and training microbatch, \textit{TopoEP} converts the current routing result into hot-expert replication and token-rerouting decisions and executes the resulting plan without data-dependent host synchronization, reducing critical-path overhead. To generate these decisions, \textit{TopoEP} uses a deterministic GPU solver that performs inter-node placement followed by intra-node refinement, allowing all ranks to independently produce bitwise-identical plans. On a 32-GPU NVIDIA H800 cluster, integrating \textit{TopoEP} with Megatron-LM improves end-to-end training throughput by 6.2%--11.4% across three representative MoE models.
Figures & tables
| Model | MoE layers | Experts | Top- | Hidden | Interm. size | PP / EP |
| Qwen3-30B-A3B [ 31 ] | 5 | 128 | 8 | 768 | 1/32 | |
| GLM-4.5-Air [ 36 ] | 5 | 128 | 8 | 2/16 | ||
| DeepSeek-V2 [ 3 ] | 4 | 160 | 6 | 2/16 |