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

作者/机构: Dharma Teja Vooturi, Dhiraj Kalamkar, Dipankar Das, Bharat Kaul Parallel Computing Lab (India) Intel Corporation

速读

一句话结论 本文基于 Aurora 超级计算机的数千个 Intel PVC GPU 瓦片开发了 Optimus 训练库，通过定制化算子与分片优化策略，成功实现了最高 220B 参数混合专家（MoE）语言模型在万亿 token 级别的大规模预训练，并在万卡规模下保持了约 90% 的扩展效率。

要解决什么问题 训练大规模模型面临着严峻的显存容量与并行效率双重瓶颈。使用 AdamW 以 BF16 训练 $P$ 个参数的模型至少需要 $16P$ 字节显存（含权重、梯度及 FP32 状态），导致单张 64 GB 显存的 GPU 连 7B 模型都无法装载。为打破显存限制，业界通常叠加数据并行（DP）、分片优化器（SO）和专家并行（EP）等技术，但这在 MoE 模型中引入了新卡点。一方面，MoE 路由机制导致非均匀的 token 通信，使标准 Hugging Face 实现在 GPU 上存在大量碎片化矩阵乘法，计算极为低效。另一方面，当在 EP 之上应用标准分片优化器时，由于标准优化器仅在 DP 维度切分状态，而 EP 会在各 rank 间复制非专家参数（如注意力层），导致非专家参数的优化器状态被毫无必要地冗余复制了 EP 次，极大浪费显存。此外，万卡集群硬件故障率极高，传统数据加载和单点检查点极易压垮文件系统，导致训练频繁中断。

怎么做的 核心思路是针对 MoE 架构的路由特征重构底层 GPU 算子，结合参数复制语义重新设计优化器分片策略，并在基础设施层引入多重容错机制。关键设计由三个部件构成。一是 FastSparseMoE 模块，通过五阶段 GPU kernel 优化替换低效实现：通信阶段放弃不规则的 all2all，改用更高效的 allgather 和 reducescatter；计数与索引阶段利用前缀和在本地计算路由边界；专家计算阶段将同一 rank 上所有专家权重合并，利用分组矩阵乘法（Grouped mm）一次性完成计算，消除碎片化开销。二是 EP 感知的分片优化器（EPSO）。它将模型参数划分为专家参数 $P_E$ 和非专家参数 $P_{NE}$，采用二维分片策略：$P_E$ 的优化器状态仅在 DP 维度分片，而 $P_{NE}$ 的状态在 DP 和 EP 两个维度同时分片，彻底消除冗余复制。三是高可靠容错基础设施。包含将数据预处理为 numpy 数组支持 mmap 懒加载、通过广播机制加载模型避免文件系统挂起，以及设计双检查点（交替保留最近两个以防写入崩溃）和 DP 分散式检查点（将写入负载均匀分散到各节点），确保软硬件故障时能自动剔除坏节点并快速重启。

效果如何 实验在 Aurora 超级计算机上搭建，硬件规模从 384 个扩展至 12288 个 Intel PVC GPU 瓦片，使用包含 4 万亿 token 的 OLMoE-Mix-0924 数据集进行预训练。对比基线包含两个明确对象：同等计算量下的 1B 稠密模型 Mula-1B（代表传统稠密扩展路线），以及官方的 allenai/OLMoE-1B-7B-0924 模型（代表开源 SOTA MoE 路线）。量化结果显示，在等计算量设置下，7B MoE 模型仅用 1/8 的训练数据（约 500B tokens）就达到了 1B 稠密模型训练 4 万亿 token 的最终准确率，平均准确率高出 6.7%，在困难的 MMLU 基准上更是打破了稠密模型的随机表现瓶颈。在算子性能上，FastSparseMoE 为前向和后向计算带来最高 2.83 倍加速，EPSO 为优化器步骤带来最高 1.36 倍加速，两者结合使端到端训练速度提升最高 1.71 倍。在极限计算规模扩展测试中，将结合了 8 路流水线并行和 12 路专家并行的 220B 超大 MoE 模型从 384 个 GPU 瓦片扩展到 12288 个时，依然保持了约 90% 的扩展效率。作者也坦承了规模扩展的局限性：当扩展到超过一千个 GPU 瓦片时，扩展效率会出现约 10% 的下降，但通过强制均匀路由（FUR）的消融实验证明，这种效率下降是超大规模分布式系统的固有损耗，而非 MoE 模型专家选择不均匀所导致。

A1 主要贡献

