mlsanityData to 14 Oct 2026

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

arXiv 2305.13525 ↗First posted 22 May 2023cs.LG cs.AI cs.DCcited by 6

Not on the map. Why?

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.

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

26 of 45 references.

References · 26

  1. 2017Attention Is All You Needcited by 60,361
  2. 2020Language Models are Few-Shot Learnerscited by 32,419
  3. 2015U-Net: Convolutional Networks for Biomedical Image Segmentationcited by 19,713
  4. 2021High-Resolution Image Synthesis with Latent Diffusion Modelscited by 18,863
  5. 2015Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shiftcited by 7,005
  6. 2020Scaling Laws for Neural Language Modelscited by 6,586
  7. 2022Hierarchical Text-Conditional Image Generation with CLIP Latentscited by 6,585
  8. 2016Layer Normalizationcited by 5,669
  9. 2021Improved Denoising Diffusion Probabilistic Modelscited by 3,761
  10. 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,868
  11. 2018Group Normalizationcited by 1,524
  12. 2019ZeRO: Memory Optimizations Toward Training Trillion Parameter Modelscited by 1,301
14 more
  1. 2015Aligning Books and Movies: Towards Story-like Visual Explanations by Watching Movies and Reading Bookscited by 1,043
  2. 2018Reconciling modern machine learning practice and the bias-variance trade-offcited by 927
  3. 2016Training Deep Nets with Sublinear Memory Costcited by 848
  4. 2022GPT-NeoX-20B: An Open-Source Autoregressive Language Modelcited by 713
  5. 2020Image Segmentation Using Deep Learning: A Surveycited by 660
  6. 2022Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Modelcited by 559
  7. 2015Empirical Evaluation of Rectified Activations in Convolutional Networkcited by 519
  8. 2019A Survey of Deep Learning-based Object Detectioncited by 106
  9. 2020Reducing Communication in Graph Neural Network Trainingcited by 50
  10. 2021Breaking the Computation and Communication Abstraction Barrier in Distributed Machine Learning Workloadscited by 40
  11. 2021Maximizing Parallelism in Distributed Training for Huge Neural Networkscited by 23
  12. 2021Tesseract: Parallelize the Tensor Parallelism Efficientlycited by 19
  13. 2022Merak: An Efficient Distributed DNN Training Framework with Automated 3D Parallelism for Giant Foundation Modelscited by 11
  14. 2023ATP: Adaptive Tensor Parallelism for Foundation Modelscited by 3

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2024: 2; 2025: 3; 2026: 1.

6 papers.

  1. 2024Loki: Low-rank Keys for Efficient Sparse Attentioncited by 89
  2. 2024Be like a Goldfish, Don't Memorize! Mitigating Memorization in Generative LLMscited by 50
  3. 2025Gemstones: A Model Suite for Multi-Faceted Scaling Lawscited by 20
  4. 2025Understanding and Improving Communication Performance in Multi-node LLM Inferencecited by 2
  5. 2025Power Law Guided Dynamic Sifting for Efficient Attentioncited by 1
  6. 2026The Life of a Token: from Words to Bits on the Wirecited by 0