发布 Olmo-core 3：面向大型 MoE 的开源、可扩展训练基础设施

Source: https://huggingface.co/blog/allenai/olmocore3 📄Tech Report (https://allenai.org/papers/olmocore3) | 💻Code (https://github.com/allenai/olmo-core) | 🧩Interactive demo (https://narrative.allen.ai/scaling-up-training) Olmo-core 3 blog draft - Google Docs-image-1 (2) (https://cdn-uploads.huggingface.co/production/uploads/638e39b249de7ae552d977b5/gZSxRv34VBkkUCXDoq7UE.png)

今天，我们正式发布Olmo-core 3（https://github.com/allenai/olmo-core），这是对大语言模型开发框架的一次重大升级，其中重新设计了开源的混合专家（MoE）训练系统。

Olmo-core 3 旨在将 MoE 训练扩展到万亿参数规模，同时保持计算效率。它是下一代 Olmo 背后的核心系统之一，也是我们持续致力于开放每个新模型背后的工具与训练基础设施的一部分。

训练大型 AI 模型需要大量算力，推高了成本和能耗，也让许多学术研究人员和小型实验室难以开展前沿模型开发。MoE 模型提供了一种更高效的路径——它们可以容纳更多的学习组件（即参数），而无需每次输入都动用所有参数。但整个模型仍需存储在 GPU 显存中并在训练期间更新，而在集群范围内将输入路由到正确的专家（MoE 中的专用组件）也会产生通信与协调开销。随着 MoE 规模增大，这些开销可能抵消仅使用模型部分组件进行计算所带来的大部分效率优势。

Olmo-core 3 正是为了弥合这一差距而构建的。在一项基准测试中，我们把专家池从 8 个扩展到 128 个，同时每个 token（语言模型处理的文本最小单位）仍只选取 4 个专家，每个 token 的活跃参数数量大致保持在约 3.2B。总参数容量从 4.6B 增长到 47B，而训练吞吐量的下降不到 5%。

同一套基础设施已在总参数超过一万亿的配置下完成基准测试。

围绕 MoE 的实际运作方式构建训练栈

Olmo-core 随着每一代 Olmo 的演进而不断发展。

我们对稀疏模型的研究可以追溯到 OlmoE，它采用了配备 64 个路由专家的 MoE 架构。相比之下，Olmo 3 采用的是稠密架构，即对每个 token 几乎整个模型都是活跃的，其训练栈也围绕这一设计构建。Olmo-core 3 则扩展了该框架，引入了专为更大规模 MoE 模型设计的训练系统。

我们早期在 Olmo-core 中实现的 MoE 采用完全分片数据并行（FSDP），会在每一批小的训练数据到来时收集并重新分片模型权重。Olmo-core 3 则切换到基于分布式数据并行（DDP）的系统：它将专家常驻于 GPU 上，并将相关数据路由到对应专家，避免了反复的权重收集过程。

NVIDIA 的 Megatron-Core 是训练大型 MoE 的成熟方案。Olmo-core 3 为 Olmo 背后的框架带来了集成化的 MoE 训练栈，其重新设计后的吞吐量优于我们此前基于 FSDP 的实现。在八块 NVIDIA B300 GPU 上的初步测试中，一个 470 亿参数的 MoE 使用新栈时每 GPU 每秒可处理 52,000 个 token，而旧实现仅为 19,400——吞吐量约为 2.7 倍。

扩展并优化 MoE 训练

Olmo-core 3 将多种在 GPU 集群上分布式部署大型 MoE 的技术，与提升路由和计算效率的优化手段结合在一起。

有三项技术决定了模型及其训练状态如何在硬件上划分：

• 专家并行（Expert parallelism）：将专家分散到各 GPU 上，使每个 GPU 只需存储完整专家池的一部分。
• 流水线并行（Pipeline parallelism）：将模型的各层——即依次对输入进行变换的连续阶段——划分到不同的 GPU 组中，降低每个 GPU 需要在显存中保存的模型份额。
• 分布式优化器（A distributed optimizer）：将优化器状态——训练期间用于计算并应用更新的额外数据——分散到各 GPU 上，而不是在每块 GPU 上都存储完整副本。

这些技术共同作用，使 MoE 能够扩展规模，而无需每块 GPU 都在显存中保存整个模型及其训练状态。