本文展示了在拥有127,488个Intel PVC (Ponte Vecchio) GPU瓦片（tiles）的ExaScale级别机器——Aurora超级计算机【1, Aurora: Architecting Argonne’s First Exascale Supercomputer for Accelerated Scientific Discovery, 2025, arXiv】上，进行千级别GPU瓦片规模的大型语言模型（LLM）预训练。为此，研究团队开发了一个名为Optimus的内部训练库，该库支持标准的大模型训练技术。

核心研究与贡献如下：
1. 模型预训练与验证：使用Optimus库，在3072个GPU瓦片上，利用OLMoE-mix-0924数据集的全部4万亿个token，从头预训练了一个10亿参数的稠密模型Mula-1B和一个70亿参数的混合专家（MoE）模型Mula-7B-A1B。
2. 模型规模扩展：在同一数据集上，通过预训练三个更大的MoE模型——Mula-20B-A2B、Mula-100B-A7B和Mula-220B-A10B，直到1000亿个token，成功展示了模型规模的扩展能力。
3. 计算规模扩展：在最大的模型Mula-220B-A10B上，将计算规模从384个GPU瓦片扩展到12288个GPU瓦片，并在12288个GPU瓦片规模下观察到约90%的扩展效率。
4. 性能优化：通过为专家计算定制的GPU核（custom GPU kernels）和一个新颖的EP感知分片优化器（EP-Aware sharded optimizer），显著提升了MoE模型的运行时性能，实现了高达1.71倍的训练加速。
5. 可靠性与容错：作为Optimus库的一部分，开发了一套强大的可靠性和容错功能，以提高大规模训练的稳定性和连续性。

A3 背景知识

训练大型模型

大规模模型训练的内存挑战。使用AdamW优化器以BF16混合精度训练一个拥有P个参数的模型，至少需要16P字节的内存，其中包括：2P用于权重，2P用于梯度，4P用于FP32主权重，以及8P用于FP32优化器状态。因此，在单个PVC GPU瓦片上训练一个70亿参数的小模型是不可行的，因为其112 GB（16 × 7）的内存需求超过了PVC GPU瓦片64 GB的可用内存容量。为了克服这一限制并实现大型模型训练，学术界提出了多种技术，例如分片数据并行（sharded data parallelism，如ZeRO 【2, ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, 2020, SC】和FSDP 【3, PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel, 2023, arXiv】）、模型并行（张量并行【4, Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, 2019, arXiv】、专家并行【5, GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding, 2020, arXiv】和流水线并行【6, GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, 2019, Advances in Neural Information Processing Systems】）、上下文并行【7, Ring Attention with Blockwise Transformers for Near-Infinite Context, 2024, ICLR】、激活检查点【8, Training Deep Nets with Sublinear Memory Cost, 2016, arXiv】等。本节将介绍我们在Optimus训练库中采用的部分技术及其具体实现。

分片优化器（Sharded Optimizer, SO）。在PyTorch DDP中，优化器状态在每个数据并行（DP）rank上是完全复制的。在后向传播之后，梯度通过all-reduce操作在所有rank间同步。每个rank在本地拥有完整的梯度后，独立地更新其模型权重的副本。相比之下，当使用带有分片优化器的DP时，优化器状态是在DP rank之间进行分区（分片）而不是复制。因此，使用reduce-scatter操作来同步梯度，只将相关的梯度分片分发给每个rank——具体来说，是该rank所拥有的优化器状态对应的参数的梯度。每个rank只更新其分配到的参数分片。最后，使用all-gather操作来共享更新后的参数分片，以便所有DP rank再次拥有一个一致且完整的模型副本。对优化器状态进行分片使我们能够仅用DP就训练小型MoE模型。

张量并行（Tensor Parallelism, TP）。在TP中，模型被水平切分。一个稠密Transformer模型中的解码器块包含两个模块：注意力（attention）和多层感知机（MLP）。在TP中，注意力模块中的注意力头被划分，MLP模块中的中间层大小被划分。为了累积部分输出激活和输入梯度，需要在前向/后向传播中，在注意力层和MLP层之后/之前对激活/梯度执行all-reduce操作。除了减少GPU的内存需求外，TP还允许我们在不增加全局批量大小的情况下，将训练扩展到更多的GPU上。

专家并行（Expert Parallelism, EP）。混合专家（MoE）模型中的解码器块包含两个模块：注意力和稀疏MoE（SparseMoE）。SparseMoE模块包含多个专家和一个路由器（router）。在专家并行（EP）中，专家被划分到各个EP rank上，而非专家参数则在每个rank上被复制。由于专家参数占了模型权重的大部分，分发它们使得我们可以训练中等规模的MoE模型。在EP中，全局批量大小会像DP一样随着GPU数量的增加而扩展。

