PaperlayerView exact revision on arXiv

Source-linked AI summary

DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale

Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Ammar Ahmad Awan, Jeff Rasley, Yuxiong He

arXiv:2201.05596v2cs.LGcs.AIcs.DC

TL;DR

Large dense-model training increasingly strains practical hardware resources, while MoE models reduce training cost but create parameter, memory, and inference challenges. DeepSpeed-MoE combines expanded autoregressive MoE training, parameter-efficient architectures and compression, and optimized inference, reporting up to 3.7x smaller models and up to 4.5x faster, 9x cheaper inference than quality-equivalent dense models.

  • Problem

    Large dense-model training increasingly strains available hardware, while MoE models introduce substantially larger parameter and memory requirements and challenging inference performance.

  • Method

    DeepSpeed-MoE combines autoregressive MoE training with PR-MoE, Mixture-of-Students distillation, and an optimized end-to-end MoE inference system.

  • Results

    3.7x model-size reduction and up to 4.5x faster, 9x cheaper inference than quality-equivalent dense models are reported.

  • Takeaways & Limitations

    DeepSpeed-MoE supports using sparse MoE models to train and deploy higher-quality models with fewer resources.

  • Takeaways & Limitations

    MoE models remain constrained by larger parameter counts and memory-bandwidth requirements that complicate training and inference.

Abstract

from arXiv · show

As the training of giant dense models hits the boundary on the availability and capability of the hardware resources today, Mixture-of-Experts (MoE) models become one of the most promising model architectures due to their significant training cost reduction compared to a quality-equivalent dense model. Its training cost saving is demonstrated from encoder-decoder models (prior works) to a 5x saving for auto-aggressive language models (this work along with parallel explorations). However, due to the much larger model size and unique architecture, how to provide fast MoE model inference remains challenging and unsolved, limiting its practical usage. To tackle this, we present DeepSpeed-MoE, an end-to-end MoE training and inference solution as part of the DeepSpeed library, including novel MoE architecture designs and model compression techniques that reduce MoE model size by up to 3.7x, and a highly optimized inference system that provides 7.3x better latency and cost compared to existing MoE inference solutions. DeepSpeed-MoE offers an unprecedented scale and efficiency to serve massive MoE models with up to 4.5x faster and 9x cheaper inference compared to quality-equivalent dense models. We hope our innovations and systems help open a promising path to new directions in the large model landscape, a shift from dense to sparse MoE models, where training and deploying higher-quality models with fewer resources becomes more widely possible.

1 Introduction

Rapidly growing dense-model costs motivate MoE, which reduces training compute but introduces major memory and inference challenges. DeepSpeed-MoE addresses these challenges through expanded applications, parameter-efficient architectures, and optimized inference.

  • Motivation: 1000x growth in the largest trained model has intensified pressure on hardware resources as larger models generally improve quality.Megatron-Turing NLG 530B reached half a trillion parameters, while further dense-model scaling is becoming difficult.
  • Motivation: MoE offers sub-linear compute scaling and improved quality without increased training cost, but its application to autoregressive NLG remains less explored.Prior NLP work focused primarily on encoder-decoder and sequence-to-sequence models.
  • Challenges: MoE models can require substantially more parameters and device memory than quality-equivalent dense models, creating training and inference challenges.Switch-Base has 10x more parameters than T5-large, while fitting large MoE models can require thousands of GPUs.
  • Solutions: 5x training-cost reduction is demonstrated for autoregressive NLG models such as GPT-3 and MT-NLG while achieving the same model quality.This extends MoE cost-reduction results beyond encoder-decoder models.
  • Solutions: 3.7x smaller MoE models are achieved with comparable quality through PR-MoE and its distilled Mixture-of-Students variant.PR-MoE reduces parameters by up to 3x, while MoS reduces model size by up to 3.7x with minimal or staged-distillation changes.
  • Solutions: 7.3x lower inference latency and cost is provided by DeepSpeed-MoE compared with existing MoE inference solutions.The system also supports under-25-ms inference for trillion-parameter MoE models and up to 4.5x faster, 9x cheaper inference than quality-equivalent dense models.
  • Implications: DeepSpeed-MoE positions MoE as a more economical alternative to dense models while supporting higher-quality models under fixed hardware resources.The paper frames this as a shift from dense to sparse MoE models.

