DeepSpeed‑MoE: End‑to‑End Training and Inference Solutions for Mixture‑of‑Experts Models

Preface : DeepSpeed finally addresses large‑scale Mixture‑of‑Experts (MoE) models, with a particular focus on improving inference performance.

Motivation : Training massive models such as Megatron‑Turing NLG 530B consumes millions of GPU‑hours; MoE’s sparse routing can achieve comparable convergence with far lower compute, but inference acceleration remains a challenge.

DeepSpeed‑MoE Overview : An end‑to‑end solution embedded in the DeepSpeed library that provides novel MoE structures, model‑compression methods, and a highly optimized inference system.

Model Optimisation – PR‑MoE

PR‑MoE combines two new structures: Pyramid‑MoE and Residual‑MoE.

1. Pyramid‑MoE

Observes that placing the same number of experts deeper in the network yields better loss; thus it allocates few experts in shallow layers and more in deeper layers, forming a pyramid shape that reduces parameters while preserving accuracy.

2. Residual‑MoE

Finds Top‑2 gating outperforms Top‑1 because a second expert can correct the first. PR‑MoE fixes the first expert and lets only the second participate in gating, achieving Top‑2‑like performance with lower latency.

3. PR‑MoE (Combined)

The hybrid architecture inherits the parameter efficiency of Pyramid‑MoE and the accuracy of Residual‑MoE, delivering fewer parameters, higher throughput, and comparable precision to standard MoE.

Model Distillation – PR‑MoS

PR‑MoS is a student version of PR‑MoE with reduced depth. It retains the MoE structure during knowledge distillation, unlike prior methods that collapse MoE to dense models. A staged KD trick (early‑stop when teacher‑student performance intersect) prevents over‑distillation.

Distributed Strategy – Inference Serving

DeepSpeed employs a mixed parallelism scheme: uniform Data Parallelism (DP) across all layers, and variable Expert Parallelism (EP) per layer, supplemented by Expert Slicing when EP cannot match DP.

DP degree is constant N.

EP degree varies per layer; if EP < N, additional DP compensates.

Communication optimisations include hierarchical All‑to‑All (intra‑node then inter‑node) and specialised routing for Expert Parallel + Expert Slicing.

Kernel Optimisation

Two main changes: (a) replace sparse one‑hot routing tables with dense mapping tables and fuse the sparse einsum; (b) fuse the entire gating logic into a single kernel to reduce memory traffic.

Performance

Benchmarks on Azure A100 clusters show DeepSpeed‑MoE reduces inference latency by up to 7.3× compared with dense baselines, cuts resource consumption by 9×, and achieves 4.5× speed‑up while maintaining quality. Scaling experiments up to 2‑trillion‑parameter models confirm the trend.

Overall, DeepSpeed‑MoE pushes the limits of both system‑level optimisation and algorithmic innovation for MoE models, delivering substantial inference efficiency gains.