Aurora 超级计算机上大型混合专家语言模型的可扩展预训练

从头开始预训练 大语言模型 (LLM) 需要大量计算。 Aurora 超级计算机是一台 ExaScale 机器，拥有 127,488 个 Intel PVC (Ponte Vechio) GPU 块。在这项工作中，我们展示了 Aurora 上 1000 个 GPU 块规模的 LLM 预训练。为了实现这一目标，我们开发了 Optimus，这是一个内部训练库，支持标准大型 模型训练 技术。使用 Optimus，我们首先在 3072 个 GPU 块上从头开始对 Mula-1B（一个 10 亿密集模型）和 Mula-7B-A1B（一个 70 亿混合专家 (MoE) 模型）进行预训练，以获取 OLMoE-mix-0924 数据集的全部 4 万亿个标记。然后，我们通过在同一数据集上预训练三个大型 MoE 模型 Mula-20B-A2B、Mula-100B-A7B 和 Mula-220B-A10B 直至 1000 亿个词元来演示模型扩展。在我们最大的模型 Mula-220B-A10B 上，我们将计算扩展从 384 个 GPU 块增加到 12288 个 GPU 块，并观察到在 12288 个 GPU 块时扩展效率约为 90%。我们使用用于专家计算的自定义 GPU 内核和新颖的 EP-Aware 分片优化器显着提高了 MoE 模型的运行时性能，从而使训练速度提高了 1.71 倍。作为 Optimus 库的一部分，我们还开发了一组强大的可靠性和容错功能，以大规模提高训练稳定性和连续性。