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

TL;DR

This work presents AxoNN, a scalable, open-source framework implementing a four-dimensional hybrid parallelism to train GPT-style LLMs across thousands of GPUs. By combining data parallelism with a 3D tensor-parallel matrix-multiplication scheme and a data-grouped 4D grid, the approach achieves unprecedented weak scaling and sustained bf16 flop/s on diverse hardware (NVIDIA and AMD) and assesses a broad range of model sizes. Key innovations include a predictive performance model for near-optimal configurations, automated BLAS kernel tuning, and overlapping asynchronous collectives with computation, yielding substantial speedups and reduced time-to-solution for state-of-the-art models. Beyond performance, the work investigates memorization risks in large LLMs, demonstrating catastrophic memorization at large scales and showing mitigation via Goldfish Loss, highlighting both the opportunities and privacy considerations of scalable open training. The combination of AxoNN’s scalable training, experimental measurement across leading HPC systems, and memorization analyses provides a practical path toward faster training cycles and safer deployment of extremely large language models on commodity HPC resources.

Abstract

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.

Figures (11)

• Figure 1: Parallelization of a matrix multiply in an FC layer with Agarwal's 3D parallel matrix multiplication algorithm agarwal-3d on eight GPUs organized in a $2\times2\times2$ topology. We use $G_x$, $G_y$, and $G_z$ to refer to the number of GPUs along the three dimensions of the virtual grid topology.
• Figure 2: Plots validating the performance model by comparing the observed time per batch and the rank ordered by the model for two neural networks: GPT-20B (left) and GPT-40B (right).
• Figure 3: Creation of a ring among eight GPUs on two nodes for a collective communication operation (all-reduce/reduce-scatter/all-gather).
• Figure 4: Two rings among four GPUs each across two nodes for performing collective operations simultaneously.
• Figure 5: The impact of overlapping non-blocking collectives with computation on the training times of different sized models on 8,192 GCDs of Frontier.