Olmo-core 3 还降低了将数据路由到正确专家并执行专家计算的开销。*逐行专家并行（Rowwise expert parallelism）*将路由数据直接放入专家输入缓冲区，最大限度地减少重排所需的额外工作。*GPU 常驻路由（GPU-resident routing）将路由元数据保留在 GPU 上，使 CPU 可以排队处理工作，而无需等待这些信息被拷回。而分组 GEMM（grouped GEMM）*将大量小型专家计算合并在一起，使 GPU 能够更高效地执行它们。

最后，Olmo-core 3 支持MXFP8，这是一种更低精度的数值格式，用更少的比特表示部分数值。只要数值转换的代价小于由此节省的开销，它就能减少计算量以及 GPU 之间传输的数据量。

我们在四块 NVIDIA B300 GPU 上进行了一项受控基准测试，对端到端训练吞吐量中 MXFP8 的效果进行了测量，测试中各专家的工作负载均匀分布。在系统中受助最大的部分开启 MXFP8 后，训练吞吐量比基线所用的高精度格式 BF16 高出约 21%，而峰值活跃显存从 103 GiB 降至 95 GiB。大部分收益来自前馈计算和专家间的数据搬运，而非仅来自注意力机制。

这些技术和优化必须协同工作。加速训练的某一部分可能会在其他地方产生代价：更快的计算可能需要更多的数据搬运，而减少比特传输如果数据转换耗时过长也可能收效甚微。Olmo-core 3 正是围绕这些跨越整个训练流程的权衡来构建的，这让我们——以及使用这套开源栈的研究人员——能够灵活掌控各环节的组合方式。

探索我们的交互式演示（https://narrative.allen.ai/scaling-training），了解数据并行、专家并行和流水线并行如何协同扩展 MoE 训练——从单 GPU 到多 GPU。

向万亿参数规模扩展

我们已在 NVIDIA B300 GPU 上对 Olmo-core 3 进行了多种配置的基准测试，其中包括一个 1.2 万亿参数、每个 token 活跃参数 583.6 亿、部署在 512 块 GPU 上的模型。其观测到的最高吞吐量为每 GPU 每秒 858 TFLOP/s——即每块 GPU 每秒可完成的有效模型计算量。这些测试采用随机路由来测量系统性能，而非测量训练后模型的质量。

我们还测试了 DeepEP v2——一种在 GPU 间处理专家间通信的替代方案，达到了总参数 2.38 万亿的配置。这是一次短容量测试，而非完整的训练运行，因此它展示的是 Olmo-core 3 能够达到的规模，而非持续训练性能。

在这些规模下，系统性能只是故事的一部分。我们的技术报告还记录了那些对我们如何训练 MoE 及如何评估其性能产生影响的实验。例如：

• 一个本意鼓励均衡路由的评分指标，可能在实际负载变得愈发不均衡时反而持续提升。我们将这种失效称为token 选区划分（token gerrymandering）。
• 因为专家处理的 token 更少而调低其学习率——即其训练更新的幅度——在我们测试的模型家族中并没有改善结果。
• 即使矩阵维度相同，当被处理的数值发生变化时，GPU 计算所耗时间也不同。因此，性能对比不仅需要形状匹配，也需要输入数值匹配。
• 在不同的 GPU 流上重叠通信与计算并不总是能让训练更快。在某些测试中，它反而拖慢了端到端执行——这提醒我们，更多的重叠并不一定意味着更高的吞吐量。

报告在阐述这些发现的同时，也介绍了我们测试过但最终未采纳的方案。

为下一代 Olmo 而建，对所有人开放

Olmo-core 3 是我们下一步构建工作的基础。下一代 Olmo 将采用 MoE 架构，我们希望它成为迄今能力最强的 Olmo——在我们最大的数据集上训练，并拥有我们最长的上下文窗口。

新的训练栈让我们能够超越此前的 MoE 工作成果，同时在模型和硬件演进过程中为灵活调整训练提供更大空间。而且它是完全开放的——研究人员和开发者可以使用 Olmo-core 3 训练自己的 MoE，将其适配到不同硬件上，并针对路由、并行策略和系统其他部分开展实验。

这也是我们看待开源模型开发方式的一部分——当模型背后的基础设施和训练决策同样开放时，模型权重才会更有价值。

若想深入了解系统设计、实验、消融实验，以及我们一路探索过的方案，请阅读我们的技术报告，并在 GitHub 上探索 Olmo-core 3（https://github.com/allenai/olmo-core）。