2 Related Work

Related work established MoE’s scaling and training benefits, including autoregressive NLG applications, but left large-scale inference comparatively underdeveloped. DeepSpeed-MoE extends the scope to integrated architecture, training, and inference optimization.

  • Dense-model scaling: GPT-3 and Megatron-Turing NLG demonstrated quality gains at larger scales, but Megatron-Turing NLG required three months on over 2000 A100 GPUs.These results motivated more resource-efficient approaches to large-model training.
  • MoE scaling: MoE enabled scaling models to large parameter counts with sparse computation, including a 600B-parameter multilingual translation model.Earlier work also applied convolutional MoE layers to a 127B LSTM-based model.
  • Autoregressive NLG: Recent parallel studies applied MoE to autoregressive NLG, whereas this work additionally investigates model design and inference.The paper distinguishes its scope from studies primarily focused on MoE training.
  • Inference systems: Before this work, no MoE system was specifically designed and optimized for inference, despite the development of several MoE training systems.This gap was identified as a barrier to serving large-scale MoE models.
  • Training systems: DeepSpeed MoE training provided a flexible API and supported models up to 8x larger, including a 10-billion-parameter multilingual model trained with 5x less time than an equal-quality dense model.The earlier system targeted scalable MoE training rather than the broader inference focus of this paper.
  • Training systems: FastMoE, Fairseq-MoE, and Tutel advanced MoE training support, with Tutel reporting up to 40 percent improvement over Fairseq.FastMoE combined data and expert parallelism, while Fairseq-MoE offered a generic language-model training pipeline.

3 DeepSpeed-MoE for NLG: Reducing the Training Cost of Language Models by 5 Times

DeepSpeed-MoE applies sparse expert routing to autoregressive NLG models, improving quality at similar compute or matching larger dense models with substantially lower training cost.

  • Training-cost reduction: MoE can significantly improve NLG quality at the same training cost as a dense model, extending sparse modeling to autoregressive language generation.The study evaluates GPT-like models at 350M, 1.3B, and 6.7B scales.
  • Model architecture: Top-1 gating activates only the selected expert for each token, so 1.3B+MoE-128 uses roughly the same per-token training computation as a 1.3B dense model.The MoE layers add 128 experts on every other feedforward layer.
  • Quality evaluation: MoE versions achieve validation loss and six-task zero-shot quality comparable to dense models with 4–5x larger bases.The reported comparisons include 350M+MoE-128 versus 1.3B dense and 1.3B+MoE-128 versus 6.7B dense.
  • Quality evaluation: Related-work comparisons indicate that the proposed MoE models can reach on-par or better quality with fewer parameters on certain tasks, subject to differing data and hyperparameters.The comparison includes 1.3B+MoE-128 and several PR-MoE variants against an 8B+MoE-64 model.
  • Training-cost reduction: 5x lower training cost achieves the model quality of a 6.7B dense NLG model using a 1.3B+MoE-128 model.Each token activates one expert, keeping compute similar to the dense base while increasing total parameters.

4 PR-MoE and MoS: Reducing the Model Size and Improving Parameter Efficiency

PR-MoE combines pyramid allocation of experts with residual routing to reduce MoE parameter counts while preserving model quality. DeepSpeed-MoE extends parallel training to accommodate the varying expert counts across PR-MoE layers.

  • Motivation: 8x more parameters arise in the 1.3B+MoE-128 model despite achieving the same quality as its dense counterpart, motivating parameter-efficiency improvements.The comparison is between 52B parameters in 1.3B+MoE-128 and 6.7B parameters in the dense model.
  • Two Observations and Intuitions: Second-Half-MoE performs significantly better than First-Half-MoE, indicating that deeper layers benefit more from larger numbers of experts.The comparison uses 350M+MoE models with MoE layers placed in either the first or second half of the network.
  • Pyramid Residual MoE Architecture: PR-MoE combines Pyramid-MoE, which uses more experts in later layers, with Residual-MoE, which routes each token through one fixed MLP and one selected expert.The resulting architecture replaces standard MoE layers with PR-MoE layers.
  • System Design: DeepSpeed-MoE supports different expert- and data-parallelism degrees across model regions, enabling efficient PR-MoE training without compromising efficiency or memory requirements.This design addresses the absence of a single suitable expert-parallelism degree when layers contain different numbers of experts.
  • Evaluation of PR-MoE: PR-MoE uses less than one-third of Standard-MoE parameters at 350M scale and about 60% at 1.3B scale while achieving similar accuracy.These comparisons evaluate 350M+PR-MoE-32/64 against 350M+MoE-128 and 1.3B+PR-MoE-64/128 against 1.3B+MoE-128.
  • Evaluation of PR-MoE: PR-MoE reduces the validation-loss gap to around 0.01 relative to 350M+MoE-128, demonstrating parameter efficiency with minimal quality impact.Pyramid-MoE and Residual-MoE each reduce the gap, while their combination reduces it further.

