Scalable Pretraining of Large Mixture of Experts Language Models on Aurora Super Computer

📄 arXiv: 2604.00785v1 📥 PDF

作者: Dharma Teja Vooturi, Dhiraj Kalamkar, Dipankar Das, Bharat Kaul

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

发布日期: 2026-04-01

💡 一句话要点

在Aurora超算上预训练大规模混合专家语言模型，实现高效扩展。

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

关键词: 大规模语言模型 混合专家模型 预训练 超级计算机 GPU加速 分布式训练 模型扩展 优化器

📋 核心要点

  1. 大规模语言模型预训练需要巨大的计算资源，现有方法在大规模GPU集群上的扩展性和效率面临挑战。
  2. 本文提出Optimus训练库，并优化了MoE模型的专家计算和优化器，以提高训练速度和扩展性。
  3. 实验表明，在Aurora超算上，Mula-220B-A10B模型在12288个GPU tiles上实现了90%的扩展效率，训练速度提升高达1.71倍。

📝 摘要（中文）

本文展示了在Aurora超级计算机上进行大规模语言模型（LLM）预训练的工作，Aurora拥有127,488个Intel PVC（Ponte Vechio）GPU tiles。我们利用自研训练库Optimus，支持标准的大模型训练技术。首先，在3072个GPU tiles上，使用OLMoE-mix-0924数据集的4万亿tokens，从头预训练了10亿参数的稠密模型Mula-1B和70亿参数的混合专家(MoE)模型Mula-7B-A1B。然后，通过在相同数据集上预训练Mula-20B-A2B、Mula-100B-A7B和Mula-220B-A10B三个大型MoE模型至1000亿tokens，展示了模型扩展能力。在最大的模型Mula-220B-A10B上，我们将计算规模从384扩展到12288个GPU tiles，并在12288个GPU tiles上观察到约90%的扩展效率。我们使用定制的GPU kernels进行专家计算，以及一种新颖的EP-Aware分片优化器，显著提高了MoE模型的运行时性能，训练速度提升高达1.71倍。作为Optimus库的一部分，我们还开发了一套强大的可靠性和容错功能，以提高大规模训练的稳定性和连续性。

🔬 方法详解

问题定义：本文旨在解决大规模混合专家（MoE）语言模型在ExaScale超级计算机上的高效预训练问题。现有方法在扩展到数千个GPU时，常常面临通信瓶颈、专家负载不均衡以及优化器效率低下的问题，导致训练速度慢、资源利用率低。

核心思路：本文的核心思路是通过优化MoE模型的计算和通信，并设计高效的分片优化器，从而提高大规模预训练的扩展性和效率。具体而言，通过定制GPU kernels加速专家计算，并提出EP-Aware分片优化器来减少通信开销。

技术框架：整体框架基于Optimus训练库，支持标准的大模型训练技术。主要流程包括：数据加载和预处理、模型初始化、前向传播（包括专家选择和计算）、反向传播、优化器更新。针对MoE模型，特别优化了专家计算和通信过程。

关键创新：本文的关键创新在于：1) 针对MoE模型的定制GPU kernels，显著加速了专家计算；2) 提出的EP-Aware分片优化器，通过感知专家并行（Expert Parallelism）来减少优化器状态的通信开销。与现有方法相比，这些创新能够更有效地利用大规模GPU集群的计算资源。

关键设计：在专家计算方面，采用了优化的矩阵乘法和激活函数实现。EP-Aware分片优化器将优化器状态按照专家进行分片，并仅在需要时进行通信。具体参数设置和损失函数选择与标准LLM预训练方法一致，网络结构则采用了MoE架构，具体专家数量和容量根据模型大小进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，在Aurora超算上，Mula-220B-A10B模型在12288个GPU tiles上实现了约90%的扩展效率。通过定制GPU kernels和EP-Aware分片优化器，MoE模型的训练速度提升高达1.71倍。这些结果验证了本文提出的方法在大规模MoE模型预训练方面的有效性。

🎯 应用场景

该研究成果可应用于各种需要大规模语言模型支持的领域，如自然语言处理、机器翻译、文本生成、对话系统等。通过高效地预训练更大规模的MoE模型，可以提升这些应用的性能和效果，并为未来的AI发展奠定基础。此外，Optimus训练库的开发也为其他研究者提供了有价值的工具。

📄 摘要（原文）

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.