用于 MoE 训练的 Olmo-core 3（约 8 分钟阅读）

# 推出 Olmo-core 3：面向大规模 MoE 的开放、可扩展训练基础设施 | Ai2 来源：https://allenai.org/blog/olmocore3 今天我们发布了 **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 是语言模型处理的文本最小单元——使每个 token 的激活参数数量大致稳定在约 3.2B。总参数容量从 4.6B 增长到 47B，而训练吞吐量下降不到 5%。 同一套基础设施已在超过一万亿总参数的规模下完成了基准测试。 ### **围绕 MoE 的实际运行方式构建训练栈** Olmo-core 随着每一代 Olmo 不断演进。 我们对稀疏模型的研究可追溯至 [OlmoE](https://allenai.org/blog/olmoe-an-open-small-and-state-of-the-art-mixture-of-experts-model-c258432d0514)，它采用了配备 64 个路由专家的 MoE 架构。相比之下，[Olmo 3](https://allenai.org/blog/olmo3) 使用的是稠密架构，即几乎模型的全部参数在每个 token 上都被激活，其训练栈也是围绕这一设计构建的。Olmo-core 3 在框架中加入了专为更大型 MoE 模型设计的训练系统。 我们在 Olmo-core 中早期实现的 MoE 使用了全分片数据并行（FSDP），针对每个小批次的训练数据重新收集并重新分片模型权重。Olmo-core 3 则切换为基于[分布式数据并行（DDP）](https://narrative.allen.ai/scaling-up-training#data-parallelism)的系统。它让专家常驻于 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 将多种将大型 MoE 分布到 GPU 集群上的技术，与让路由和计算更高效的优化结合在一起。 三种技术决定了模型及其训练状态如何分布在硬件上： - **专家并行**（https://narrative.allen.ai/scaling-up-training#expert-parallelism）将专家分散到各 GPU 上，每块 GPU 只保存完整专家池的一部分。 - **流水线并行**（https://narrative.allen.ai/scaling-up-training#pipeline-parallelism）将模型的层——即依次对输入进行变换的各个阶段——划分到不同的 GPU 组上，减少每块 GPU 需要保存在内存中的模型部分。 - **分布式优化器**将优化器状态——即训练过程中用于计算和应用更新的附加数据——分散到各 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-up-training)，了解数据并行、专家并行和流水线并行如何协同工作以扩展 MoE 训练——从单块 GPU 到多块 GPU。 ### **扩展到万亿参数量级** 我们在一系列 NVIDIA B300 GPU 配置下对 Olmo-core 3 进行了基准测试，包括一个 1.2 万亿参数的模型，在 512 块 GPU 上每个 token 激活 583.6 亿参数。观测到的最高吞吐量为每 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)。 ## 订阅获取 Ai2 最新动态的月度更新。