Optimizing Mixture of Experts using Dynamic Recompilations
This addresses performance bottlenecks for researchers and practitioners using large-scale Mixture of Experts models, offering incremental improvements in speed and scalability.
The paper tackles the inefficiency of current DNN frameworks in supporting dynamic data flow in Mixture of Experts models, presenting DynaMoE, a library that uses dynamic recompilations to optimize resource use, achieving up to 1.8x speedup and 2.3x larger model sizes compared to existing systems.
The Mixture of Experts architecture allows for outrageously large neural networks by scaling model parameter size independently from computational demand (FLOPs). However, current DNN frameworks cannot effectively support the dynamic data flow in Mixture of Experts, and implementations on top of these frameworks need to use workarounds that introduce significant overheads. To address the limitation of these frameworks, we present DynaMoE, a DNN library that uses dynamic recompilations to optimize and adapt the use of computational resources to the dynamic needs of Mixture of Experts models. Our evaluation shows that DynaMoE achieves a 1.8x speedup and supports 2.3x larger model sizes when compared to existing MoE systems, even when not using recompilations. We then present further optimizations enabled by dynamic recompilations that yield an additional 1.7x speedup while simultaneously reducing memory pressure and improving model quality.