# Democratizing AI: Open-source Scalable LLM Training on GPU-based Supercomputers

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

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 50 references that are in the map.
- First posted on arXiv 12 Feb 2025.

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

- [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)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Extracting Training Data from Large Language Models](https://mlsanity.com/paper/2012.07805/) (2020)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling](https://mlsanity.com/paper/2304.01373/) (2023)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [Quantifying Memorization Across Neural Language Models](https://mlsanity.com/paper/2202.07646/) (2022)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [TinyLlama: An Open-Source Small Language Model](https://mlsanity.com/paper/2401.02385/) (2024)
- [Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model](https://mlsanity.com/paper/2201.11990/) (2022)
- [A Study of BFLOAT16 for Deep Learning Training](https://mlsanity.com/paper/1905.12322/) (2019)
- [MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs](https://mlsanity.com/paper/2402.15627/) (2024)
- [Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training](https://mlsanity.com/paper/2110.14883/) (2021)
- [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)
- [AxoNN: An asynchronous, message-driven parallel framework for extreme-scale deep learning](https://mlsanity.com/paper/2110.13005/) (2021)
- [Optimizing Distributed Training on Frontier for Large Language Models](https://mlsanity.com/paper/2312.12705/) (2023)
- [Scaling Distributed Deep Learning Workloads beyond the Memory Capacity with KARMA](https://mlsanity.com/paper/2008.11421/) (2020)
- [Exploiting Sparsity in Pruned Neural Networks to Optimize Large Model Training](https://mlsanity.com/paper/2302.05045/) (2023)

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

- [Characterizing the Efficiency of Distributed Training: A Power, Performance, and Thermal Perspective](https://mlsanity.com/paper/2509.10371/) (2025)
- [Understanding and Improving Communication Performance in Multi-node LLM Inference](https://mlsanity.com/paper/2511.09557/) (2025)
- [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)
- [The Big Send-off: Scalable and Performant Collectives for Deep Learning](https://mlsanity.com/paper/2504.18658/) (2025)
- [FictionalQA: A Dataset for Studying Memorization and Knowledge Acquisition](https://mlsanity.com/paper/2506.05639/) (2025)
- [Communication-free Sampling and 4D Hybrid Parallelism for Scalable Mini-batch GNN Training](https://mlsanity.com/paper/2604.02651/) (2026)

## Abstract (arXiv)

Training and fine-tuning large language models (LLMs) with hundreds of billions to trillions of parameters requires tens of thousands of GPUs, and a highly scalable software stack. In this work, we present a novel four-dimensional hybrid parallel algorithm implemented in a highly scalable, portable, open-source framework called AxoNN. We describe several performance optimizations in AxoNN to improve matrix multiply kernel performance, overlap non-blocking collectives with computation, and performance modeling to choose performance optimal configurations. These have resulted in unprecedented scaling and peak flop/s (bf16) for training of GPT-style transformer models on Perlmutter (620.1 Petaflop/s), Frontier (1.381 Exaflop/s) and Alps (1.423 Exaflop/s). While the abilities of LLMs improve with the number of trainable parameters, so do privacy and copyright risks caused by memorization of training data, which can cause disclosure of sensitive or private information at inference time. We highlight this side effect of scale through experiments that explore "catastrophic memorization", where models are sufficiently large to memorize training data in a single pass, and present an approach to prevent it. As part of this study, we demonstrate fine-tuning of a 405-billion parameter LLM using AxoNN on Frontier.