Democratizing AI: Open-source Scalable LLM Training on GPU-based Supercomputers
Fuente:
arXiv
Gespeichert in:
| Hauptverfasser: | , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Veröffentlicht: |
2025
|
| Schlagworte: | |
| Online-Zugang: | |
| Tags: |
Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
|
| _version_ | 1866929711776530432 |
|---|---|
| author | Singh, Siddharth Singhania, Prajwal Ranjan, Aditya Kirchenbauer, John Geiping, Jonas Wen, Yuxin Jain, Neel Hans, Abhimanyu Shu, Manli Tomar, Aditya Goldstein, Tom Bhatele, Abhinav |
| author_facet | Singh, Siddharth Singhania, Prajwal Ranjan, Aditya Kirchenbauer, John Geiping, Jonas Wen, Yuxin Jain, Neel Hans, Abhimanyu Shu, Manli Tomar, Aditya Goldstein, Tom Bhatele, Abhinav |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2502_08145 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Democratizing AI: Open-source Scalable LLM Training on GPU-based Supercomputers Singh, Siddharth Singhania, Prajwal Ranjan, Aditya Kirchenbauer, John Geiping, Jonas Wen, Yuxin Jain, Neel Hans, Abhimanyu Shu, Manli Tomar, Aditya Goldstein, Tom Bhatele, Abhinav Machine Learning Artificial Intelligence Distributed, Parallel, and Cluster Computing 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. |
| title | Democratizing AI: Open-source Scalable LLM Training on GPU-based Supercomputers |
| topic | Machine Learning Artificial Intelligence Distributed, Parallel, and Cluster Computing |
| url | https://arxiv.org/abs/2502.08145 |