在1万2千块GPU上训练超大混合专家模型,效率达90%。
Scalable Pretraining of Large Mixture of Experts Language Models on Aurora Super Computer
- 自研训练库Optimus支持大规模模型训练,适配混合专家架构。
- 在12288块GPU上训练2200亿参数模型,效率达90%,速度提升1.71倍。
- 适合关注超大规模语言模型训练与系统优化的研究者。
从头预训练大型语言模型需要巨大算力。Aurora超级计算机是拥有127,488块Intel PVC GPU的埃克萨级系统。本文展示了在数千块GPU上进行大模型预训练的能力。为此,我们开发了Optimus——一个支持主流大模型训练技术的内部训练库。利用Optimus,我们在3072块GPU上从头训练了Mula-1B(10亿参数稠密模型)和Mula-7B-A1B(70亿参数混合专家模型),覆盖4万亿标记的OLMoE-mix-0924数据集。随后,我们进一步扩展至三个大型MoE模型:Mula-20B-A2B、Mula-100B-A7B和Mula-220B-A10B,均在相同数据集上训练至1000亿标记。在最大模型Mula-220B-A10B中,计算规模从384扩展到12288块GPU,12288块时实现约90%的缩放效率。通过定制的GPU专家计算内核和新型EP-Aware分片优化器,显著提升MoE模型运行效率,训练速度最高提升1.71倍。作为Optimus的一部分,我们还构建了完善的可靠性与容错机制,保障大规模训练的稳定性和连续性。
原文摘要 · Abstract (English)
Pretraining Large Language Models (LLMs) from scratch requires massive amount of compute. Aurora super computer is an ExaScale machine with 127,488 Intel PVC (Ponte Vechio) GPU tiles. In this work, we showcase LLM pretraining on Aurora at the scale of 1000s of GPU tiles. Towards this effort, we developed Optimus, an inhouse training library with support for standard large model training techniques. Using Optimus, we first pretrained Mula-1B, a 1 Billion dense model and Mula-7B-A1B, a 7 Billion Mixture of Experts (MoE) model from scratch on 3072 GPU tiles for the full 4 trillion tokens of the OLMoE-mix-0924 dataset. We then demonstrated model scaling by pretraining three large MoE models Mula-20B-A2B, Mula-100B-A7B, and Mula-220B-A10B till 100 Billion tokens on the same dataset. On our largest model Mula-220B-A10B, we pushed the compute scaling from 384 to 12288 GPU tiles and observed scaling efficiency of around 90% at 12288 GPU tiles. We significantly improved the runtime performance of MoE models using custom GPU kernels for expert computation, and a novel EP-Aware sharded optimizer resulting in training speedups up to 1.71x. As part of the Optimus library, we also developed a robust set of reliability and fault tolerant features to improve training stability and continuity at scale.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。