4.2 Mixture-of-Students: Distillation for Even Smaller Model Size and Faster Inference

Mixture-of-Students (MoS) compresses PR-MoE models by reducing expert depth and applying staged knowledge distillation, preserving accuracy while reducing inference cost. Staged distillation avoids the underfitting and late-training accuracy degradation observed with full-process distillation.

  • Distillation for Smaller Models: Layer reduction through knowledge distillation compresses MoE models while reducing model size and computation and preserving the MoE structure.The approach focuses on reducing model depth rather than quantization or pruning.
  • Staged Knowledge Distillation: Full-process distillation initially helps but later hurts validation accuracy because the reduced-capacity student struggles to optimize KD and language-model losses simultaneously.The proposed explanation is that PR-MoE already reduces capacity, so further depth reduction can push the student into underfitting.
  • Mixture-of-Students Construction: MoS reduces each teacher expert branch’s depth, then trains the student to imitate the teacher’s outputs while retaining sparse MoE inference benefits.The student keeps the teacher’s sparsely gated architecture but uses shallower expert branches.
  • PR-MoS: PR-MoS combines PR-MoE with knowledge distillation to further reduce MoE model size.Both the teacher and student are selected as PR-MoE models.
  • Evaluation: Staged KD gives the student a validation curve similar to its teacher and improves zero-shot accuracy on 5 of 6 tasks.Reported gains include 1.1 points on LAMBADA, 6.5 on BoolQ, 1.7 on RACE-h, and 4.5 on TriviaQA; PIQA is the exception.
  • Evaluation: 42.87 and 47.96 average accuracy retain 99.5% and 99.1% of the 350M and 1.3B teacher performance despite 12.5% fewer layers.The teacher accuracies are 43.08 and 48.37, respectively.
  • Inference Implications: PR-MoE reduces training memory by up to 3x, while PR-MoE with MoS reduces inference model size by up to 3.7x while retaining model accuracy.The section connects these reductions to lower inference latency and cost.

5 DeepSpeed-MoE Inference: Serving MoE Models at Unprecedented Scale and Speed

MoE inference is dominated by parameter loading and achievable memory bandwidth when batches are small. DeepSpeed-MoE combines model-size reduction with a multi-GPU system that aggregates bandwidth, achieving substantially better latency and cost than baseline MoE and dense alternatives.

  • Inference Performance: Small-batch MoE inference latency depends primarily on loading model parameters from main memory rather than on compute alone.The two main performance factors are overall model size and achievable memory bandwidth.
  • System Results: 7.3x better latency and lower cost than baseline MoE systems, with up to 4.5x faster and 9x cheaper inference than quality-equivalent dense models.DeepSpeed-MoE uses a multi-GPU system that aggregates memory bandwidth across distributed GPUs.

5.1 Design of DeepSpeed-MoE Inference System

DeepSpeed-MoE targets the MoE inference performance paradox: individual tokens activate few parameters, but batches may require many experts. Its system coordinates routing, parallelism, communication, and kernels to approach the favorable case.

  • Performance Paradox: Top-1 gating activates one expert per MoE layer for each token, but different tokens can activate different experts and collectively require the full model.This creates a gap between the per-token critical path and batch-level parameter access.
  • Design Goal: The system aims to steer inference toward the best case by minimizing each device’s critical data path and maximizing aggregate memory bandwidth.The design goal is explicitly framed as reducing per-device data access while sustaining scale.
  • Parallelism and Routing: Model partitioning and token grouping route tokens with the same critical data path together to reduce per-device data access.This is coordinated with multiple forms of parallelism.
  • Communication: Communication scheduling is coordinated with parallelism to group and route tokens efficiently.The optimization is one of three coordinated system-level strategies.
  • Kernels: Transformer and MoE kernels are optimized to improve per-device performance.Kernel optimization complements the routing and communication optimizations.