流水线并行（Pipeline Parallelism, PP）。在流水线并行（PP）中，模型按层进行垂直切分。每个rank处理其对应层块的前向/后向计算，并将激活/梯度发送给下一个/上一个rank。在PP中，输入批次被划分为微批次（micro-batches），并使用不同的调度策略在PP rank之间进行调度。PP真正使得我们可以训练任意大的模型，因为它是按层进行划分的。

选择性激活检查点（Selective Activation Checkpointing, SAC）。在前向传播中，需要存储中间激活以便在后向传播中计算梯度。这些中间激活会占用大量内存。在SAC中，对于选定的前向传播块，只存储该块的输入激活，而不是该块中所有的激活。在对选定块进行后向传播时，会从存储的输入重新计算前向传播，然后再计算后向传播。

Optimus训练库的实现细节。在我们的Optimus训练库中，我们实现了以下功能：
* 我们为稠密模型和MoE模型分别实现了两个参考huggingface模型allenai/OLMo-1B-hf和allenai/OLMoE-1B-7B-0924，并加入了模型并行（张量、专家和流水线并行）支持。
* 我们为流水线并行（PP）实现了gpipe【6, GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, 2019, Advances in Neural Information Processing Systems】、1f1b【9, PipeDream: Generalized Pipeline Parallelism for DNN Training, 2019, SOSP】和interleaved-1f1b【10, Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, 2021, MLSys】调度策略，并启用了在使用PP训练MoE模型时进行专家负载均衡辅助损失计算的支持。
* 我们为三个模块实现了选择性激活检查点：规范化层（norm layer）、注意力块（attention block）和稀疏MoE块（SparseMoE block），并提供了在训练期间选择一个或多个模块的选项。
* 我们开发了一个可直接替换的优化模块FastSparseMoEBlock，与Hugging Face OLMoE模型中的SparseMoEBlock模块相比，它能显著加速MoE模型的训练。（详见3.1节）
* 我们为AdamW实现了分片优化器，并开发了一种新颖的、专家并行（EP）感知的分片优化器EPSO，当使用EP进行训练时，该优化器能带来更好的性能。（详见3.2节）

A2 方法细节

MoE性能优化

与稠密模型相比，MoE模型的计算更为复杂。因此，Hugging Face提供的现成实现性能不佳。本节将介绍我们为加速MoE模型训练而开发的两项性能优化。第一项是FastSparseMoE，它是对OLMoE Hugging Face实现中SparseMoE模块的一个优化替代模块。第二项是EP感知的分片优化器（EP-Aware sharded optimizer），它利用EP的语义，提供了比标准分片优化器更好的性能。

Table 3: MoE性能优化在组件级和完整训练级的加速效果。FSMOE影响训练的前向（F）和后向（B）组件。EPSO仅影响使用EP时的优化器组件。

3.1 FastSparseMoE

实现概述。在MoE Transformer模型中，SparseMoEBlock由N个专家和一个路由器组成。当对MoE模型应用专家并行（EP）时，SparseMoEBlock中的专家被划分，而路由器在EP ranks之间被复制，这导致每个rank拥有N/EP个专家和完整的路由器。算法1描述了我们在GPU上针对带有EP的FastSparseMoEBlock实现的五阶段优化方案。

阶段1 (Token通信)。SparseMoEBlock的输入首先被送到路由器，这是一个线性层，它接收S个输入token并产生大小为[S,N]的logits。然后对logits应用Softmax，再对Softmax的输出应用TopK来为每个输入token选择K个专家。被选中的专家可能不在同一个rank上，因此输入token需要被传输到相应的rank。这种通信是非均匀和不规则的，因为一个输入token只需与它所选专家所在的rank通信。理想情况下，应使用all2all操作来执行此通信。但我们观察到，尽管通信量更大，使用allgather操作将一个EP rank的所有输入token发送给所有其他EP rank会更快。这可以归因于OneCCL库中allgather操作的高效实现，因为它具有规则和均匀的通信模式。在算法1的第11-13行，我们在EP ranks之间通信输入和TopK的输出（权重和索引）。在前向传播中，我们对输入、权重和索引执行allgather。在后向传播中，我们对输入和权重应用reducescatter以累积部分梯度。

