Olmo-core 3发布:面向大型混合专家模型的开源可扩展训练基础设施
我们发布了 Olmo-core 3,这是我们大型语言模型开发框架的重大升级,其核心亮点是重新设计的开源混合专家(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 的实现大幅提升。在 8 块 NVIDIA B300 GPU 上的初步测试中,一个 470 亿参数的 MoE 模型在新训练栈下每块 GPU 每秒可处理 52,000 个 token,而旧实现只有 19,400 个——吞吐量约为原来的 2.7 倍。
扩展并优化 MoE 训练
Olmo-core 3 采用多种技术将大型 MoE 分布在 GPU 集群上,并通过优化使路由和计算效率更高。
三种技术决定了模型及其训练状态在硬件上的分割方式:
- 专家并行 将专家分布在不同的 GPU 上,每个 GPU 仅存储完整专家池的一部分。
- 流水线并行 将模型层——即处理输入的各个连续阶段——分配给不同的 GPU 组,从而减少每个 GPU 内存中需要保留的模型数据量。
- 分布式优化器 将优化器状态——训练期间用于计算和应用更新所需的额外数据——分布在多个 GPU 上,而不是在每个 GPU 上存储完整副本。
这些技术共同使得 MoE 能够扩展规模,而无需每个 GPU 都在内存中保持完整的模型及其训练状态。
Olmo-core 3 还降低了将数据路由到正确专家以及执行其计算的成本。按行专家并行直接将路由数据放入专家输入缓冲区,最小化重新排列数据所需的额外工作。GPU 驻留路由将路由元数据保留在 GPU 上,使 CPU 无需等待该信息复制回来即可排队任务。此外,分组 GEMM将许多小型专家计算合并,以便 GPU 能更高效地执行。
最后,Olmo-core 3 支持 MXFP8,这是一种低精度数值格式,能用更少的位表示某些数值。只要这种节省抵消了在不同数值格式间转换的成本,它就能减少计算量和 GPU 之间移动的数据量。
我们在一个受控基准测试中测量了 MXFP8 对端到端训练吞吐量的影响,该测试使用四块 NVIDIA B300 GPU,并在专家之间均匀分配工作量。在系统中受益最大的部分启用 MXFP8 后,与作为基准的高精度格式 BF16 相比,训练吞吐量提高了约 21%,而峰值活动内存从 103 GiB 降至 95 GiB。大部分收益来自前馈计算和在专家之间移动数据,而非仅来自注意力机制。
这些技术与优化必须协同工作。加速训练流程中的某一部分,往往会在其他环节产生额外开销;更快的计算可能意味着更多的数据搬运,而如果数据转换耗时过长,单纯减少搬运的数据量也未必有益。Olmo-core 3 正是围绕这些贯穿训练全流程的权衡关系构建的,从而赋予我们——以及使用这套开放技术栈的研究者——对各组件如何拼合的控制权。
查看我们的交互式演示,了解数据并行、专家并行和流水线并行如何协同工作以扩展 MoE 训练——从单卡扩展到多卡。
迈向万亿参数级扩展
我们在 NVIDIA B300 GPU 上对 Olmo-core 3 进行了多种配置的性能基准测试,其中包括一个总参数量达 1.2 万亿、每个 token 激活参数量为 583.6 亿、且运行在 512 张 GPU 上的模型。其观察到的最高吞吐量为 858 TFLOP/s/GPU——该指标衡量的是每张 GPU 每秒执行的有效模型计算量。这些测试使用随机路由来衡量系统性能,而非已训练模型的质量。
我们还尝试了 DeepEP v2,这是一种处理专家间跨 GPU 通信的替代方案,达到了总参数 2.38 万亿的配置。这是一次短周期测试而非完整训练,旨在展示 Olmo-core 3 能够达到的规模,而非持续训练性能。
在此类规模下,系统性能只是全貌的一部分。我们的技术报告还记录了指导我们如何训练 MoE 以及评估其性能的实验,例如:
- 一个旨在鼓励路由均衡的得分指标,可能在实际工作负载变得越不均衡时反而有所改善。我们将这种失效现象称为token 重划选区(token gerrymandering)。
- 在处理 token 数量较少的专家时,降低其学习率(即训练更新步长)并未在我们测试的模型家族中带来更好的结果。
- 当处理的数据值变化时,即使矩阵维度相同,GPU 的计算耗时也会有所不同。因此,性能对比不仅需要匹配的形状,还需要匹配的输入值。
- 在独立的 GPU stream 上让通信与计算重叠,并不总能加快训练速度。有些测试中,它反而拖慢了端到端执行——这提醒我们,重叠更多并不等于吞吐更高。
报告中详细介绍了这些发现,以及我们测试过但最终没有采用的方案。
为下一代 Olmo 而建,向所有人开放
Olmo-core 3 是我们后续工作的基石。下一代 Olmo 将采用 MoE 架构,我们的目标是让它成为迄今最强的 Olmo——使用我们最大的数据集、最长的上下文窗口进行训练。
新框架让我们能在以往 MoE 工作的基础上进一步扩展,同时随着模型和硬件的演进,保留更大的灵活度来调整训练方案。而且它完全开源——研究者和开发者可以用 Olmo-core 3 训练自己的 MoE,适配不同的硬件,也可以在路由、并行等系统环节上做实验。
这也是我们对开放模型开发的理解之一——如果背后的基础设施和训练决策同样是开放的,模型权重才会更有价值。
想深入了解系统设计、实验、消融研究以及我们尝试过的各种方案,请阅读我们的技术报告,并在 GitHub 上探索 Olmo-core 3。



