AI2发布开源MoE训练框架Olmo-core 3，支持万亿参数级扩展

10月1日，AI2正式发布Olmo-core 3，这是对其大语言模型开发框架的一次重大升级，核心是一个重新设计的开放混合专家（MoE）训练系统。该框架旨在将MoE训练扩展至万亿参数规模，同时保持计算效率，也是其下一代Olmo模型背后的核心系统之一。

在大模型训练成本高企的背景下，MoE模型通过仅激活部分参数来提升效率，但随着专家数量增长，通信和协调成本会侵蚀其计算优势。Olmo-core 3针对这一问题进行了重构。在一项基准测试中，专家池从8个增至128个，每个token仍只选择4个专家，总参数容量从4.6B增至47B，而训练吞吐量下降不到5%。该基础设施还已在超万亿总参数规模上完成测试。

技术实现上，Olmo-core 3从全分片数据并行（FSDP）转向基于分布式数据并行（DDP）的系统，让专家常驻GPU并直接路由数据，避免重复的参数收集。同时结合专家并行、流水线并行和分布式优化器，降低单卡内存压力；并通过行级专家并行、GPU驻留路由和分组GEMM等优化提升路由与计算效率。框架还支持MXFP8低精度格式，在B300 GPU上的受控测试中，端到端训练吞吐量较BF16提升约21%，峰值显存从103 GiB降至95 GiB。

在NVIDIA B300 GPU上的多组测试中，Olmo-core 3在512张GPU上实现了1.2万亿参数模型的运行，每个token激活583.6亿参数，最高观测吞吐量达858 TFLOP/s/GPU。此外，团队还通过DeepEP v2探索了2.38万亿总参数的配置。AI2表示，Olmo-core 3将用于训练下一代MoE架构的Olmo模型，相关技术报告和代码已开源。