A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training

Siddharth Singh, Olatunji Ruwase, Ammar Ahmad Awan, Samyam Rajbhandari, Yuxiong He +1 more
3/11/2023
cs.LGcs.AIcs.DC

No official code found. Best implementation: deepspeedai/DeepSpeed (43,229 stars, Apache-2.0, PyTorch, last commit 2026-04-29).

Abstract

Mixture-of-Experts (MoE) is a neural network architecture that adds sparsely activated expert blocks to a base model, increasing the number of parameters without impacting computational costs. However, current distributed deep learning frameworks are limited in their ability to train high-quality MoE models with large base models. In this work, we present DeepSpeed-TED, a novel, three-dimensional, hybrid parallel algorithm that combines data, tensor, and expert parallelism to enable the training of MoE models with 4 to 8x larger base models than the current state-of-the-art. We also describe memory optimizations in the optimizer step, and communication optimizations that eliminate unnecessary data movement. We implement our approach in DeepSpeed and achieve speedups of 26% over a baseline (i.e. without our communication optimizations) when training a 40 billion parameter MoE model (6.7 billion base model with 16 experts) on 128 V100 GPUs.

View on arXivView PDF

Code Implementations(1)

deepspeedai/DeepSpeedCommunity100%
PyTorchApache-2.0

DeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.

43,2295,039C, C++Jan 23, 20205 months agoApache-2.0
billion-parameterscompressiondata-parallelismdeep-learninggpu+8 more

Similar Papers

MoNTA: Accelerating Mixture-of-Experts Training with Network-Traffc-Aware Parallel Optimization

90% match2024

Accelerating MoE Model Inference with Expert Sharding

88% match2025

Pipeline MoE: A Flexible MoE Implementation with Pipeline Parallelism

87% match2023

MoE Parallel Folding: Heterogeneous Parallelism Mappings for Efficient Large-Scale MoE Model Training with Megatron Core

87% match2025

X-MoE: Enabling Scalable Training for Emerging Mixture-of-Experts Architectures on HPC Platforms

87% match1 repo292025

MoE-DisCo:Low Economy Cost Training Mixture-of-Experts Models

85% match2026

Cite this paper

@article{singh2023hybrid,
  title  = {A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training},
  author = {Siddharth Singh and Olatunji Ruwase and Ammar Ahmad Awan and Samyam Rajbhandari and Yuxiong He and Abhinav Bhatele},
  year   = {2023},
  eprint = {2303.06318},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2303.06318v2}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2303.06318.svg)](https://codeofpaper.com/paper/2303.06318)

Discussion