阶段2 (Token计数)。每个rank必须计算两种类型的计数：1) token计数：路由到该rank上某个专家的输入token数量。2) 专家计数：在该rank上本地选择的专家数量。算法1的第25-37行展示了我们进行token计数的GPU核。该核的输入是索引张量，其中包含来自所有EP ranks的所有输入token的K个选定专家ID。我们启动TH个线程，每个线程被映射来处理索引张量的一个行块。每个线程扫描其行块，并且只处理该rank本地的专家（第29行的条件）。如果一个选定的专家是该rank本地的，线程就将相应的token和专家计数加1。在第39-43行，我们对所有线程的部分token计数求和以得到最终计数，并执行前缀和操作来计算后续阶段所需的边界信息。

阶段3 (索引生成)。每个rank计算两种类型的索引：1) 输入索引：为该rank上的专家选择的输入token的索引。2) 输出索引：为该rank上选定专家的输入索引的索引。算法1的第53-72行展示了我们进行索引生成的GPU核。该核的输入是我们在阶段2中计算出的索引、token和专家计数。线程划分、线程到索引张量的映射以及本地专家的过滤与阶段2中的token计数核保持一致。此外，我们初始化了一个计数器张量，用于跟踪给定专家和线程路由了多少token。如果一个选定的专家是该rank本地的，线程首先从token计数中获取基准值（第61行），从计数器中获取偏移值（第62行）。然后输入token的ID被设置在输入索引的(base+offset)位置，而输出索引存储(base+offset)的值。图5展示了一个有无EP情况下的索引生成示例。


图5：在四个输入token（T=4）、四个专家（N=4）、每个输入token选择两个专家（K=2）的情况下，索引生成（输入索引和输出索引）的示例。在右侧的EP情况下，有两个rank，专家0和1放在rank 0上，专家2和3放在rank 1上。

阶段4 (专家计算)。每个专家的输入是为其选择的输入token的一个子集。在算法1的第74行，我们使用从阶段3生成的输入索引来为专家准备输入。一个专家就是一个包含三个线性层的MLP块：gate_proj、up_proj和down_proj。专家之间的计算没有依赖关系，它们可以被一起处理以提高效率，即通过一个核运行多个GEMM，而不是分开运行。这在GEMM规模小且专家数量多时尤其有用。在我们的实现中，我们将一个rank内跨专家的某个线性层的权重合并成一个单一的权重张量。例如，第76行的门控权重（gate weight）是通过连接该rank上所有专家的gate_proj.weight张量得到的。这个合并后的权重，连同合并后的输入和来自累积token计数的边界信息，使我们能够使用分组矩阵乘法（Grouped mm）操作，如第76、77和79行所示。

阶段5 (输出归约)。一个输入token被路由到K个专家，这K个专家的输出必须使用权重（TopK的输出）进行加权平均，以产生最终输出。在EP中，专家被划分到各个EP rank上，所以K个被选中的专家也分布在不同的EP rank上（$K = K_1 + .. + K_{EP}$）。因此，rank r需要对$K_r$个本地专家进行加权平均以生成部分输出。我们通过两个核来实现这个阶段：一个用于前向传播（第84-94行），另一个用于后向传播（第101-111行）。对于前向核，我们启动T×H个线程，每个线程映射到一个输出元素。每个线程遍历相应的本地选定专家（第88行），并使用阶段2计算的索引信息来缩放专家输出（第91行）。在第116行，我们执行reducescatter来归约输出并适当地拆分输出。对于后向传播，我们对梯度执行allgather，以确保所有rank都拥有完整的输出梯度。对于后向核，我们启动RT个线程，每个线程通过缩放和归约分别计算MLP输出和权重的梯度。

性能增益。FastSparseMoE (FSMOE) 优化的性能增益如表3所示。一个训练步骤可分为三个部分：前向、后向和优化器。FSMOE只影响前向和后向部分。我们为前向和后向组合部分实现了1.33倍到2.83倍的加速。在端到端训练层面，我们实现了1.11倍到1.71倍的加速。Mula-7B-A1B的加速比最高，因为它是在没有EP的情况下训练的，因此没有EP通信的开销。

算法1：带有专家并行（EP）的FastSparseMoEBlock

