SYMI: Efficient Mixture-of-Experts Training via Model and Optimizer State Decoupling
This addresses a critical bottleneck for researchers and practitioners scaling large MoE models efficiently, though it is an incremental improvement on existing systems.
The paper tackles the performance-accuracy trade-off in Mixture-of-Experts (MoE) model training caused by load imbalance across experts, by introducing SYMI, an adaptive system that decouples expert parameters from optimizer state placement, achieving 30.5% and 25.9% faster time-to-convergence compared to DeepSpeed and FlexMoE.
Mixture-of-Experts (MoE) models have become a widely-adopted solution to continue scaling model sizes without a corresponding linear increase in compute. During MoE model training, each input token is dynamically routed to a subset of experts -- sparsely-activated feed-forward networks -- within each transformer layer. The distribution of tokens assigned to each expert varies widely and rapidly over the course of training. To handle the wide load imbalance across experts, current systems are forced to either drop tokens assigned to popular experts, degrading convergence, or frequently rebalance resources allocated to each expert based on popularity, incurring high state migration overheads. To break this performance-accuracy tradeoff, we introduce SYMI, an adaptive MoE training system. The key insight of SYMI is to decouple the placement of expert parameters from their large optimizer state. SYMI statically partitions the optimizer of each expert across all training nodes. Meanwhile, SYMI dynamically adjusts the placement of expert parameters by repurposing existing weight updates, avoiding migration overheads. In doing so, SYMI right-sizes the GPU resources allocated to each expert, on a per-iteration basis, with minimal overhead. Compared to state-of-the-art MoE training systems, DeepSpeed and FlexMoE, SYMI is able to achieve a 30.5% and 25.9% faster time-to-convergence, respectively.