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.

翻译：重通信（特别是集合通信操作）在将十亿参数神经网络的训练扩展至大规模并行系统时，可能成为关键的性能瓶颈。本文提出一种四维（4D）方法以优化并行训练中的通信。该四维方法是三维张量并行与数据并行的混合，并在AxoNN框架中实现。此外，我们采用两项关键策略进一步降低通信开销：第一，将高开销的集合操作（reduce-scatter、all-gather及all-reduce）与计算进行激进重叠；第二，开发分析模型以识别四维算法定义的大搜索空间中的高性能配置。该模型通过简化特定训练负载的调优流程赋能实践者。在Perlmutter系统的1024个GPU上训练800亿参数的GPT时，AxoNN以26%的显著优势超越先进框架Megatron-LM，同时实现了高达理论峰值FLOP/s的57%（即总计算量182 PFLOP/s）。