1: T = EP * S                 ▷ S - 序列长度, T - Token数量
2: NR = N/EP                  ▷ NR - 每个rank的专家数量
3: n_start = r * NR           ▷ EP rank r 的专家起始索引
4: n_end = (r + 1) * NR - 1   ▷ EP rank r 的专家结束索引
5: 
6: 阶段 1 : Token通信
7: logits = Router(input)              ▷ Shape(input)=[S,H]
8: probs = Softmax(logits)             ▷ Shape(logits)=[S,N]
9: weights, indices = TopK(probs)      ▷ Shape(weights/indices)= [S,K]
10: 
11: input = 前向allgather 后向reducescatter (input)
12: weights = 前向allgather 后向reducescatter (weights)
13: indices = 前向allgather (indices)
14: 
15: 阶段 2 : Token计数
16: TBS = 8                       ▷ TBS - Token块大小
17: TH = T/TBS                    ▷ TH - 线程数量
18: partial_token_counts = Zeros(NR*TH)
19: partial_cum_token_counts = Zeros(NR*TH+1)
20: cum_token_counts = Zeros(NR+1)
21: expert_counts = Zeros(T)
22: cum_expert_counts = Zeros(T+1)
23: 
24: 核函数 : 计算部分token计数
25: for tid in range(TH) do             ▷ 映射到GPU线程
26:   for i in range(TBS) do
27:     for k in range(K) do
28:       n = indices[t,k]
29:       if (n_start ≤ n ≤ n_end) then
30:         t = tid*TBS + i
31:         ln = n - n_start
32:         partial_token_counts[ln*TH+t] += 1
33:         expert_counts[t] += 1
34:       end if
35:     end for
36:   end for
37: end for
38: 
39: partial_cum_token_counts[1:] = PrefixSum(partial_token_counts)]
40: cum_expert_counts[1:] = PrefixSum(expert_counts)]
41: for n in range(NR + 1) do
42:   cum_token_counts[n] = partial_cum_token_counts[n*TH]
43: end for
44: 
45: 阶段 3 : 索引生成
46: RT = cum_token_counts[-1]      ▷ EP rank上专家的路由token数
47: selected_experts_indices = Zeros(RT)
48: input_indices = Zeros(RT)
49: output_indices = Zeros(RT)
50: counter = Zeros(NR,TH)
51: 
52: 核函数 : 生成输入和输出索引
53: for tid in range(TH) do             ▷ 迭代映射到GPU上的一个线程
54:   for i in range(TBS) do
55:     t = tid*TBS + i
56:     o_ind = cum_expert_counts[t]
57:     for k in range(K) do
58:       n = indices[t,k]
59:       if (n_start ≤ n ≤ n_end) then
60:         ln = n - n_start
61:         base = partial_cum_token_counts[ln*TH + tid]
62:         offset = counter[ln][tid]
63:         i_ind = base + offset
64:         input_indices[i_ind] = t
65:         output_indices[o_ind] = i_ind
66:         selected_expert_indices[o_ind] = k
67:         counter[ln][tid] += 1
68:         o_ind += 1
69:       end if
70:     end for
71:   end for
72: end for
73: 
74: 阶段 4 : 专家计算
75: mlp_in = input[input_indices]
76: gate_out = Grouped_mm(mlp_in, gate_weight, cum_token_counts)
77: up_out = Grouped_mm(mlp_in, up_weight, cum_token_counts)
78: mul_out = Silu(gate_out) * up_out
79: mlp_out = Grouped_mm(mul_out, down_weight, cum_token_counts)
80: 
81: 阶段 5 : 输出归约
82: function ExpertOutputReductionForward( )
83:   output = Zeros(T,H)
84:   for (t, h) in range(T)Xrange(H) do   ▷ 映射到GPU线程
85:     acc = 0
86:     base = cum_expert_counts[t]
87:     size = cum_expert_counts[t+1] - cum_expert_counts[t]
88:     for i in range(size) do
89:       k = selected_expert_indices[base+i]
90:       index = output_indices[base+i]
91:       acc += weights[t][k] * mlp_out[index][h]
92:     end for
93:     output[t][h] = acc
94:   end for
95:   return output
96: end function
97: 
98: function ExpertOutputReductionBackward(output_grad)
99:   mlp_out_grad = Zeros(RT,H)
100:  weights_grad = Zeros(EP*S,K)
101:  for rt in range(RT) do           ▷ 映射到GPU线程
102:    o_ind = output_indices[rt]
103:    t = input_indices[index]
104:    k = selected_expert_indices[rt]
105:    weight_grad_acc = 0
106:    for h in range(H) do
107:      mlp_out_grad[o_ind][h] = weights[t][k] * output_grad[t][h]
108:      weight_grad_acc += mlp_out[o_ind][h] * output_grad[t][h]
109:    end for
110:    weights_grad[t][k] = weight_grad_acc
111:  end for
112:  return mlp_out_grad, weights_grad
113: end function
114: 
115: output = ExpertOutputReduction(mlp_out, weights)
116: output = 前向reducescatter 后向allgather (output)
117: return output

3.2 EP感知的（EP-Aware）分片优化器

