Introducing Olmo-core 3: 为大规模 MoE 设计的开放、可扩展训练基础设施

📌 One-Sentence Summary

Olmo-core 3 是一个开源框架，它重新设计了 Mixture-of-Experts (MoE) 训练基础设施，在保持计算效率的同时将规模扩展至万亿参数级别，吞吐量较前代提升 2.7 倍，并随附一份记录关键系统级发现的技术报告。

📝 Summary

Olmo-core 3 是 Allen AI 用于开发大语言模型开放框架的一次重大升级。它采用重新设计的开放 MoE 训练系统，旨在将 MoE 训练扩展到万亿参数规模且不牺牲计算效率。该框架用分布数据并行 (DDP) 系统取代了早期的 FSDP 实现，使专家模型常驻 GPU 并将相关数据路由至专家，从而避免了重复的权重聚合。在 8 块 NVIDIA B300 GPU 的基准测试中，一个 470 亿参数的 MoE 每块 GPU 每秒可处理 52，000 个 token，约为之前实现方案 (19，400 tokens/sec/GPU) 的 2.7 倍。系统结合了专家并行、流水线并行和分布优化器三种并行技术，无需每块 GPU 存储完整模型和训练状态即可实现扩展。其他优化包括行式专家并行、GPU 常驻路由、grouped GEMM 以及 MXFP8 支持，后者在将峰值活跃内存从 103 GiB 降低至 95 GiB 的同时，吞吐量比 BF16 高出 21%。该基础设施已在总参数量超过一万亿的规模上通过基准测试，包括在 512 块 GPU 上运行的 1.2 万亿参数模型（峰值吞吐量 858 TFLOP/s/GPU）以及使用 DeepEP v2 的 2.38 万亿参数配置。随附的技术报告记录了一些非直观的系统行为，例如「token gerrymandering」（即路由平衡分数提升但实际工作负载反而失衡），以及发现将通信与计算在独立 GPU 流中重叠并不总能提高吞吐量。Olmo-core 3 将作为下一代 Olmo 模型（将采用 MoE 架构）的基础，并完全向研究人员和开发人员开放，支持使用、适配和扩展。

💡 Main Points

Olmo-core 3 使用分布数据并行 (DDP) 系统取代了早期的 FSDP-based MoE 训练实现，使专家模型常驻 GPU 并引导数据流向，实现了 2.7 倍的吞吐量提升。

在 8 块 NVIDIA B300 GPU 的基准测试中，采用新栈的 470 亿参数 MoE 每块 GPU 每秒处理 52，000 个 token，而早期实现仅为 19，400 个。这避免了 FSDP 重复权重聚合和重新分片的开销。

通过专家并行、流水线并行和分布优化器三种并行技术，实现了 MoE 的扩展，而无需每块 GPU 存储完整模型和训练状态。

专家并行将专家分布在不同 GPU 上，流水线并行将模型层跨 GPU 组拆分，分布优化器则将优化器状态分散在各 GPU 中。这些技术共同作用，使 MoE 能扩展至万亿参数规模，同时保持单块 GPU 的内存需求在可控范围内。

包括行式专家并行、GPU 常驻路由、grouped GEMM 和 MXFP8 支持在内的路由与计算优化，提升了训练效率。

行式专家并行将路由数据直接放入专家输入缓冲区，GPU 常驻路由将元数据保留在 GPU 中以避免 CPU-GPU 同步，grouped GEMM 将小型专家计算合并以提高 GPU 利用率，而 MXFP8 将峰值活跃内存从 103 GiB 降至 95 GiB，且吞吐量比 BF16 高 21%。

系统已在总参数量超过一万亿的规模上完成基准测试，包括在 512 块 GPU 上的 1.2 万亿参数模型和使用 DeepEP v2 的 2.38 万亿参数配置。

1.2 万亿参数配置在每个 token 激活 583.6 亿参数的情况下，达到了 858 TFLOP/s/GPU 的峰值吞吐量。这些测试使用随机路由来衡量系统性能而非训练模型质量，而 2.38T 配置是一次短期容量测试，旨在证明扩展能力而非持续训练性能。

技术报告记录了非直观的系统行为，包括「token gerrymandering」以及通信-计算重叠的非单调影响。

「token gerrymandering」描述了路由平衡分数提升但实际工作负载反而更不平衡的现象。此外，由于处理 token 减少而降低专家的学习率并未改善结果，GPU 计算时间随输入值（而非仅随形状）而变化，且在独立 GPU 流上重叠通信与计算并不总能提高吞吐量。

💬 Key Quotes

Olmo-core 3 旨在将 MoE 训练扩展到万亿参数范围，同时保持计算效率。

在 8 块 NVIDIA B300 GPU 的初步测试中，一个 470 亿参数的 MoE 使用新栈每块 GPU 每秒处理 52，000 个 token，而使用我们早期实现则为 19，400 个——吞吐量大约提升了 2.7 倍。

在系统中帮助最大的部分启用 MXFP8 后，训练吞吐量比我们作为基准的高精度格式 BF16 高出约 21%，同时峰值活跃内存从 103 GiB 降至 95 GiB。

旨在鼓励平衡路由的分数可能会在实际工作负载变得更不平衡时反而提高。我们将这种失效称为「token gerrymandering」。

在独立的 GPU 流上重叠通信和计算并不总是能加快训练速度。在某些测试中，它反而减慢了端到端执行速度——这提醒我们，更多的重叠并不必然意味着更高的吞吐量。

📊 Article Meta

AI Screening： 92

Source： Hugging Face - Blog

Author： Kyle Wiggers

Category： 人工智能

Language： 英文

Read Time： 6 min

Word Count： 1349

Tags：

Mixture of Experts , 模型训练与推理 , 开源 , AI 基础设施 , LLM