Olmo-core 3 开源混合专家训练基础设施
TL;DR
Olmo-core 3 是一个新发布的开源训练堆栈，可在万亿参数规模下高效训练混合专家（MoE）模型，相比之前的实现，令牌吞吐量最高提升 2.7 倍，并支持 MXFP8 精度等高级优化。
Olmo‑core 3 是什么以及为何重要
Olmo‑core 3 是 Allen Institute 的 MoE 训练框架的重新设计，旨在在专家池从几十个扩展到数百个的同时，保持每个令牌的活跃参数数量大致恒定，从而维持计算效率。通过消除传统上侵蚀 MoE 优势的通信和路由开销，该堆栈使学术实验室和小型研究团队也能负担得起万亿参数的稀疏模型。
核心架构变更
从 FSDP 切换到 DDP
- 早期的 Olmo‑core 使用全分片数据并行（FSDP），每次小批量都会反复收集和重新共享模型权重。
- Olmo‑core 3 采用分布式数据并行（DDP），将专家权重保留在 GPU 上，仅路由必要的令牌数据。这消除了代价高昂的权重收集步骤。
集成的 MoE 堆栈 vs. Megatron‑Core
- NVIDIA 的 Megatron‑Core 仍是大型 MoE 的参考实现。
- Olmo‑core 3 提供了一个端到端、高度集成的 MoE 堆栈，性能优于之前的基于 FSDP 的版本。在八块 NVIDIA B300 GPU 上的基准测试中，一个 470 亿参数的 MoE 实现了 每 GPU 每秒 52,000 个令牌，而之前仅为 每 GPU 每秒 19,400 个令牌——吞吐量提升了 2.7 倍。
采用的扩展技术
并行策略
| 技术 | 目的 |
|---|---|
| 专家并行 | 将专家池分布在 GPU 上，使每个 GPU 仅存储部分专家。 |
| 流水线并行 | 将模型层拆分到 GPU 组中，降低每个 GPU 的内存占用。 |
| 分布式优化器 | 将优化器状态分布在 GPU 上，避免在每个设备上都进行完整复制。 |
这三种机制共同作用，使 MoE 模型能够扩展，而无需每个 GPU 都持有整个模型或优化器状态。
路由与计算优化
- 按行专家并行 – 直接将路由后的令牌写入专家输入缓冲区，减少数据重排开销。
- GPU 本地路由 – 将路由元数据保留在 GPU 上，避免 CPU 成为瓶颈。
- 分组 GEMM – 将来自不同专家的多个小矩阵乘法批量合并为单个更大的 GEMM 调用，提升 GPU 利用率。
使用 MXFP8 的精度优化
- MXFP8 是一种低精度数值格式，可减少计算量和 GPU 间通信流量。
- 在四块 NVIDIA B300 GPU 的受控基准测试中，启用 MXFP8 后，端到端吞吐量相比 BF16 基线提升了 约 21 %，同时峰值活跃内存从 103 GiB 降低至 95 GiB。
- 大部分性能提升来自更快的前馈计算和减少的专家间数据移动。
展示的规模
| 配置 | 总参数量 | 每令牌活跃参数量 | GPU 数量 | 峰值吞吐量 |
|---|---|---|---|---|
| 1.2 T 参数 MoE | 1.2 T | 58.36 B | 512 | 858 TFLOP/s / GPU |
| 2.38 T 参数（DeepEP v2） | 2.38 T | — | — | — |
1.2 T 参数的运行使用随机路由以压力测试系统性能；不代表训练后模型的质量。
超越原始速度的实证发现
- 令牌操纵 – 路由平衡得分可能提升，而实际工作负载平衡却恶化，凸显了指标设计的陷阱。
- 专家学习率缩放 – 在测试的模型族中，降低稀疏使用专家的学习率并未带来更好结果。
- 输入值依赖的运行时间 – 当输入值不同时，相同矩阵形状的 GPU 内核表现出不同的运行时间，强调了需要使用值匹配的基准测试。
- 通信与计算重叠 – 有时重叠这些数据流反而会减慢整体执行速度，表明并非所有情况下重叠越多越好。
这些观察结果在附带的技术报告中有详细说明，并为未来的 MoE 训练研究提供指导。
开源可用性与社区影响
- 代码 – 完全开源，托管于 GitHub： https://github.com/allenai/olmo-core
- 技术报告 – 提供深入的系统设计、消融实验和实验方法： https://allenai.org/papers/olmocore3
- 交互式演示 – 可视化讲解数据、专家和流水线并行： https://narrative.allen.ai/scaling-up-training
通过以开源许可证发布该堆栈，Allen Institute 使研究人员能够训练自己的万亿参数 MoE 模型，将系统适配到其他硬件，并实验新型路由或精度方案。
展望
Olmo‑core 3 将成为下一代 Olmo 模型的基石，这些模型将结合新的 MoE 架构、更大的数据集和更长的上下文窗口。该堆栈的模块化设计确保其能随着硬件进步而演进，使研究社区始终处于可扩展稀疏模型训练的前沿。
关键要点
- Olmo‑core 3 在保持每令牌活跃参数恒定的同时，将 MoE 训练扩展至万亿参数级别。
- 通过 DDP、专家/流水线并行和路由优化，其令牌吞吐量相比之前的基于 FSDP 的堆栈最高提升 2.7×。
- MXFP8 精度进一步带来 约 21 % 的速度提升并降低内存使用。
- 该框架完全开源，为社区提供了一个生产级工具，用于大规模稀疏模型研究。