Revisiting MoE and Dense Speed-Accuracy Comparisons for LLM Training

TL;DR

This work tackles the unfair comparison between Mixture-of-Experts (MoE) and dense LLMs by arguing that FLOPs understate MoE communication overhead and total compute. It proposes measuring model complexity with step time and using the Chinchilla compute-optimal budget, coupled with a 3D Data/Expert/Model sharding strategy implemented via GShard/GSPMD, to enable scalable MoE training. Across scales of 6.4B, 12.6B, and 29.6B parameters, MoEs consistently outperform dense models on speed-accuracy benchmarks, with the 3D sharding approach keeping dense-to-MoE step-time overhead under $20 ext{%}$. The results, together with an accessible implementation, suggest MoE can achieve meaningful gains in practical large-scale training settings and provide a framework for fair, hardware-aware comparisons.

Abstract

Mixture-of-Experts (MoE) enjoys performance gain by increasing model capacity while keeping computation cost constant. When comparing MoE to dense models, prior work typically adopt the following setting: 1) use FLOPs or activated parameters as a measure of model complexity; 2) train all models to the same number of tokens. We argue that this setting favors MoE as FLOPs and activated parameters do not accurately measure the communication overhead in sparse layers, leading to a larger actual training budget for MoE. In this work, we revisit the settings by adopting step time as a more accurate measure of model complexity, and by determining the total compute budget under the Chinchilla compute-optimal settings. To efficiently run MoE on modern accelerators, we adopt a 3D sharding method that keeps the dense-to-MoE step time increase within a healthy range. We evaluate MoE and dense LLMs on a set of nine 0-shot and two 1-shot English tasks, as well as MMLU 5-shot and GSM8K 8-shot across three model scales at 6.4B, 12.6B, and 29.6B. Experimental results show that even under these settings, MoE consistently outperform dense LLMs on the speed-accuracy trade-off curve with meaningful gaps. Our full model implementation and sharding strategy has been released at~\url{https://github.com/apple/axlearn}

Figures (1)

• Figure 1: Comparing MoEs designed from a wider range of choices to the 6.4B dense. We vary the dense backbone scale and the number of sparse layers. On the right figure, L2, E4, E2 represent last-2, every-4, every-2 number of layers, respectively.