设计动机与原理。数据并行（DP）通过划分输入批次并在DP ranks间复制模型来减少训练时间。标准的DP实现，如PyTorch DDP（分布式数据并行），也在DP ranks间复制优化器状态，这会占用大量内存。分片优化器（Sharded Optimizers, SO）通过在DP ranks间分片优化器状态来解决这个问题。专家并行（EP）则划分MLP块中的专家参数，并在EP ranks间复制非专家参数（注意力、嵌入、lm_head、规范化层）。当在EP之上应用DP和SO时，由于SO只在DP ranks间划分优化器状态，对应于非专家参数的优化器状态会被复制EP次。为了解决这个问题，我们提出了EP感知的（Expert Parallel Aware）分片优化器（EPSO）。在EPSO中，我们首先根据参数的复制方式将rank r的参数$P_r$分为两组：$P_{Er}$（专家参数）和$P_{NE}$（非专家参数）。$P_E$中的参数在DP ranks间复制，而$P_{NE}$中的参数在DPxEP ranks间复制。在MoE的DP+EP模式中，$P_{Er}$和$P_{NEr}$分别对应专家参数和非专家参数。EPSO采用一种更细粒度的优化器分片策略：$P_{Er}$的优化器状态仅在DP维度上分片，而$P_{NEr}$的优化器状态则在DP和EP两个维度上都进行分片。图6展示了EPSO与标准分片优化器的对比示例。

性能增益。EPSO优化的性能增益如表3所示。EPSO仅影响优化器步骤，不影响训练步骤中的前向和后向传播。在优化器层面，我们实现了1.07倍到1.36倍的加速。在端到端训练中，我们实现了高达1.19倍的加速。


图6：用于MoE模型的EP感知分片优化器（EP=2, DP=2）。参数 P = [P E, P N E]。专家参数 P E = [P E1, P E2]。在EP=2中，PE1和PE2被放置在不同的专家rank上，PNE被复制。

A4 实验环境与结果

实验环境

  • 数据集：OLMoE-Mix-0924数据集【11, OLMoE: Open Mixture-of-Experts Language Models, 2024, arXiv】，包含4万亿个token。
  • 模型架构：Mula系列模型。稠密模型遵循OLMo【12, OLMo: Accelerating the Science of Language Models, 2024, ACL】架构，MoE模型遵循OLMoE【11, OLMoE: Open Mixture-of-Experts Language Models, 2024, arXiv】架构。具体模型配置参数见表1。

  • 硬件配置：

    • 平台：Aurora超级计算机【1, Aurora: Architecting Argonne’s First Exascale Supercomputer for Accelerated Scientific Discovery, 2025, arXiv】。
    • GPU：Intel PVC (Ponte Vecchio) GPU瓦片。
    • 规模：实验规模从32个节点（384个GPU瓦片）扩展至1024个节点（12288个GPU瓦片）。

  • 软件配置：

    • 训练框架：自研的Optimus训练库。
    • 通信库：OneCCL。

  • 训练超参数：

    • 优化器：AdamW，使用分片优化器（SO或EPSO）。
    • 超参数：beta1=0.9, beta2=0.99, eps=1e-8。
    • 学习率：最小4e-5，峰值4e-4，前2500步线性预热，之后余弦衰减。
    • 权重衰减：0.1，应用于所有参数。
    • 梯度裁剪：范数为1.0，仅在预热后应用。
    • 精度：bfloat16用于梯度归约。
    • 上下文长度：2048。
    • 全局批大小：630万个token。

Table 1: 模型配置。

实验结果

2.1 Mula-1B 与 Mula-7B-A1B 对比

实验内容：在256个节点（3072个GPU瓦片）上，使用DP=3072和AdamW分片优化器，对Mula-1B（稠密）和Mula-7B-A1B（MoE）模型在4万亿token的OLMoE-Mix-0924数据集上进行了完整预训练。

实验结果与分析：

  • 训练损失与准确率：如图1a所示，MoE模型Mula-7B-A1B的训练损失持续低于稠密模型Mula-1B。根据表2，在计算量相同的情况下，Mula-7B-A1B的平均准确率比Mula-1B高6.7%。
  • 训练效率：如图2a所示，Mula-7B-A1B模型仅用约5000亿个token（即1/8的训练数据）就达到了Mula-1B训练4万亿token后的最终准确率。
  • 困难基准表现：在最具挑战性的MMLU基准上（图2b），Mula-1B的准确率一直处于随机水平，而Mula-7B-A1B的准确率在训练约1万亿token后开始显著提升并稳步增长。
  • 与基线模型对比：如表2和图3所示，Mula-7B-A1B的性能与同样在OLMoE-Mix-0924数据集上训练1.3个epoch的OLMoE-1B-7B-0924模型【11, OLMoE: Open Mixture-of-Experts Language Models, 2024, arXiv】相当，验证了我们软件栈的正确性。

