DeepSpeed-Megatron MoE initialization flow creating expert-parallel groups during model wrapping

MoE Internals: DeepSpeed-Megatron Expert Parallel Implementation

MoE Internals: DeepSpeed-Megatron Expert Parallel Implementation DeepSpeed-Megatron MoE is easy to misread if you only inspect one initialization function. Megatron builds the usual TP, PP, and DP topology. DeepSpeed then wraps the model, finds MoE modules, and gives those modules expert-parallel groups. The topology is split across the training framework and the MoE layer implementation. This post connects the pieces: deepspeed.initialize(), EP groups, expert-DP groups, MoELayer, and the difference between a DeepSpeed MoE wrapper and a Megatron-integrated SwitchMLP-style module. For routing, capacity, and All-to-All fundamentals, start with MoE Parallelism Principles. ...

May 17, 2025 · 8 min · Duo An
GShard-style MoE layer with gate, experts, and top-k combine

MoE Parallelism Principles: GShard, Expert Parallel, and All-to-All

MoE Parallelism Principles: GShard, Expert Parallel, and All-to-All Mixture-of-Experts is sparse compute wrapped in a communication problem. The model idea is simple: replace a dense FFN with many expert FFNs and route each token to a few of them. The distributed-training problem starts when those experts live on different GPUs. Shazeer et al. introduced the sparsely-gated MoE layer in Outrageously Large Neural Networks. GShard made the pattern practical at scale with top-2 routing, capacity, load-balancing loss, and all-to-all dispatch in GShard. Switch simplified the router to top-1 in Switch Transformers. This post explains the shared systems model underneath them. ...

May 10, 2025 · 8 min · Duo An