Unlocking Sparse MoE Large Model Training with Megatron-Core on Alibaba Cloud

Overview

With the rapid evolution of large‑model technology, model size and structure are expanding quickly, but training and inference costs remain a major challenge. Model sparsification, especially Mixture‑of‑Experts (MoE), can dramatically reduce computation and storage while preserving or even improving performance. Alibaba Cloud PAI and NVIDIA collaborated to build a scalable MoE training solution based on the Megatron‑Core MoE framework.

MoE Algorithm Introduction

MoE (Mixture of Experts) replaces the traditional Feed‑Forward Network (FFN) in Transformers with a set of expert networks. During training, a router selects one expert per token, allowing the model to increase parameters without increasing FLOPs. This yields significant training and inference throughput gains compared to dense models of similar quality.

Sparse MoE Layer : Replaces the FFN with multiple expert networks, each typically a small FFN.

Router : Directs tokens to specific experts, e.g., token "More" may go to expert 2 while "Parameters" goes to expert 1.

Load‑balancing is critical because some experts may receive many tokens while others receive few, reducing training efficiency. An auxiliary loss encourages equal token distribution across experts, and token‑capacity mechanisms can drop or pad tokens to maintain balance.

Mixtral‑8x7B and Dropless MoE

Mixtral‑8x7B adopts the dropless MoE algorithm from the Megablocks paper, storing token‑to‑expert assignments in expert_indices and probabilities in Probabilities. Tokens exceeding expert capacity are dropped, and the remaining tokens are processed with a batched matrix multiplication, followed by a weighted sum using normalized softmax.

Megatron‑Core MoE Training Framework

Megatron‑Core is NVIDIA’s lightweight, production‑ready framework for large‑scale LLM training. Version 0.5 adds native support for massive MoE models, offering features such as:

Parallelism

Expert Parallelism (EP) – each rank handles one or more experts.

3D Parallelism (Data, Tensor, Pipeline, Sequence).

Future Context Parallelism for longer sequences.

Token Dispatch Mechanism

Supports dropless MoE (no token dropping) with upcoming token‑drop MoE support.

Router and Load Balancing

Provides Top‑K router and upcoming Expert Choice router, with load‑balancing algorithms like Sinkhorn (S‑BASE), Z‑Loss, and Load‑Balancing Loss.

Grouped GEMM

Utilizes CUTLASS 2.8’s Grouped GEMM to efficiently handle variable‑length inputs from multiple experts, improving SM utilization and performance.

Upcoming Features

Context Parallelism

FP8 low‑precision training

FP8 Grouped GEMM

Token‑Drop MoE

MoE Platform Tools on Alibaba Cloud

The training stack consists of three layers:

PAI platform (DSW for interactive notebooks and DLC for multi‑node, multi‑GPU training).

PAI‑Megatron‑Patch – bridges open‑source LLMs and Megatron.

NVIDIA Megatron‑Core – provides the core training engine.

Users can develop data pipelines in DSW, convert HuggingFace weights to Megatron format, and launch training with scripts that pass hyper‑parameters (e.g., learning rate, Flash Attention) to the engine.

HF to Megatron Weight Conversion

Converting HuggingFace checkpoints to Megatron requires careful splitting and merging of weight matrices, especially for tensor‑parallelism (TP>1). For example, the MLP gate_proj and up_proj are split across TP=2 and recombined into Megatron’s dense_h_to_4h format. Expert distribution parameters (expert_model_parallel_size, world_size) are also needed for MoE models.

Training and Evaluation Results

Experiments on Mixtral‑8x7B demonstrate:

Zero‑shot loss alignment between HuggingFace and Megatron checkpoints.

Training loss convergence after ~2k steps for both pre‑training and fine‑tuning.

Code‑generation downstream task (Human‑Eval) shows significant improvement after instruction fine‑tuning.

Throughput tests on A800 GPUs (2 × 16 cards) indicate Megatron‑Core MoE is ~10% faster than Megablocks under comparable settings.

Conclusion

The Megatron‑Core MoE training toolkit, integrated with Alibaba Cloud PAI, provides a reliable and efficient pipeline for sparse large‑model training, weight conversion, continued pre‑training, fine‑tuning, and downstream tasks such as code generation. Future releases will add more high‑quality LLM best‑practice guides.