  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.
