[Paper Club] Upcycling Large Language Models into Mixture of Experts
Oct 29, 2024 · 39:02
NVIDIA's Ethan He presents Megatron-Core MoE and the upcycling of dense LLMs into Mixture-of-Experts models, showing that upcycling a 15B-param dense model into a 64-expert MoE and training on 1T tokens yields 5% lower validation loss and 4% higher MMLU than continued dense training on the same compute. Key techniques include swapping the router order (softmax-then-top-K instead of top-K-then-softmax) with a 4× output scaling to preserve initial forward pass behavior, and initializing fine-grained MoE routers by duplicating half the weights so each shard group selects identically. A high learning rate matching the original pre-training peak is critical; constant low LR causes catastrophic forgetting. The upcycled model matches the compute of a 1.7× larger dense model per scaling laws, but data quality remains paramount—continued dense training still delivered a 20% MMLU jump. Megatron-Core's fused permutation, GroupedGEMM, and expert parallelism are available as a standalone library.
