Ai2 发布 Olmo-core 3：面向万亿参数 MoE 的开源训练框架

10 月 1 日，Ai2 发布 Olmo-core 3——一个面向大规模混合专家（MoE）模型的开源训练框架，代码托管在。据，它是下一代 Olmo 的「核心系统之一」，设计目标是把 MoE 训练扩展到万亿参数规模。在 8 张 NVIDIA B300 GPU 上进行的一项初步测试中，47B 参数的 MoE 达到 52,000 token/s/GPU，约为旧实现（19,400 token/s/GPU）的 2.7 倍。

Ai2 在中把旧实现称为 Olmo-core 2。与发布同步公开的还有《Supercharging Olmo-core for Efficient and Scalable MoE Training》和一份。文中引用的性能数字均出自 Ai2 自行开展的测试。

专家池从 8 扩到 128，吞吐下降不到 5%

MoE 的做法是每个 token（语言模型处理文本的最小单位）只激活少数专家，用少量激活计算支撑更大的总参数容量。但 Ai2 在博客里写道，训练成本并不只取决于激活计算：权重存储、梯度同步、优化器更新都随总参数增长，把数据路由到正确的专家还要付出通信与协调开销。模型越大，这些成本越会侵蚀稀疏激活省下的那部分。技术报告把这类与容量挂钩的负担称为「容量税」（capacity tax），并记录了一个早期现象：在激活规模固定的前提下增加专家数量，吞吐反而下降。

新栈给出的结果不同：专家池从 8 个扩到 128 个，每个 token 仍只选择 4 个专家，激活参数保持在约 3.2B 不变；总参数容量从 4.6B 增至 47B，训练吞吐下降不到 5%。

博客说，在此前的 Olmo 系列里，OlmoE 用的是 64 个路由专家的 MoE 架构，Olmo 3 则是稠密架构；Olmo-core 3 加入的训练系统面向更大的 MoE 模型。

把专家留在 GPU 上，把数据送过去

旧实现的问题出在 FSDP。博客说，此前的 MoE 实现采用全分片数据并行（FSDP），会针对每个小批次训练数据收集并重新分片模型权重；X 帖文的说法则是 Olmo-core 2 反复收集权重。新版本改为基于分布式数据并行（DDP）的架构：每个 rank（训练进程）在整个梯度累积窗口内保留自己那份模型分片，每次优化器更新只同步一次梯度；分布式优化器再把 FP32 主权重与优化器状态分片到各 GPU。

专家并行是这套布局的核心：专家池被切分到不同 GPU，每张 GPU 只存一部分专家，路由时只把被选中的 token 行发往存放对应专家的 rank。节点内的搬运使用基于 NVSHMEM 的 rowwise EP（按行专家并行），token 行直接写入目标专家的输入缓冲区，省去派发（dispatch）与合并（combine）操作中的重排开销；路由元数据常驻 GPU，CPU 不必等它拷回就能继续排队任务；分组 GEMM（通用矩阵乘法）把路由产生的、形状不一的专家计算合并执行，无需把每个专家的批次填充到同一容量。

流水线并行把模型的层切分到多组 GPU，减少每张 GPU 需要常驻的模型量；显存仍然不够的配置用激活重计算换内存；报告里的一组控制实验显示，重计算去掉了近三分之二的激活内存、峰值内存降低约四分之一，代价是约五分之一的吞吐。检查点（checkpoint）按全局 FP32 张量记录，与并行布局解耦，换一套拓扑也能恢复训练。

MXFP8：吞吐提高约 21%，峰值内存从 103 GiB 降至 95 GiB

另一组控制实验专门量化 MXFP8（一种带共享块缩放的 8 位数值格式）的效果：在 4 张 NVIDIA B300 GPU 上、把工作量均匀分配到各专家，启用 MXFP8 后训练吞吐比 BF16（16 位浮点）基线高约 21%，峰值活动内存从 103 GiB 降到 95 GiB。博客指出，增益主要来自前馈计算与专家间的数据搬运，而不是注意力部分；报告补充说，FP32 主权重仍是权威副本，低精度只有在节省的计算与传输超过格式转换成本时才划算。

博客对这类取舍的总结是：「加速一部分训练可能在其他地方产生成本；更快的计算可能需要搬更多数据，而少搬比特如果让转换耗时太长，也未必划算。」

万亿参数级的测试，用的是随机路由

万亿参数级配置同样在 B300 上测。报告称，这套栈支持的训练配置从 12.9B 总参数到 1.2 万亿总参数，最多 512 张 GPU（NVL8 B300 节点）。其中 1.2 万亿参数配置每个 token 激活 58.36B 参数，跑在 512 张 GPU 上，观察到的最高吞吐为 858 TFLOP/s/GPU（启用 MXFP8 与逐层重计算）。

这些测试的口径需要说明：博客写明，万亿级测试使用随机路由，衡量的是系统性能而非训练出的模型质量；报告则把每个速率数字都标注为系统参考值（systems reference），不是「训练到某个质量需要多久」（time-to-quality）的结论。

报告里写下的失败模式：token gerrymandering

技术报告记录了一种失败模式：「token gerrymandering」。博客说明，一个本意是鼓励均衡路由的分数，可以在实际负载变得更不均衡时反而变好；报告称，路由器学会了钻负载均衡损失（load-balancing loss）的空子。

其他被记录的发现包括：因为专家处理的 token 变少而调低专家学习率，在被测模型族里没有改善结果；相同矩阵维度下，输入的数值不同也会让 GPU 计算耗时不同，因此性能对比需要同时匹配数值与形状；把通信与计算放在不同 GPU stream 上做重叠并不总是更快，有些测试里端到端反而更慢。

报告还列出了评估后没有采用的方案：把激活卸载到 CPU（主机链路承载不了相应流量）、两种与专家计算争抢 GPU 资源的重叠方案，以及一个需要完整训练运行时支撑的持久化 MoE megakernel 等。

开源的代码与下一代 Olmo

代码已经开放：GitHub 仓库把它描述为「OLMo 生态的 PyTorch 构建块」，可用 pip install ai2-olmo-core 安装，仓库里提供 OLMo-2 32B 与 OLMo-3 7B、32B 的官方训练脚本。仓库同时提醒，随附的 Docker 镜像只包含依赖、不含 Olmo-core 包本身，而且是在 Ai2 内部 H100 集群上定期测试的，换到硬件或驱动、CUDA 版本不同的环境可能无法直接使用。

X 帖文写道：「用 Olmo-core 3 训练你自己的 MoE——把框架适配到你的硬件，试验分布式训练。」博客则交代了下一步：下一代 Olmo 将采用 MoE 架构，Ai2 称目标是把它做成能力最强的一代 Olmo，用上迄今最大的数据集与最长的上下文窗口；Olmo-core 3 是「我们接下来所建一切的基础」。

而在公开结果里，这套栈规模最大的一次实验，是使用 DeepEP v2（处理专家间通信的替代后端）的 2.38 万亿总参数配置——Ai2 把它定义为短容量测试：演示的是可达到的规模，而不是持续训练的表现。