A 4D Hybrid Algorithm to Scale Parallel Training to Thousands of GPUs

TL;DR

This paper presents a 4D hybrid communication algorithm for large-scale neural network training that significantly reduces communication overhead and improves performance on thousands of GPUs.

Contribution

It introduces a novel 4D hybrid parallelism approach combined with strategies for overlapping communication and computation, along with an analytical model for efficient configuration tuning.

Findings

AxoNN outperforms Megatron-LM by 26% on 80-billion parameter GPT training.

Achieves 57% of the theoretical peak FLOP/s, totaling 182 PFLOP/s.

Effectively reduces communication bottlenecks in large-scale GPU training.

Abstract

Heavy communication, in particular, collective operations, can become a critical performance bottleneck in scaling the training of billion-parameter neural networks to large-scale parallel systems. This paper introduces a four-dimensional (4D) approach to optimize communication in parallel training. This 4D approach is a hybrid of 3D tensor and data parallelism, and is implemented in the AxoNN framework. In addition, we employ two key strategies to further minimize communication overheads. First, we aggressively overlap expensive collective operations (reduce-scatter, all-gather, and all-reduce) with computation. Second, we develop an analytical model to identify high-performing configurations within the large search space defined by our 4D algorithm. This model empowers practitioners by simplifying the tuning process for their specific training workloads. When training an 80-billion parameter GPT on 1024 GPUs of Perlmutter, AxoNN surpasses Megatron-LM, a state-of-the-art framework, by a significant 26%. Additionally, it achieves a significantly high 57% of the theoretical peak FLOP/s or 182 PFLOP/s in total.