SMILE: Scaling Mixture-of-Experts with Efficient Bi-level Routing

TL;DR

SMILE improves the training efficiency of Mixture-of-Experts models by introducing bi-level routing that better utilizes heterogeneous network bandwidth, achieving significant speedups without sacrificing convergence.

Contribution

The paper proposes SMILE, a novel bi-level routing method that enhances MoE training scalability and efficiency by addressing network congestion issues.

Findings

Achieves 2.5x pretraining throughput speedup over Switch Transformer.

Maintains convergence speed despite increased experts.

Effectively exploits heterogeneous network bandwidth.

Abstract

The mixture of Expert (MoE) parallelism is a recent advancement that scales up the model size with constant computational cost. MoE selects different sets of parameters (i.e., experts) for each incoming token, resulting in a sparsely-activated model. Despite several successful applications of MoE, its training efficiency degrades significantly as the number of experts increases. The routing stage in MoE relies on the efficiency of the All2All communication collective, which suffers from network congestion and has poor scalability. To mitigate these issues, we introduce SMILE, which exploits heterogeneous network bandwidth and splits a single-step routing into bi-level routing. Our experimental results show that the proposed method obtains a 2.5x speedup over Switch Transformer in terms of pretraining throughput on the Colossal Clean Crawled Corpus without losing any convergence speed.