# AxoNN: An asynchronous, message-driven parallel framework for extreme-scale deep learning

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

## Key facts

- Cited by 15 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 30 references that are in the map.
- First posted on arXiv 25 Oct 2021.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Decoupled Weight Decay Regularization](https://mlsanity.com/paper/1711.05101/) (2017)
- [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [Horovod: fast and easy distributed deep learning in TensorFlow](https://mlsanity.com/paper/1802.05799/) (2018)
- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://mlsanity.com/paper/2104.07857/) (2021)
- [Mesh-TensorFlow: Deep Learning for Supercomputers](https://mlsanity.com/paper/1811.02084/) (2018)
- [Revisiting Distributed Synchronous SGD](https://mlsanity.com/paper/1702.05800/) (2017)
- [Poseidon: An Efficient Communication Architecture for Distributed Deep Learning on GPU Clusters](https://mlsanity.com/paper/1706.03292/) (2017)
- [PipeMare: Asynchronous Pipeline Parallel DNN Training](https://mlsanity.com/paper/1910.05124/) (2019)
- [Automatic Graph Partitioning for Very Large-scale Deep Learning](https://mlsanity.com/paper/2103.16063/) (2021)

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

- [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)
- [A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training](https://mlsanity.com/paper/2303.06318/) (2023)
- [Gemstones: A Model Suite for Multi-Faceted Scaling Laws](https://mlsanity.com/paper/2502.06857/) (2025)
- [A Survey and Empirical Evaluation of Parallel Deep Learning Frameworks](https://mlsanity.com/paper/2111.04949/) (2021)
- [Democratizing AI: Open-source Scalable LLM Training on GPU-based Supercomputers](https://mlsanity.com/paper/2502.08145/) (2025)
- [Exploiting Sparsity in Pruned Neural Networks to Optimize Large Model Training](https://mlsanity.com/paper/2302.05045/) (2023)
- [Understanding and Improving Communication Performance in Multi-node LLM Inference](https://mlsanity.com/paper/2511.09557/) (2025)
- [FAST: Factorizable Attention for Speeding up Transformers](https://mlsanity.com/paper/2402.07901/) (2024)
- [Plexus: Taming Billion-edge Graphs with 3D Parallel Full-graph GNN Training](https://mlsanity.com/paper/2505.04083/) (2025)
- [Power Law Guided Dynamic Sifting for Efficient Attention](https://mlsanity.com/paper/2506.05300/) (2025)
- [MPipeMoE: Memory Efficient MoE for Pre-trained Models with Adaptive Pipeline Parallelism](https://mlsanity.com/paper/2506.22175/) (2025)
- [The Big Send-off: Scalable and Performant Collectives for Deep Learning](https://mlsanity.com/paper/2504.18658/) (2025)
- [Communication-free Sampling and 4D Hybrid Parallelism for Scalable Mini-batch GNN Training](https://mlsanity.com/paper/2604.02651/) (2026)
- [Transforming the Use of Earth Observation Data: Exascale Training of a Generative Compression Model with Historical Priors for up to 10,000x Data Reduction](https://mlsanity.com/paper/2605.08633/) (2026)

## Abstract (arXiv)

In the last few years, the memory requirements to train state-of-the-art neural networks have far exceeded the DRAM capacities of modern hardware accelerators. This has necessitated the development of efficient algorithms to train these neural networks in parallel on large-scale GPU-based clusters. Since computation is relatively inexpensive on modern GPUs, designing and implementing extremely efficient communication in these parallel training algorithms is critical for extracting the maximum performance. This paper presents AxoNN, a parallel deep learning framework that exploits asynchrony and message-driven execution to schedule neural network operations on each GPU, thereby reducing GPU idle time and maximizing hardware efficiency. By using the CPU memory as a scratch space for offloading data periodically during training, AxoNN is able to reduce GPU memory consumption by four times. This allows us to increase the number of parameters per GPU by four times, thus reducing the amount of communication and increasing performance by over 13%. When tested against large transformer models with 12-100 billion parameters on 48-384 NVIDIA Tesla V100 GPUs, AxoNN achieves a per-GPU throughput of 49.4-54.78% of theoretical peak and reduces the training time by 22-37 days (15-25% speedup) as compared to the state-of-the-art.