Ai2 于 10 月 1 日发布 Olmo-core 3。这是一个重新设计了混合专家(MoE,每个输入只使用部分专家模块的模型结构)训练系统的框架。

Ai2 将其介绍为下一代 Olmo 的核心基础。这次发布不包含新训练的模型权重。

在一项测试中,每个 token 只选 4 个专家,同时把专家数从 8 个增加到 128 个。活跃参数保持在约 3.2B,总参数从 4.6B 增至 47B,据称训练吞吐量下降不到 5%。

此前的 MoE 实现使用完全分片数据并行(FSDP),每个小批次都要收集并重新分片权重。Olmo-core 3 改为基于分布式数据并行(DDP),让专家常驻 GPU,并把数据路由到对应专家。

在 Ai2 用 8 块 NVIDIA B300 进行的初步测试中,47B 参数的 MoE 每块 GPU 每秒处理 52,000 个 token,约为旧实现 19,400 个的 2.7 倍。

模型和训练状态通过专家并行、流水线并行和分布式优化器分布到多块 GPU 上。分布式优化器把用于计算更新的优化器状态分散存放,而不是在每块 GPU 上保存完整副本。

它还支持低精度格式 MXFP8。在 4 块 B300 的测试中,吞吐量比 BF16 高约 21%,峰值活跃内存从 103GiB 降到 95GiB。

代码以 Apache License 2.0 发布在 GitHub 的 Olmo-core 仓库中。

注意事项:区分系统测量与模型质量

Ai2 在 512 块 B300 上测试了每个 token 激活 58.36B 参数的 1.2 万亿参数模型,观测到最高 858 TFLOP/s/GPU。该测试为测量系统性能而使用随机路由,并不衡量训练后模型的质量。

借助 DeepEP v2 达到的 2.38 万亿参数配置是短时容量测试,而非完整训练。Ai2 也说明它展示的是可达到的规模,而不是持续训练性能。

所有数字都是 Ai2 的测量,本文没有复现。Ai2 自己也报告,在部分测试中让通信与计算重叠反而拖慢了端到端执行。

评估实现时应一并看的指标

如果要亲自评估,可以用目标路由、精度和网络搭建一个有代表性的小型训练任务,同时比较内存和有效吞吐量。

仅凭容量数字或单个内核的速度,很难判断完整训练是否可行。

Olmo-core 3 的万亿级数字展示的是开放训练基础设施的系统极限,下一代 Olmo 的质量仍是另一个问题。