推出 Olmo-core 3：面向大型 MoE 的开放、可扩展训练基础设施

背景与动机

训练大型 AI 模型需要大量算力，导致成本和能耗上升，使许多学术研究者和较小实验室难以进行先进模型开发。MoE 模型提供了一种更高效的方法：它们可以包含更多学习组件（参数），而无需每个输入都使用所有参数。但完整模型仍需存储在 GPU 内存中并在训练期间更新，且跨集群将输入路由到正确的专家会带来通信和协调成本。随着 MoE 规模增长，这些成本可能侵蚀使用部分模型带来的计算优势。Olmo-core 3 旨在缩小这一差距。

• MoE 模型通过仅激活部分参数来提升效率，但存储和路由成本随规模增长而增加。
• Olmo-core 3 的目标是支持万亿参数级 MoE 训练，同时保持计算效率。

性能基准

在一个基准测试中，专家池从 8 个增加到 128 个，但每个 token 仍只选择 4 个专家，每个 token 的活跃参数大致固定在约 3.2B。总参数容量从 4.6B 增长到 47B，而训练吞吐量下降不到 5%。同一基础设施已在超过一万亿总参数的规模上进行了基准测试。

• 专家池从 8 扩展到 128，每 token 活跃参数约 3.2B，总参数从 4.6B 增至 47B，吞吐量下降不到 5%。
• 该基础设施已测试超过 1 万亿总参数。

训练栈设计

Olmo-core 随每代 Olmo 演进。早期稀疏模型工作可追溯到 OlmoE，使用 64 个路由专家的 MoE 架构。Olmo 3 则使用密集架构，训练栈围绕该设计构建。Olmo-core 3 扩展了框架，设计了面向更大 MoE 模型的训练系统。

早期 MoE 实现使用完全分片数据并行（FSDP），为每个小批量训练数据收集和重新分片模型权重。Olmo-core 3 转向基于分布式数据并行（DDP）的系统，将专家常驻在 GPU 上并将相关数据路由到它们，避免重复的权重收集。

NVIDIA 的 Megatron-Core 是训练大型 MoE 的成熟选项。Olmo-core 3 为 Olmo 背后的框架带来了集成的 MoE 训练栈，其重新设计相比早期基于 FSDP 的实现提高了吞吐量。在八块 NVIDIA B300 GPU 上的初步测试中，一个 47B 参数的 MoE 使用新栈每 GPU 每秒处理 52,000 个 token，而早期实现为 19,400 个，吞吐量约为 2.7 倍。

• 从 FSDP 转向 DDP，专家常驻 GPU，避免重复权重收集。
• 在 8 块 B300 GPU 上，47B MoE 吞吐量从 19,400 token/s/GPU 提升至 52,000 token/s/GPU，约 2.7 倍。

扩展与优化技术

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 间移动的数据量，只要节省超过格式转换的成本。在四块 NVIDIA B300 GPU 上的受控基准测试中，工作均匀分布在各专家上，启用 MXFP8 后训练吞吐量比 BF16 基线高约 21%，峰值活跃内存从 103 GiB 降至 95 GiB。大部分收益来自前馈计算和专家间数据移动，而非仅注意力部分。

这些技术和优化必须协同工作。加速训练的一部分可能在其他地方产生成本；更快的计算可能需要更多数据移动，而移动更少比特如果转换数据耗时过长可能无益。Olmo-core 3 围绕整个训练过程的这些权衡构建，让用户控制各部分如何配合。

• 专家并行、流水线并行和分布式优化器共同实现内存高效的扩展。
• 行式专家并行、GPU 常驻路由和分组 GEMM 降低路由和计算开销。
• MXFP8 在 B300 上提升吞吐量约 21%，峰值内存从 103 GiB 降至 95 GiB。
• 优化需权衡计算速度、数据移动和转换成本。

万亿参数级扩展

Olmo-core 3 已在 NVIDIA B300 GPU 上跨多种配置进行基准测试，包括一个 1.2 万亿参数模型，每 token 活跃参数 58.36B，跨 512 块 GPU。其最高观测吞吐量为 858 TFLOP/s/GPU。这些测试使用随机路由来测量系统性能，而非训练模型的质量。

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

在这些规模下，系统性能只是部分情况。技术报告还记录了影响我们训练 MoE 和衡量其性能的实验。例如：旨在鼓励平衡路由的分数可能在实际工作负载变得不均衡时反而改善，我们称这种失败为 token gerrymandering；降低专家的学习率（因其处理更少 token）在测试的模型家族中未改善结果；GPU 计算在处理值变化时耗时不同，即使矩阵维度相同，因此性能比较需要匹配输入值和形状；在单独 GPU 流上重叠通信和计算并不总是加速训练，某些测试中反而减慢端到端执行。报告解释了这些发现以及我们测试并选择不采用的方法。

• 1.2 万亿参数模型在 512 块 B300 上达到 858 TFLOP/s/GPU。
• DeepEP v2 实验达到 2.38 万亿总参数，但为短容量测试。
• 发现包括 token gerrymandering、学习率调整无效、输入值影响 GPU 计算时间、重叠通信不一定加速。

面向下一代 Olmo

Olmo-core 3 是我们下一步构建的基础。下一代 Olmo 将使用 MoE 架构，目标成为我们最有能力的 Olmo，在我们最大的数据集上训练，并具有最长的上下文窗口。

• 下一代 Olmo 将采用 MoE 架构，使用最大数据集和最长上下文窗口。