开源框架AxoNN实现千亿级大模型在超算上的高效训练,解决训练规模带来的隐私泄露问题。
Democratizing AI: Open-source Scalable LLM Training on GPU-based Supercomputers
- 提出四维混合并行算法,优化计算与通信重叠,提升GPU超算性能
- 在Frontier超算上实现1.381 Exaflop/s峰值算力,成功训练4050亿参数模型
- 揭示大模型单次训练就可能完全记忆数据,提出防止灾难性记忆的新方法
训练拥有数百亿至万亿参数的大语言模型需要数万张GPU及高度可扩展的软件栈。本文提出一种新型四维混合并行算法,集成于可移植、开源的AxoNN框架中。通过优化矩阵乘法内核、非阻塞集合通信与计算重叠、性能建模选择最优配置,实现了在Perlmutter(620.1 Petaflop/s)、Frontier(1.381 Exaflop/s)和Alps(1.423 Exaflop/s)上的空前可扩展性与峰值浮点性能。尽管模型参数量越大能力越强,但也会引发隐私与版权风险——即训练数据被模型记忆,在推理时可能泄露敏感信息。我们通过实验揭示‘灾难性记忆’现象:模型足够大时可在单次训练中完全记忆数据,并提出防范策略。作为研究的一部分,我们使用AxoNN在Frontier上成功完成了4050亿参数大模型的微调。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。