Ai2 Olmo-core 3：MoE训练吞吐量提升2.7倍的秘诀

一次针对稀疏模型训练瓶颈的系统级重构

Allen Institute for AI（Ai2）发布了 Olmo-core 3，这是其用于训练 Olmo 系列语言模型的开源训练框架。此次发布以 Apache 2.0 许可证开源，并附带了完整的技术报告。核心变化只有一个方向：把混合专家（MoE）模型训练中那些随总参数量膨胀、却不带来等量有用计算的开销，从系统层面拆掉。

据技术报告披露，Olmo-core 3 记录的可运行配置点覆盖范围从 16 块 GPU 上的 129 亿参数模型，一直延伸到 512 块 GPU 上的 1.2 万亿参数模型，并配有官方预训练脚本和导出到 Hugging Face 格式的检查点转换工具。

容量税：稀疏模型的结构性负担

技术报告给这类开销起了一个清晰的名字：容量税（capacity tax）。它指的是跟随总存储参数量产生的成本——权重存储、梯度同步、优化器状态更新——即便每个 token 实际激活的参数保持不变，这些成本仍随总参数量线性增长。

稠密模型不会遇到这个问题，因为更多参数意味着更多有效计算，成本与收益同步。稀疏模型打破了这个绑定关系。早期的完全分片数据并行（FSDP）实现在每个 micro-batch 都要 gather 和 reshard 权重，结果是：在固定激活规模下增加专家数量，吞吐量必然下降。

Olmo-core 3 的方案是切换到分布式数据并行（DDP）：每个 rank 在整个累积窗口内持有自己的模型分区，每个优化器步骤只同步一次梯度，并由分布式优化器对 FP32 主权重和优化器状态做分片。在此基础上，专家并行将路由专家分散到各 rank，只把被选中的 token 行通过 NVSHMEM 直接写入目标 rank 的缓冲区。放不下的层由流水线并行覆盖。

混合专家（MoE）模型的核心思路是把前馈网络拆成多个并行的"专家"子网络，每个 token 经路由函数只激活其中少数几个（例如从 128 个专家中选 4 个）。这样总参数量可以极大膨胀，但单个 token 的实际计算量（即激活参数）保持恒定。问题在于，权重存储、梯度同步和优化器状态这些成本是按总参数量而非激活参数计算的，因此总参数量越大，这部分与计算无关的开销就越高——这正是"容量税"的由来。稀疏模型因此与稠密模型有本质区别：稠密模型参数量与有效计算同步增长，稀疏模型则把两者解耦，从而把负担转移到了系统实现层面。

完全分片数据并行（FSDP） 与 分布式数据并行（DDP） 的关键差异在于权重同步时机。FSDP 把模型参数切分到各 GPU，但在每次前向和反向传播前都要把权重 gather 回来、用完再 reshard，对 MoE 而言，micro-batch 越多，这种 gather/reshard 往返就越频繁，成为吞吐量瓶颈。DDP 则让每个 rank 在整个梯度累积窗口内持有自己的模型分区，只在每个优化器步骤同步一次梯度，再由分布式优化器对 FP32 主权重和优化器状态做分片，从而摊薄了同步开销。

专家池扩展的实测数据

这套重构要回答的核心问题是：专家池能扩多大而不吃掉吞吐量？

Ai2 给出了两组关键测量。第一组中，专家池从 8 扩展到 128 个专家，每个 token 只选择 4 个专家，激活参数维持在约 32 亿不变。总参数容量从 46 亿增长到 470 亿——扩大了十倍——而训练吞吐量下降不到 5%。

第二组对比在 8 块 NVIDIA B300 GPU 上进行：一个 470 亿参数的 MoE 模型在新框架上达到每 GPU 每秒约 52,000 tokens，而旧实现只有约 19,400，吞吐量提升约 2.7 倍。Ai2 将这次运行标注为"初步结果"。

路由不只是模型问题，更是系统问题

报告将路由视为一个工程界面而非建模细节。其中一节深入分析了负载均衡损失函数到底在优化什么，并命名了一种称为"Token Gerrymandering"的失效模式：损失函数数值在改善，但专家间的不平衡反而在加剧。任何调过路由均匀性的团队都遇到过类似的结果。

