Search papers, labs, and topics across Lattice.
TopoEP eliminates dynamic-routing stragglers in expert-parallel MoE training by performing per-layer, per-microbatch expert replication and token rerouting entirely on GPUs. By using a deterministic solver that accounts for hierarchical interconnect topologies, it removes device-host data transfers and critical-path cross-rank synchronization. Benchmarked on a 32-GPU NVIDIA H800 cluster within Megatron-LM, the system improves end-to-end training throughput by 6.2% to 11.4% across three representative MoE architectures.
Dynamic MoE load balancing can run at every single microbatch with zero CPU synchronization by having GPUs deterministically solve bitwise-identical expert replication plans locally, boosting cluster throughput by up to 11.4%.
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.