# Maximizing Parallelism in Distributed Training for Huge Neural Networks

arXiv: https://arxiv.org/abs/2105.14450 · HTML: https://mlsanity.com/paper/2105.14450/

## Key facts

- Cited by 23 papers in the map (data through 14 Oct 2026).
- Builds on 15 of its 25 references that are in the map.
- First posted on arXiv 30 May 2021.

## Builds on (most cited first, 15 of 15)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [Large Batch Training of Convolutional Networks](https://mlsanity.com/paper/1708.03888/) (2017)
- [Don't Decay the Learning Rate, Increase the Batch Size](https://mlsanity.com/paper/1711.00489/) (2017)
- [Train longer, generalize better: closing the generalization gap in large batch training of neural networks](https://mlsanity.com/paper/1705.08741/) (2017)
- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://mlsanity.com/paper/2104.07857/) (2021)
- [Mesh-TensorFlow: Deep Learning for Supercomputers](https://mlsanity.com/paper/1811.02084/) (2018)
- [PipeDream: Fast and Efficient Pipeline Parallel DNN Training](https://mlsanity.com/paper/1806.03377/) (2018)

## Built on by (most cited first, 23 of 23)

- [A Survey of Large Language Models](https://mlsanity.com/paper/2303.18223/) (2023)
- [Efficient Large Language Models: A Survey](https://mlsanity.com/paper/2312.03863/) (2023)
- [Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training](https://mlsanity.com/paper/2110.14883/) (2021)
- [A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training](https://mlsanity.com/paper/2303.06318/) (2023)
- [RollPacker: Mitigating Long-Tail Rollouts for Fast, Synchronous RL Post-Training](https://mlsanity.com/paper/2509.21009/) (2025)
- [On Efficient Training of Large-Scale Deep Learning Models: A Literature Review](https://mlsanity.com/paper/2304.03589/) (2023)
- [The Efficiency Spectrum of Large Language Models: An Algorithmic Survey](https://mlsanity.com/paper/2312.00678/) (2023)
- [PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management](https://mlsanity.com/paper/2108.05818/) (2021)
- [LLMem: Estimating GPU Memory Usage for Fine-Tuning Pre-Trained LLMs](https://mlsanity.com/paper/2404.10933/) (2024)
- [Merak: An Efficient Distributed DNN Training Framework with Automated 3D Parallelism for Giant Foundation Models](https://mlsanity.com/paper/2206.04959/) (2022)
- [Training and Serving System of Foundation Models: A Comprehensive Survey](https://mlsanity.com/paper/2401.02643/) (2024)
- [A 4D Hybrid Algorithm to Scale Parallel Training to Thousands of GPUs](https://mlsanity.com/paper/2305.13525/) (2023)
- [Colossal-Auto: Unified Automation of Parallelization and Activation Checkpoint for Large-scale Models](https://mlsanity.com/paper/2302.02599/) (2023)
- [Rethinking Memory and Communication Cost for Efficient Large Language Model Training](https://mlsanity.com/paper/2310.06003/) (2023)
- [Pro-Prophet: A Systematic Load Balancing Method for Efficient Parallel Training of Large-scale MoE Models](https://mlsanity.com/paper/2411.10003/) (2024)
- [A Survey on Memory-Efficient Transformer-Based Model Training in AI for Science](https://mlsanity.com/paper/2501.11847/) (2025)
- [TrainVerify: Equivalence-Based Verification for Distributed LLM Training](https://mlsanity.com/paper/2506.15961/) (2025)
- [The Rise of Sparse Mixture-of-Experts: A Survey from Algorithmic Foundations to Decentralized Architectures and Vertical Domain Applications](https://mlsanity.com/paper/2602.08019/) (2026)
- [Dive into Big Model Training](https://mlsanity.com/paper/2207.11912/) (2022)
- [TPLA: Tensor Parallel Latent Attention for Efficient Disaggregated Prefill and Decode Inference](https://mlsanity.com/paper/2508.15881/) (2025)
- [Mixed Sparsity Training: Achieving 4× FLOP Reduction for Transformer Pretraining](https://mlsanity.com/paper/2408.11746/) (2024)
- [Comprehensive Performance Modeling and System Design Insights for Foundation Models](https://mlsanity.com/paper/2410.00273/) (2024)
- [You Don't Need All Attentions: Distributed Dynamic Fine-Tuning for Foundation Models](https://mlsanity.com/paper/2504.12471/) (2025)

## Abstract (arXiv)

The recent Natural Language Processing techniques have been refreshing the state-of-the-art performance at an incredible speed. Training huge language models is therefore an imperative demand in both industry and academy. However, huge language models impose challenges to both hardware and software. Graphical processing units (GPUs) are iterated frequently to meet the exploding demand, and a variety of ASICs like TPUs are spawned. However, there is still a tension between the fast growth of the extremely huge models and the fact that Moore's law is approaching the end. To this end, many model parallelism techniques are proposed to distribute the model parameters to multiple devices, so as to alleviate the tension on both memory and computation. Our work is the first to introduce a 3-dimensional model parallelism for expediting huge language models. By reaching a perfect load balance, our approach presents smaller memory and communication cost than existing state-of-the-art 1-D and 2-D model parallelism. Our experiments on 64 TACC's V100 GPUs show that our 3-D parallelism outperforms the 1-D and 2-D parallelism with 2.32x and 1.57x speedup, respectively.