Layered gradient accumulation and modular pipeline parallelism: fast and efficient training of large language models

Abstract

The advent of the transformer has sparked a quick growth in the size of language models, far outpacing hardware improvements. (Dense) transformers are expected to reach the trillion-parameter scale in the near future, for which training requires thousands or even tens of thousands of GPUs. We investigate the challenges of training at this scale and beyond on commercially available hardware. In particular, we analyse the shortest possible training time for different configurations of distributed training, leveraging empirical scaling laws for language models to estimate the optimal (critical) batch size. Contrary to popular belief, we find no evidence for a memory wall, and instead argue that the real limitation -- other than the cost -- lies in the training duration. In addition to this analysis, we introduce two new methods, \textit{layered gradient accumulation} and \textit{modular pipeline parallelism}, which together cut the shortest training time by half. The methods also reduce data movement, lowering the network requirement to a point where a fast InfiniBand connection is not necessary. This increased network efficiency also improve on the methods introduced with the ZeRO optimizer, reducing the memory usage to a tiny fraction of the available GPU memory.

arXiv comment22 pages, 8 figures

References in corpus (16)

• Language Models are Few-Shot Learners
• Longformer: The Long-Document Transformer
• Scaling Laws for Neural Language Models
• HOGWILD!: A Lock-Free Approach to Parallelizing Stochastic Gradient Descent
• Mixed Precision Training
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• Training Deep Nets with Sublinear Memory Cost
• Generating Long Sequences with Sparse Transformers
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding
• Measuring the Effects of Data Parallelism on Neural Network Training
• An Empirical Model of Large-Batch Training
• On the Computational Inefficiency of Large Batch Sizes for Stochastic Gradient Descent
• Asynchronous stochastic convex optimization
• 1-bit Adam: Communication Efficient Large-Scale Training with Adam's Convergence Speed
• Critical Parameters for Scalable Distributed Learning with Large Batches and Asynchronous Updates