位于西雅图的艾伦人工智能研究所(Ai2)10月2日发布大语言模型开发框架 Olmo-core 3。据 SiliconANGLE 报道,该框架改进了混合专家(MoE)大模型的训练方式,可在保持计算效率、压低成本的前提下把训练规模推到万亿参数级别。
混合专家模型与稠密模型的运行方式不同。稠密模型在生成内容时会激活整个模型,而混合专家模型把计算分散到模型中专司其职的各个部分,每生成一个 token 只调用其中一小部分。也就是说,一个 MoE 模型的总参数量可以远大于它实际参与计算的参数量。这种机制降低了训练与推理的总体算力消耗,但完整模型仍需分布存放在多张 GPU 的显存中,训练期间专家之间的网络协调也会产生额外开销。
Ai2 表示,Olmo-core 3 意在弥合稠密模型与混合专家模型之间的差距,它允许专家池从 8 个扩展到 128 个,同时每个 token 仍然只选择 4 个专家。使用同一套基础设施,模型可以扩展到超过一万亿个参数。
在基准测试中,Ai2 称 Olmo-core 3 在英伟达 B3000 GPU 上训练一个 470 亿参数的模型时,每秒可处理 52,000 个 token。相比之下,英伟达 Megatron-core 训练架构在同一指标上约为每秒 19,400 个 token,新框架的吞吐量约为其 2.7 倍。Megatron-core 是目前训练大型混合专家模型的常用方案。
Ai2 在一份关于该项目的白皮书中说明,新架构采用专家并行,把专家分散到多张 GPU 上,每张卡只需存放完整专家池的一部分;它还把模型的各个层切分到不同的 GPU 组,以减少每张 GPU 必须常驻内存的模型部分;分布式优化器则把优化器状态分散到多张 GPU,而不是在每张卡上保存完整副本。这些做法合起来降低了模型规模扩大时的内存开销,因为整个模型及其训练状态不必一次性全部装进内存。此外,框架支持 MXFP8 这种用更少比特表示部分数值的格式,可减少计算量以及 GPU 之间的数据传输量。
Ai2 称,万亿参数级别的模型通常超出那些无法接触国家级或企业级基础设施的研究者的能力范围,新框架还允许研究者针对不同硬件调整混合专家训练,并对路由、并行等系统环节进行实验,以培育新的生态。该项目及相关系统已在 GitHub 上面向开发者和开源社区开放。