Llama 3 Meets MoE: Efficient Upcycling
Aditya Vavre, Ethan He, Dennis Liu, Zijie Yan, June Yang, Nima Tajbakhsh, Ashwath Aithal
TL;DR
This work tackles the prohibitive compute costs of scaling LLMs by leveraging Mixture-of-Experts (MoE) to expand capacity efficiently. It introduces an upcycling-based recipe that initializes an 8-Expert Top-2 MoE from a pre-trained Llama 3-8B dense checkpoint and trains on a modest data blend, using under $1\%$ of pre-training compute. The approach, implemented with Megatron-Core 5-D parallelism and MoE Parallel Folding, achieves a notable $2\%$ absolute gain on 0-shot MMLU and reaches a Model FLOPs Utilization of $46.8\%$ during training, while significantly reducing GPU-hour costs via online upcycling in NeMo. The paper also provides empirical guidance on capacity-factor settings and router choices, demonstrating a practical, cost-effective pathway to high-capacity MoE models.
Abstract
Scaling large language models (LLMs) significantly improves performance but comes with prohibitive computational costs. Mixture-of-Experts (MoE) models offer an efficient alternative, increasing capacity without a proportional rise in compute requirements. However, training MoE models from scratch poses challenges like overfitting and routing instability. We present an efficient training recipe leveraging pre-trained dense checkpoints, training an 8-Expert Top-2 MoE model from Llama 3-8B with less than $1\%$ of typical pre-training compute. Our approach enhances downstream performance on academic benchmarks, achieving a $\textbf{2%}$ improvement in 0-shot accuracy on MMLU, while reaching a Model FLOPs Utilization (MFU) of $\textbf{46.8%}$ during training using our framework. We also integrate online upcycling in NeMo for seamless use of pre-trained weights, enabling cost-effective development of high-capacity MoE models.