5.2 Flexible Combination of Tensor-Slicing, Expert-Slicing, Data Parallelism, and Expert Parallelism

DeepSpeed-MoE combines expert and non-expert parallelism to distribute both parameter types across devices. This multidimensional design supports low-latency inference and scaling to multi-trillion-parameter MoE models.

  • System Design: The inference system uses expert parallelism and expert-parameter slicing for experts, plus data parallelism and tensor-slicing for non-expert parameters.The combination is designed to minimize critical-path parameters while increasing aggregate memory and bandwidth.
  • Expert Parameters: Expert parallelism partitions experts across devices and groups tokens assigned to the same experts.This reduces the expert parameters processed on each device during batch inference.
  • Expert Parallelism: 1.3B+MoE-128 can provide a 1.3-billion-parameter sequential path per device, 5x smaller than its 6.7-billion-parameter quality-equivalent dense model.The stated up-to-5x speed potential assumes no communication overhead.
  • Expert Parameters: Expert-slicing partitions parameters within an expert across multiple GPUs, adding a parallelism dimension for latency-sensitive deployments.This is useful when scaling to more devices than the number of experts.
  • Non-expert Parameters: Non-expert parameters remain limited by single-device memory and bandwidth unless tensor-slicing distributes them within a node.Tensor-slicing enables hundreds of billions of non-expert parameters using aggregate GPU resources.
  • Multidimensional Parallelism: Combining expert parallelism, expert-slicing, tensor-slicing, and data parallelism scales multi-trillion-parameter MoE models across dozens or hundreds of devices.The design uses data-parallel replicas across nodes to avoid communication overhead for non-expert parameters.

5.3 Optimized Communication Subsystem: Grouping and Routing Tokens More Efficiently

DeepSpeed-MoE optimizes expert-parallel communication through hierarchical and parallelism-coordinated all-to-all schemes. These designs reduce communication hops and latency overhead when scaling across many GPUs.

  • Hierarchical All-to-all: Hierarchical all-to-all reduces communication hops from O(p) to O(G + p/G) through intra-node and inter-node exchanges.The process uses data-layout transformations around intra-node and inter-node all-to-all operations.
  • Parallelism Coordinated Communication Optimization: Combining expert and tensor parallelism coordinates all-to-all communication instead of treating each parallelism independently.Tensor-slicing replication lets all-to-all operate within subsets of devices sharing the same tensor-slicing rank.
  • Parallelism Coordinated Communication Optimization: Coordinated communication reduces all-to-all latency from O(p) to O(p/L) when tensor-parallel replication is available.The optimization applies when tensor-parallel operators precede expert-parallel operators.
  • Parallelism Coordinated Communication Optimization: For expert-parallel operators followed by tensor-slicing, all-to-all plus allgather reduces overhead from O(p) to O(p/L) + O(L).The final all-to-all uses the same subset strategy, followed by allgather across tensor-parallel ranks.
  • Parallelism Coordinated Communication Optimization: At 128 GPUs with 8-way tensor-slicing, all-to-all latency overhead falls from (128C1 + C2) to (16C1 + C2).C1 and C2 depend on point-to-point latency, message size, and bandwidth.

5.4 Highly Optimized Transformer and MoE Related Kernels

DeepSpeed-MoE replaces costly sparse MoE operations with fused kernels and explicit data-layout transformations. These changes reduce computation and substantially lower MoE-kernel latency.

  • Kernel Design: DeepSpeed-MoE optimizes gating and token-layout operations by replacing sparse-dense einsums with explicit data transformations and fused kernels.The approach reduces compute complexity from cubic to quadratic for these operations.
  • MoE-Related Computation: The MoE pipeline comprises gating, token sorting by expert assignment, and restoring tokens to their original order.Gating produces the token-to-expert assignment used by the two ordering operations.
  • MoE-Related Computation: Sparse tensor representations create latency through mask construction, top-k selection, cumulative sums, sparse multiplication, and many kernel launches.The passage identifies both representation overhead and excessive kernel invocation as bottlenecks.
  • Kernel Optimization: A fused gating kernel and dense token-to-expert mapping table reduce kernel-launch, memory, and compute overhead.The fused kernel includes top-k, cumulative-sum, and scatter operations.
  • Results: 6x reduction in MoE-kernel-related latency results from the combined kernel and data-layout optimizations.Sparse einsums are replaced by sorting and restoring tokens using the mapping table, reducing complexity from S × E × M × c_e to S × M × c_e.