结论：在等计算量（iso-compute）的条件下，MoE模型比稠密模型更准确、更高效。

Table 2: Mula-1B稠密模型、Mula-7B-A1B MoE模型和allenai/OLMoE-1B-7B-0924 MoE模型的基准测试性能。


图1：(左) 在4万亿token的OLMoE-Mix-0924数据集上的训练损失。(右) 在同一数据集上训练至1000亿token的训练损失。


图2：Mula-1B和Mula-7B-A1B模型在4万亿token的OLMoE-Mix-0924数据集上的基准测试性能进展。


图3：Mula-7B-A1B和allenai/OLMoE-1B-7B-0924在中间训练检查点上评估的基准测试性能进展。

2.2 模型规模扩展

实验内容：以Mula-7B-A1B为基础，扩展出三个更大的MoE模型：Mula-20B-A2B、Mula-100B-A7B和Mula-220B-A10B（配置见表1）。所有模型均在256个节点（3072个GPU瓦片）上训练至1000亿个token。不同模型使用了不同的并行策略组合：
* Mula-20B-A2B (20B)：节点内12路专家并行（EP=12）。
* Mula-100B-A7B (100B)：跨节点4路流水线并行（PP=4），节点内EP=12，并对专家块使用激活检查点。
* Mula-220B-A10B (220B)：跨节点8路流水线并行（PP=8），节点内EP=12，并对规范化层、注意力和专家块使用激活检查点。
所有实验均使用1f1b流水线调度和EPSO分片优化器。

实验结果与分析：如图1b所示，随着模型规模的增大，训练损失相应降低，这符合预期。

结论：成功展示了使用不同模型并行配置进行MoE模型规模扩展的能力。

2.3 计算规模扩展

实验内容：选择最大的模型Mula-220B-A10B，将其计算规模从32个节点（384个GPU瓦片）扩展到1024个节点（12288个GPU瓦片）。

实验结果与分析：
* 训练损失：如图4a所示，随着计算规模的扩大（从而增大了全局批大小），训练损失健康下降。这符合预期，因为更大的批大小能带来更好的梯度近似和模型权重更新。
* 扩展效率：如图4b所示，从384扩展到768个GPU瓦片时，扩展效率仅下降3%。当扩展到超过一千个GPU瓦片时，效率下降约10%，并从1538到12288个GPU瓦片期间稳定在90%左右。
* 鲁棒性验证：为了排除MoE模型中不均匀的专家选择对扩展效率数据的影响，实验还进行了强制均匀路由（Forced Uniform Routing, FUR）的训练运行。在FUR中，所有专家接收相同数量和模式的token。如图4b所示，即使在FUR下，扩展效率的动态变化与常规训练运行相似。

结论：在将Mula-220B-A10B模型从384扩展到12288个GPU瓦片的过程中，实现了约90%的扩展效率。


图4：Mula-220B-10B模型预训练从32个节点（768个GPU tile）到1024个节点（12288个GPU tile）的计算扩展。

A7 补充细节

可靠性与容错

背景。在数千个GPU的规模上进行训练，对训练基础设施的所有组件都带来了压力，不稳定性和故障率随规模增加而上升。因此，需要开发一套强大的可靠性和容错功能，以减轻基础设施的压力，确保训练稳定，并从故障中快速恢复。本节将介绍我们为解决数据加载、模型加载、模型检查点、节点故障等问题而开发的功能。

数据预处理。典型Hugging Face数据集由多个数据文件组成，每个文件包含用于训练的文档。我们执行数据预处理，以避免在训练期间进行分词，并避免跨多个文件对训练数据进行随机和非连续的内存访问。我们的数据预处理流程包括三个步骤：分词、洗牌和分片。在分词步骤中，我们通过对数据文件$D_i$中的各个文档进行分词并用EOS token连接，为每个文件生成一个token数组$T_i$。给定上下文大小C，数据文件$D_i$将有$N_i$（$N_i = T_i/C$）个大小为C的训练实例。在洗牌步骤中，我们生成一个大小为N的排列顺序P，其中N是整个数据集的训练实例总数（$N=sum(N_i)$）。在分片步骤中，我们根据排列顺序P从分词后的数据文件中收集训练实例，并将其分片成多个numpy数组文件，这些文件随后以mmap模式懒加载。我们的数据处理流程允许所有数据并行rank以连续的方式从单个文件中加载内存，从而以最小的开销来消费训练中的token。

