Maximizing Parallelism in Distributed Training for Huge Neural Networks

TL;DR

This work tackles the memory and compute bottlenecks of training very large language models by introducing 3-D intra-layer model parallelism, distributing linear operations across a $P=p^3$ processor cube to achieve near-perfect load balance. It develops a complete set of 3-D primitives, including matrix-matrix and matrix-vector operations, with all-gather and reduce-scatter communications, and applies them to a 3-D Transformer architecture where input/output directions are managed to minimize communication overhead. On a 64-GPU cluster, the approach delivers substantial speedups over prior 1-D and 2-D baselines ($2.32\times$ vs. 1-D and $1.57\times$ vs. 2-D), validating improved weak and strong scaling. These results suggest a viable path to training much larger language models by efficiently leveraging larger GPU clusters while mitigating memory and communication bottlenecks, with future work aimed at scaling to even bigger systems.

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.

Figures (6)

• Figure 1: The structure of a 3-D processing cube with $P=p^3$ processors. The colored block $(i,j,k)$ represents an example of a single processor.
• Figure 2: An example of the 3-D parallel matrix multiplication of size $2\times2\times2$.
• Figure 3: Architecture of the Transformer model that consists of multiple Transformer layers.
• Figure 4: Load-balanced 3-D matrix-matrix multiplication.
• Figure 5: Balancing vector parameters for matrix-vector operations.