Allen Institute for AI (Ai2) 近日发布了 Olmo-core 3，这是其大型语言模型训练框架的重大升级版本。该版本引入了重新设计的专家混合（MoE）训练系统，旨在将 MoE 模型的训练规模扩展至万亿参数级别，同时保持计算效率。
Olmo-core 3 是下一代 Olmo 模型背后的核心基础设施之一，延续了 Ai2 开放每一代模型背后工具与训练系统的理念。通过这一框架，研究人员和开发者可以更高效地构建和训练大规模稀疏模型。
解决 MoE 训练的效率瓶颈
训练大型人工智能模型通常伴随着高昂的计算成本和能源消耗，这限制了学术研究机构及小型实验室开发先进模型的能力。虽然 MoE 架构允许模型包含更多参数而无需在每个输入中激活所有部分，但完整的模型权重仍需存储在 GPU 内存中并在训练期间更新。此外，将输入数据路由至正确的“专家”组件会在集群间产生通信和协调开销。随着 MoE 规模的扩大，这些开销可能会抵消仅使用部分模型所带来的计算优势。
Olmo-core 3 针对这一缺口进行了优化。在一项基准测试中，团队将专家库从 8 个扩展到 128 个，同时保持每个 token 仅激活 4 个专家，使得活跃参数量大致固定在 3.2B。尽管总参数容量从 4.6B 增长至 47B，训练吞吐量的下降幅度不到 5%。
基于实际工作负载重构训练堆栈
Ai2 对稀疏模型的研究可追溯至使用 MoE 架构的 OlmoE。相比之下，Olmo 3 采用密集架构，其训练堆栈围绕此设计构建。Olmo-core 3 在此基础上扩展，专为更大规模的 MoE 模型设计了新的训练系统。
此前 Olmo-core 中的 MoE 实现采用完全分片数据并行（FSDP），配置为在每次小批量训练后收集和重新硬化模型权重。Olmo-core 3 切换至基于分布式数据并行（DDP）的系统，将专家保留在 GPU 上并将相关数据路由给它们，从而避免了重复积累带来的开销。
作为对比，NVIDIA 的 Megatron-Core 是训练大型 MoE 的主流选择。Olmo-core 3 通过重新设计，提高了之前基于 FSDP 实现的吞吐量。在八台 NVIDIA B300 GPU 上的初步测试显示，一个 47 亿参数的 MoE 模型每秒处理 52,000 个 token，相比旧实现的 19,400 个 token，吞吐量提升了约 2.7 倍。
扩展与优化技术细节
Olmo-core 3 结合了多种技术在 GPU 集群中分配大型 MoE，并优化了路由和计算效率。决定模型及其训练状态在硬件上划分的三种核心技术包括：
- 专家并行性：将专家分布在不同的 GPU 上，使每个 GPU 仅存储全部专家库的一部分。
- 管道并行性：将模型层划分为连续的阶段，由不同的 GPU 组处理输入转换，减少单个 GPU 需要存储的模型量。
- 分布式优化器：将用于计算和应用训练更新的额外数据（优化器状态）分散在 GPU 之间传播，而非在每个 GPU 上存储完整副本。
这些技术使得 MoE 能够在不要求每个 GPU 将整个模型及其训练状态存入内存的情况下进行扩展。此外，Olmo-core 3 降低了数据路由至正确专家并执行计算的成本。行式专家并行性（Rowwise expert parallelism）将路由数据直接放入专家输入缓冲区，最大限度地减少重排所需的额外工作，允许 CPU 排队工作而无需等待信息复制回来。集成的 GEMM 操作合并了许多小型专家计算，使 GPU 能更有效地执行它们。
Olmo-core 3 还支持 MXFP8 数字格式，这是一种精度较低的格式，用更少的位表示数值。这减少了计算和 GPU 间传输的数据量，前提是节省超过格式转换成本。在四台 NVIDIA B300 GPU 的控制基准测试中，MXFP8 使端到端训练吞吐量比基线 BF16 高出 21%，峰值主动内存从 103GB 降至 95GB。大部分收益来自前向计算和专家间数据移动，而非仅仅关注单一环节。
这些技术和优化必须协同工作。加速训练的某一部分可能在其他地方造成成本；更快的计算可能需要更多的数据移动，而如果数据转换耗时过长，移动更少的比特可能并无帮助。Olmo-core 3 围绕这些权衡构建，让研究人员能够控制各部分的组合方式。
迈向万亿参数规模
Ai2 在 NVIDIA B300 GPU 上对 Olmo-core 3 进行了多种配置的基准测试，其中包括一个 1.2 万亿参数的模型，在 512 个 GPU 上每 token 有 58.36 亿个参数。最高观察到的吞吐量为 858 TFLOP/s/GPU。这些测试使用随机路由来测量系统性能，而非评估模型训练质量。
团队还试验了 DeepEP v2，这是一种替代的 GPU 专家间通信方式，达到了 2.38 万亿总参数的配置。这是一次短时测试，而非完整训练运行，主要展示 Olmo-core 3 可达到的规模上限，而非持续训练性能。
除了系统性能，技术报告还记录了关于 MoE 培养及表现衡量的实验发现：
- 旨在鼓励平衡路由的分数可以提高，即使实际工作量变得不那么平衡，这种现象被称为“失败的代币策略”。
- 降低专家的学习率，以减小他们处理较少 token 时的训练更新规模。
- 当被处理的值发生变化时，即使矩阵维度相同，GPU 计算所需时间也不同。因此，性能比较需要匹配输入值以及形状。
- 在单独的 GPU 流上重叠通信和计算并不总是加快训练速度。在一些测试中，它反而减缓了端到端执行速度，表明更多的重叠不一定意味着更高的吞吐量。
面向未来的开放基础设施
Olmo-core 3 是 Ai2 下一步建设的基础。即将推出的下一代 Olmo 将使用 MoE 架构，目标是成为迄今为止能力最强的 Olmo 模型。随着模型和硬件的发展，新的堆栈使其能够超越之前的 MoE 工作。研究人员和开发人员可以使用 Olmo-core 3 训练自己的 MoE，适应不同的硬件，并尝试路由、并行化及其他系统组件。