模型广播。在数据并行（DP）中，参数在DP rank之间被复制。因此，所有DP并行rank需要加载相同的模型（训练开始时的初始化模型或训练期间的检查点模型）。由于所有DP并行rank并行加载模型，这会给文件系统带来巨大压力，并可能导致模型加载本身挂起。为避免此问题，我们只加载一次模型，然后广播给所有DP并行rank。我们通过两种方式实现了广播：使用torch.broadcast或使用torch.all_reduce分布式调用。这项技术消除了挂起现象，并显著改善了启动时间，特别是对于大型模型。

双检查点（Dual checkpointing）。在深度学习模型训练中，检查点用于确保在训练运行失败时工作损失最小，或在计算资源非连续可用时恢复训练。检查点会存储恢复训练所需的所有必要信息（模型参数、优化器状态、训练步数等）。在检查点过程中，可能因文件系统I/O问题、检查点期间资源被回收等多种原因发生故障。因此，仅有一个检查点在检查点失败时无法帮助我们恢复训练。为解决此问题，我们实现了双检查点功能，即维护两个而不是一个检查点，这确保了总有一个有效的检查点可用于恢复训练。考虑一个使用双检查点的训练场景，每1000步进行一次检查点。在第1000步，检查点被保存在ckpt-1。在第2000步，检查点被保存在ckpt-2。在第3000步，由于ckpt-1是两者中较旧的，因此选择它进行检查点。现在，如果在向ckpt-1写入检查点时发生故障，这不会成为问题，因为ckpt-2是有效的，训练可以从中恢复。如果我们使用单个检查点，一旦检查点失败，就无法恢复训练运行。

持久化模型检查点（Persistent model checkpointing）。双检查点确保了无论训练进展如何，总有一个有效的检查点可供恢复。但训练本身也可能出现问题，如梯度爆炸、数据损坏导致发散等。因此，为了能够回溯到一个良好的训练状态，我们实现了持久化仅模型检查点功能。在此功能中，我们只保存模型参数，与完整检查点相比，其内存占用显著减少。对于使用AdamW优化器的BF16混合精度训练，仅模型检查点比完整检查点少占用8倍的内存。训练可以仅从模型参数重新开始，但优化器状态需要从头初始化。我们发现这样做不会对训练产生任何显著影响。考虑一个训练了15000步的场景，在10000步后损失曲线和梯度值开始发散。因为我们每1000步存储一次仅模型检查点，我们可以从第10000步的仅模型检查点以默认优化器状态重新开始模型训练。

DP-分散式模型检查点（DP-Scattered model checkpointing）。在DP中，模型在所有rank上被复制，模型检查点由第一个dp索引对应的rank（即第一个rank）完成。但是，当我们用模型并行训练大型模型时，遵循类似的方法，即让第一个dp索引对应的rank写入模型并行分片，将导致写操作集中在少数几个节点上。在DP-分散式模型检查点中，我们分散了负载，模型并行分片m由数据并行索引d写入，其中(d = m%DP)。例如，如果我们在12个节点上训练一个12路模型并行的模型，dp索引就是节点ID，因为每个节点有12个GPU瓦片。如果只有dp索引=0进行写入，那么所有模型分片都将从单个节点写入。但使用DP-分散式模型检查点，模型分片m将由节点m写入。

硬节点故障处理。一个训练运行在一定数量的节点上启动。特定节点可能因多种原因（如ping失败、分段错误、操作系统错误等）发生故障。当这种情况发生时，必须通过排除发生故障的节点来重新启动训练运行。在硬节点故障处理功能中，我们通过启动带有一些额外缓冲节点的训练运行来自动完成此重启过程，并通过用缓冲节点之一替换失败的节点来重新启动运行。

软节点故障处理。在硬节点故障中，训练运行在节点故障后立即退出。但在软节点故障中，训练运行会继续进行，但在发生软故障的节点上会产生局部NaN值。如果未被检测到，这将导致NaN权重并用NaN污染检查点。为避免此问题，我们在每个rank中检查局部损失和梯度是否为NaN，如果出现，则标记出现NaN的rank对应的节点并退出训练运行。然后我们通过用缓冲节点之一替换失败的软节点来自动重新启动训练运行。

A5 结论

本文通过在数千个Intel GPU上对混合专家（MoE）大型语言模型进行数万亿token的预训练，展示了Aurora超级计算机的强大能力。在规模扩展方面，我们将模型规模和计算规模分别推向了2200亿参数和12288个GPU瓦片，并在将Mula-220B-A10B模型从384个GPU瓦片扩展到12288个GPU瓦片时，观察到了约90%的扩展效率。