在 Aurora 超级计算机上大规模预训练大型专家混合语言模型
机器学习
2026-04-02 v1 人工智能
分布式、并行与集群计算
摘要
从头预训练大型语言模型（LLM）需要海量计算资源。Aurora 超级计算机是一台拥有 127,488 个 Intel PVC（Ponte Vecchio）GPU 瓦片的 ExaScale 机器。在这项工作中，我们展示了在 Aurora 上以数千 GPU 瓦片规模进行 LLM 预训练。为此，我们开发了 Optimus，一个支持标准大模型训练技术的内部训练库。使用 Optimus，我们首先从头在 3072 个 GPU 瓦片上对 OLMoE-mix-0924 数据集的完整 4 万亿 token 进行了 Mula-1B（10 亿稠密模型）和 Mula-7B-A1B（70 亿专家混合模型）的预训练。然后我们通过在同一数据集上对 1000 亿 token 进行三个大型 MoE 模型 Mula-20B-A2B、Mula-100B-A7B 和 Mula-220B-A10B 的预训练来展示模型扩展。在我们的最大模型 Mula-220B-A10B 上，我们将计算扩展从 384 推进到 12288 个 GPU 瓦片，并在 12288 个 GPU 瓦片上观察到约 90% 的扩展效率。我们通过用于专家计算的自定义 GPU 核和一种新颖的 EP 感知分片优化器显著改善了 MoE 模型的运行时性能，实现了高达 1.71 倍的训练加速。作为 Optimus 库的一部分，我们还开发了一套鲁棒的可靠性和容错特性以提高大规模训练稳定性和连续性。
引用
@article{arxiv.2604.00785,
title = {Scalable Pretraining of Large Mixture of Experts Language Models on Aurora Super Computer},
author = {Dharma Teja Vooturi and Dhiraj Kalamkar and Dipankar Das and Bharat Kaul},
journal= {arXiv preprint arXiv:2604.00785},
year = {2026}
}