Maximizing Parallelism in Distributed Training for Huge Neural Networks

Abstract

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.

Technical Report of NUS HPC-AI Lab (https://ai.comp.nus.edu.sg). The leading two authors have equal contributions

References in corpus (9)

• Language Models are Few-Shot Learners
• Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• Don't Decay the Learning Rate, Increase the Batch Size
• On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima
• Large Batch Training of Convolutional Networks
• Train longer, generalize better: closing the generalization gap in large batch training of neural networks
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
• ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning

Cited by in corpus (5)

• Merak: An Efficient Distributed DNN Training Framework with Automated 3D Parallelism for Giant Foundation Models
• A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training
• PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management
• Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training
• AutoDDL: Automatic Distributed Deep Learning with Near-Optimal Bandwidth Cost