MoE LLMs: How Alibaba Cloud & NVIDIA Megatron-Core Accelerate Training

Introduction

Alibaba Cloud AI Platform PAI and NVIDIA Megatron-Core teams present their joint work on implementing and optimizing Mixture-of-Experts (MoE) large language models (LLMs).

MoE Overview

MoE splits a dense model into multiple expert sub‑networks, activating only a subset of experts per token via a routing mechanism, which reduces parameter updates and computational cost during training and inference.

Historical Development

Early 1990s concepts of multiple experts laid the theoretical foundation. In 2017 deep learning growth highlighted capacity limits of dense models. Google first combined MoE with RNNs, later introduced MoE into Transformers (Gshard) and Switch Transformers, scaling parameters to the trillion level.

Key Research Works

Sparsely‑Gated MoE introduced softmax gating, later improved with noise injection and auxiliary loss for balanced expert load.

Google Gshard added expert capacity limits, residual bypass, top‑k routing, Sinkhorn load‑balancing, and local‑group dispatching to reduce communication.

Switch Transformer adopted aggressive top‑1 routing and expert dropout to cut computation.

Megablocks expressed MoE layers as block‑sparse matrices using BCSR format for efficient GEMM.

Megatron‑Core MoE Features

Megatron‑Core provides a lightweight, modular framework supporting expert parallelism, 3‑D parallelism (data, tensor, pipeline, sequence), dropless token routing, multiple routing strategies (top‑k, Sinkhorn, z‑loss), GroupedGMM for variable‑length inputs, and optimized CUDA kernels.

PAI‑Megatron‑Patch Toolchain

The PAI‑Megatron‑Patch library converts HuggingFace model checkpoints to Megatron‑Core format, handling layernorm, attention, and MLP weight mapping, enabling seamless migration of open‑source models such as Mixtral 8×7B.

Experimental Results

Training Mixtral 8×7B with 8 experts on 16 GPUs (TP=4, EP=4) achieved loss convergence around 1.9 after 2.4K steps. Subsequent pre‑training and fine‑tuning stages showed stable loss reduction. In a code‑generation benchmark (HumanEvol), fine‑tuned models improved from 45.73 % to 53.05 % accuracy, outperforming Megablocks under comparable resources.

Best‑Practice Guide

Alibaba Cloud PAI offers a complete AI development workflow: data ingestion from OSS/NAS, distributed training on PAI‑DLC or DSW, checkpoint export, offline inference, evaluation, and one‑click deployment to EAS for online serving.

Future Directions

The teams will continue to deepen collaboration, aiming to further improve dense and sparse model training efficiency and contribute to AGI research, inviting developers to join the open‑source community.