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

arXiv: https://arxiv.org/abs/2305.13525 · HTML: https://mlsanity.com/paper/2305.13525/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 45 references that are in the map.
- First posted on arXiv 22 May 2023.

## Builds on (most cited first, 26 of 26)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Hierarchical Text-Conditional Image Generation with CLIP Latents](https://mlsanity.com/paper/2204.06125/) (2022)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Group Normalization](https://mlsanity.com/paper/1803.08494/) (2018)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Aligning Books and Movies: Towards Story-like Visual Explanations by Watching Movies and Reading Books](https://mlsanity.com/paper/1506.06724/) (2015)
- [Reconciling modern machine learning practice and the bias-variance trade-off](https://mlsanity.com/paper/1812.11118/) (2018)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [GPT-NeoX-20B: An Open-Source Autoregressive Language Model](https://mlsanity.com/paper/2204.06745/) (2022)
- [Image Segmentation Using Deep Learning: A Survey](https://mlsanity.com/paper/2001.05566/) (2020)
- [Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model](https://mlsanity.com/paper/2201.11990/) (2022)
- [Empirical Evaluation of Rectified Activations in Convolutional Network](https://mlsanity.com/paper/1505.00853/) (2015)
- [A Survey of Deep Learning-based Object Detection](https://mlsanity.com/paper/1907.09408/) (2019)
- [Reducing Communication in Graph Neural Network Training](https://mlsanity.com/paper/2005.03300/) (2020)
- [Breaking the Computation and Communication Abstraction Barrier in Distributed Machine Learning Workloads](https://mlsanity.com/paper/2105.05720/) (2021)
- [Maximizing Parallelism in Distributed Training for Huge Neural Networks](https://mlsanity.com/paper/2105.14450/) (2021)
- [Tesseract: Parallelize the Tensor Parallelism Efficiently](https://mlsanity.com/paper/2105.14500/) (2021)
- [Merak: An Efficient Distributed DNN Training Framework with Automated 3D Parallelism for Giant Foundation Models](https://mlsanity.com/paper/2206.04959/) (2022)
- [ATP: Adaptive Tensor Parallelism for Foundation Models](https://mlsanity.com/paper/2301.08658/) (2023)

## Built on by (most cited first, 6 of 6)

- [Loki: Low-rank Keys for Efficient Sparse Attention](https://mlsanity.com/paper/2406.02542/) (2024)
- [Be like a Goldfish, Don't Memorize! Mitigating Memorization in Generative LLMs](https://mlsanity.com/paper/2406.10209/) (2024)
- [Gemstones: A Model Suite for Multi-Faceted Scaling Laws](https://mlsanity.com/paper/2502.06857/) (2025)
- [Understanding and Improving Communication Performance in Multi-node LLM Inference](https://mlsanity.com/paper/2511.09557/) (2025)
- [Power Law Guided Dynamic Sifting for Efficient Attention](https://mlsanity.com/paper/2506.05300/) (2025)
- [The Life of a Token: from Words to Bits on the Wire](https://mlsanity.com/paper/2609.19924/) (2026)

## Abstract (arXiv)

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.