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

今天我们正式发布 Olmo-core 3，这是我们用于开发大语言模型框架的一次重大升级，引入了重新设计的开源混合专家（MoE）训练系统。

Olmo-core 3 的设计目标是支撑 MoE 训练扩展到万亿参数级别，同时保持计算效率。它是下一代 Olmo 的核心系统之一，也是我们持续开放每个新模型背后工具和训练基础设施承诺的一部分。

训练大型 AI 模型需要大量算力，成本和能耗随之攀升，这让许多学术研究者和小规模实验室难以触及先进模型的开发。MoE 模型提供了一种更高效的方案——它们可以包含更多可学习的组件（即参数），而无需让每个输入都动用全部参数。但完整模型仍然需要分布在 GPU 显存中，并在训练时持续更新；将输入引导到正确的专家（即 MoE 中的专门组件）——这在整个集群中会产生通信和协调开销。随着 MoE 规模扩大，这些成本可能会侵蚀"每个输入只使用部分模型"所带来的计算优势。

Olmo-core 3 就是为了弥合这一差距而构建的。在一项基准测试中，我们将专家池从 8 个扩展到 128 个，同时每个 token（语言模型处理的基本文本单元）仍只选取 4 个专家——使每个 token 的活跃参数量大致固定在约 32 亿。总参数容量从 46 亿增长到 470 亿，而训练吞吐量仅下降了不到 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 实现的吞吐量。在 8 块 NVIDIA B300 GPU 上的初步测试中，一个 470 亿参数的 MoE 使用新栈每 GPU 每秒处理 52,000 个 token，而此前实现仅为 19,400 个——约为原来的 2.7 倍吞吐量。

MoE 训练规模化与优化

Olmo-core 3 结合了多种技术，用于在 GPU 集群上分布式部署大型 MoE，并通过优化使路由和计算更加高效。

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

专家并行将专家分布到各 GPU 上，使每个 GPU 只存储完整专家池的一部分。

流水线并行将模型的层（对输入进行转换的连续阶段）划分到不同的 GPU 组，降低了每个 GPU 需要在内存中保存的模型比例。

分布式优化器将优化器状态（训练期间用于计算和应用更新的附加数据）分布到各 GPU，而不是在每个 GPU 上存储完整副本。

这些技术共同使 MoE 能够在不要求每个 GPU 保留完整模型及其训练状态的情况下进行扩展。

Olmo-core 3 还降低了将数据路由到正确专家并执行其计算的成本。行级专家并行将路由数据直接放入专家输入缓冲区，最大程度减少了重新排列数据所需的额外工作。GPU 常驻路由将路由元数据保留在 GPU 上，使 CPU 可以在无需等待该信息被复制回来的情况下排队工作。而分组 GEMM 将许多小型专家计算合并执行，使 GPU 能够更高效地处理。

最后，Olmo-core 3 支持 MXFP8，这是一种用更少位数表示部分值的低精度数值格式。只要节省的算力和 GPU 间数据移动量超过数值格式转换的成本，就能带来收益。

我们在受控环境下测量了 MXFP8 对端到端训练吞吐量的影响，基准测试在 4 块 NVIDIA B300 GPU 上进行，工作在专家间均匀分布。启用 MXFP8（作用于系统中收益最大的部分）后，训练吞吐量比使用 BF16（我们作为基准的高精度格式）高出约 21%，峰值活跃内存从 103 GiB 降至 95 GiB。大部分收益来自前馈计算和专家间数据移动，而非注意力部分。

这些技术和优化必须协同工作。加速训练的一个环节可能在其他地方产生成本；更快的计算可能需要更多的数据移动，而移动更少的位，如果数据转换耗时过长，也可能没有帮助。Olmo-core 3 围绕整个训练过程中的这些权衡而构建，让我们——以及使用该开源栈的研究者——能够控制各组件如何协同工作。

探索我们的交互式演示，了解数据并行、专家并行和流水线并行如何协同工作以扩展 MoE 训练——从单块 GPU 到大规模集群。

扩展至万亿参数级别

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

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

在这个规模下，系统性能只是问题的一个方面。我们的技术报告还记录了为informed我们如何训练 MoE 及衡量其性能而进行的实验。例如：

一个旨在鼓励均衡路由的评分机制，其分数可能提升而实际工作负载却变得更不均衡。我们称这种现象为"失败 token 选区划分"（failure token gerrymandering）。

因为专家处理的 token 较少而降低其学习率（训练更新的幅度），在我们测试的模型族中并未带来改善。

当处理的值发生变化时，GPU 计算耗时也会不同，即使矩阵维度相同。因此性能比较不仅需要匹配输入形状，还需要匹配输入值。

在独立的 GPU 流上重叠通信和计算并不总是让训练更快。在某些测试中，它反而减慢了端到端执行——这提醒我们，更多重叠不一定意味着更高吞吐量。

该报告解释了这些发现，以及我们测试和最终未采纳的方法。

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

Olmo-core 3 是我们正在构建的一切的基石。我们的下一代 Olmo 将使用 MoE 架构，我们的目标是让它成为迄今为止最强大的 Olmo，使用我们最大的数据集训练并拥有最长的上下文窗口。

新的技术栈让我们能够超越此前的 MoE 工作，同时在模型和硬件演进时给予我们更大的适应性来调整训练。而且它是完全开源的——研究者和开发者可以使用 Olmo-core 3 训练自己的 MoE、将其适配到不同硬件上，并实验路由、并行化以及系统的其他部分。

这是我们思考开源模型开发的方式的一部分——当模型权重背后的基础设施和训练决策也是开放的时候，它们会更有用。

欲深入了解系统设计、实验、消融分析以及我们一路走来测试的方法，请阅读我们的技术报告并在 GitHub 上探索 Olmo-core 3。