5.5 Performance Evaluation of DS-MoE Inference

The evaluation examines MoE inference scaling, communication, compression, and comparisons with quality-equivalent dense models. DeepSpeed-MoE achieves simultaneous latency and throughput improvements, including at trillion-parameter scale.

  • Evaluation Scope: MoE inference is evaluated for throughput, latency, scaling, compression, and comparison with quality-equivalent dense models on up to 256 A100 GPUs.The study uses multiple model configurations, including standard MoE, PR-MoE, and PR-MoE+MoS variants.
  • 5.5.1 Achieving Low Latency and Super-Linear Throughput Increase Simultaneously: DeepSpeed-MoE achieves super-linear total-throughput scaling as GPUs increase from 8 to 64 for the 52B MoE model.Throughput per GPU increases because each GPU serves fewer experts and reads less expert-parameter data from memory.
  • 5.5.2 Low Latency and High Throughput at Unprecedented Scale: 7.3x lower latency and 7.3x higher throughput are achieved versus the baseline across MoE model sizes from 107B to 2T parameters.The model-size study uses PyTorch and DeepSpeed configurations described for Figure 11.
  • 5.5.2 Low Latency and High Throughput at Unprecedented Scale: Under 25 ms inference latency is achieved for a trillion-parameter MoE model by exploiting hundreds of GPUs in parallel.This result concerns DeepSpeed-MoE inference at unprecedented scale.
  • 5.5.3 Model Compression: PR-MoE+MoS reduces the required GPUs from 32 to 16 and provides the lowest latency among the evaluated model variants.Figures 12 and 13 compare standard MoE, PR-MoE, and PR-MoE+MoS with a PyTorch standard-MoE baseline.
  • 5.5.4 Better Latency and Throughput Than Quality-Equivalent Dense Models: 4.5x faster and 9x cheaper inference are achieved for trillion-scale PR-MoE+MoS versus a 175B quality-equivalent dense model.The billion-scale comparison reports 2.4x faster and cheaper inference, while the larger-model benefit is greater.

6 Looking Forward to the Next Generation of AI Scale

Growing model sizes are approaching hardware limits for training and serving. DeepSpeed-MoE combines memory-reducing model innovations with an efficient inference framework to support larger-scale AI without increasing compute resources.

  • Motivation: Exponential model-size growth is reaching the boundary of what modern supercomputing clusters can train and serve with available hardware resources.The paper frames improving training and inference efficiency on current hardware as an alternative to waiting for new hardware.
  • Open Challenge: Prior work reduced training costs for large NLG MoE models, but efficient MoE serving with latency and cost comparable to dense models remained unaddressed.The paper identifies this serving gap as a challenge for real-world deployment.
  • Proposed Direction: PR-MoE and MoS reduce MoE memory requirements, while the DeepSpeed-MoE inference framework targets low latency and cost at large model scale.Together, these innovations make MoE models feasible to serve and potentially better than quality-equivalent dense models on latency and cost.
  • Conclusion: The combined innovations offer a path toward training and inference for next-generation AI scale without increasing compute resources.The paper presents a possible shift from dense to sparse MoE models for deploying higher-quality models with fewer resources.

Contributions

The paper assigns leadership across NLG training, MoE designs, inference-system development, and DeepSpeed software integration.

  • SR designed the NLG training experiments and architected the inference system.
  • CL led the NLG training experiments, while ZY led PR-MoE design, experiments, and system support.
  • MZ led MoS design and experiments and memory-efficient checkpointing; RYA and AAA led inference-system development and experiments.
  • JR developed, debugged, and integrated multiple software features into DeepSpeed.
Loading 2201.05596v2…