另一个发现值得注意：报告明确记录了团队尝试后放弃的优化方案。比如将通信与计算重叠——直觉上这应该隐藏通信开销——实际上使被重叠的 kernel 变慢的幅度超过了它隐藏的开销，因为两者在争夺同一批 GPU 资源。

两个工程细节对阅读代码的团队有直接价值。路由元数据保持在 GPU 上，省去了旧路径为 all-to-all 分割尺寸和专家矩阵乘法组尺寸所做的 host 往返。设备端调度的分组 GEMM 则处理路由产生的每个专家不均匀形状，报告同时测量了这与将所有专家批次 padding 到统一容量相比孰优孰劣。

NVSHMEM 是 NVIDIA 提供的 Partitioned Global Address Space（PGAS）编程模型库，允许 GPU kernel 直接读写远端 GPU 的显存，而无需先经过主机侧的 MPI 通信。这在 MoE 的专家并行中至关重要：路由决定哪些 token 要送到哪个专家所在的 rank，Olmo-core 3 因此把被选中的 token 行通过 NVSHMEM 直接写入目标 rank 的缓冲区，避免了传统的 host 往返和显式 all-to-all 调用。类似的，分组 GEMM（grouped GEMM） 用于处理 MoE 路由后每个专家收到的 token 数量不均的问题——它把不同形状的小矩阵乘法打包成一个 kernel 批量执行，设备端调度则进一步省去了为对齐形状而做 padding 的额外计算和显存占用。

吞吐量数字的边界条件

报告对这些数字的解读条件说得很清楚，引用时需要留意。测量配置在批量大小、并行策略、精度和重计算策略上各不相同，因此作者将其描述为一组"运行点"而非受控的扩展曲线。每个头条数字都是在随机路由条件下测得的系统参考值。

精度选择直接影响经济性。MXFP8——一种带共享块缩放的八位格式——减少了 GEMM 时间、保存的激活和 dispatch 负载，同时保留了一份权威的 FP32 主权重。在 1.2 万亿参数配置下，配合 MXFP8 和逐层重计算，达到了每 GPU 858 useful-model TFLOP/s。可选的 DeepEP v2 后端在实验性容量测试中触及 2.38 万亿参数。

MXFP8 是 Open Compute Project 定义的八位浮点微缩放格式，其特点是每 32 个元素共享一个缩放因子（块缩放），而非整个张量共用一个比例。相比传统的 per-tensor FP8，MXFP8 能在保持数值稳定性的同时，对 GEMM 计算、激活保存和专家 dispatch 负载都带来更细粒度的压缩。Olmo-core 3 用它来加速矩阵乘法并减少显存占用，同时保留一份权威的 FP32 主权重用于优化器更新，以确保精度不因量化而退化。逐层重计算（layer-wise recomputation） 则是在反向传播时丢弃部分中间激活，在需要时重新计算，以显存换计算，是大型模型装入有限 GPU 内存的常用手段。

代价在哪里

吞吐量的收益并非没有成本。激活重计算减少了近三分之二的激活内存和约四分之一的峰值内存，代价是约五分之一的吞吐量。这是大型配置为了装入设备内存而做出的取舍。

拓扑无关的检查点则保持了运行在不同并行布局间的可移植性——训练任务可以在不同节点数上恢复，无需重写检查点格式。对于自行维护 GPU 集群的团队来说，这降低的是运维层面的摩擦，而非计算层面的成本。

Olmo-core 3 的发布本身就是一个信号：MoE 训练的竞争正在从"模型架构设计"转向"系统实现效率"。当专家池扩大十倍而吞吐量几乎不变时，稀疏模型的规模化路径就不再被基础设施卡住了。

拓扑无关的检查点（topology-agnostic checkpointing） 指的是检查点格式不绑定于训练时的并行布局——例如原本在 512 块 GPU 上用特定张量并行和流水线并行切分训练的模型，可以在不同节点数、不同并行策略下直接恢复。传统检查点往往记录了参数在特定 rank 上的分布，换布局就需要重写格式或做离线转换，而拓扑无关方案把参数按逻辑结构存储，恢复时动态映射到新布局。这使团队能在集群故障、扩缩容或更换硬件时快速重启训练任务，降低的是运维层面的摩擦，而非计算本身的成本。