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