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

📄 arXiv: 2502.08145v1 📥 PDF

作者: Siddharth Singh, Prajwal Singhania, Aditya Ranjan, John Kirchenbauer, Jonas Geiping, Yuxin Wen, Neel Jain, Abhimanyu Hans, Manli Shu, Aditya Tomar, Tom Goldstein, Abhinav Bhatele

分类: cs.LG, cs.AI, cs.DC

发布日期: 2025-02-12

💡 一句话要点

AxoNN：开源可扩展LLM训练框架，实现GPU超算上的高效训练。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 分布式训练 GPU超算 四维混合并行 开源框架

📋 核心要点

  1. 现有LLM训练需要大量GPU资源和高度可扩展的软件栈，但开源方案不足。
  2. AxoNN框架采用四维混合并行算法，优化矩阵乘法、集合通信等，提升扩展性。
  3. 实验表明，AxoNN在多个超算平台上实现了卓越的性能和扩展性，并成功微调了4050亿参数的LLM。

📝 摘要（中文）

本文介绍了一个名为AxoNN的开源、可扩展框架，用于训练和微调具有数千亿到数万亿参数的大型语言模型（LLM）。AxoNN采用了一种新颖的四维混合并行算法，并在多个方面进行了性能优化，包括改进矩阵乘法内核性能、重叠非阻塞集合通信与计算，以及利用性能建模选择最优配置。这些优化使得AxoNN在Perlmutter（620.1 Petaflop/s）、Frontier（1.381 Exaflop/s）和Alps（1.423 Exaflop/s）上实现了前所未有的扩展性和峰值浮点运算性能（bf16），用于训练GPT风格的Transformer模型。此外，本文还探讨了模型规模增大带来的“灾难性记忆”问题，并提出了一种预防方法。最后，论文展示了使用AxoNN在Frontier上微调一个4050亿参数的LLM。

🔬 方法详解

问题定义：当前训练具有数千亿甚至数万亿参数的LLM，需要消耗大量的计算资源，依赖于大规模的GPU集群。现有的开源框架在扩展性和性能方面存在瓶颈，难以充分利用GPU超算的强大算力。此外，模型规模的增大也带来了新的问题，例如“灾难性记忆”，即模型能够记住训练数据中的敏感信息，导致隐私泄露风险。

核心思路：AxoNN的核心思路是设计一个高度可扩展、高性能的开源框架，能够充分利用GPU超算的算力来训练和微调LLM。通过采用四维混合并行算法，AxoNN能够有效地将计算任务分配到多个GPU上，并优化通信开销。此外，AxoNN还关注模型训练过程中的隐私问题，并提出了一种预防“灾难性记忆”的方法。

技术框架：AxoNN的整体架构包含以下几个主要模块：数据并行、张量并行、流水线并行和序列并行。这些并行策略共同作用，实现了四维混合并行。框架还包括性能建模模块，用于选择最优的配置参数，以最大化训练效率。此外，AxoNN还集成了优化后的矩阵乘法内核和非阻塞集合通信机制，进一步提升了性能。

关键创新：AxoNN最重要的技术创新点在于其四维混合并行算法。与传统的并行策略相比，四维混合并行能够更好地利用GPU集群的资源，实现更高的扩展性和性能。此外，AxoNN还针对GPU超算的特点进行了优化，例如优化矩阵乘法内核和重叠通信与计算，从而进一步提升了训练效率。

关键设计：AxoNN的关键设计包括：1) 四维混合并行策略的具体实现，包括如何将数据、张量、流水线和序列进行划分和分配；2) 性能建模模块的算法，用于预测不同配置下的训练性能，并选择最优参数；3) 优化后的矩阵乘法内核，例如使用更高效的算法和数据布局；4) 非阻塞集合通信机制，用于减少通信开销。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

AxoNN在Perlmutter、Frontier和Alps等GPU超算上实现了卓越的性能和扩展性。在Perlmutter上，AxoNN达到了620.1 Petaflop/s的峰值浮点运算性能（bf16）。在Frontier上，AxoNN达到了1.381 Exaflop/s的峰值浮点运算性能（bf16），并在Alps上达到了1.423 Exaflop/s。此外，AxoNN还成功地在Frontier上微调了一个4050亿参数的LLM。

🎯 应用场景

AxoNN框架可应用于各种需要大规模LLM训练和微调的场景，例如自然语言处理、机器翻译、文本生成、对话系统等。该框架的开源特性降低了LLM研究和应用的门槛，促进了AI技术的民主化。未来，AxoNN有望成为LLM训练的重要基础设施，推动AI技术的进一步发展。

📄 摘要（原文）

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.