Breaking the Training Barrier of Billion-Parameter Universal Machine Learning Interatomic Potentials

📄 arXiv: 2604.15821v1 📥 PDF

作者: Yuanchang Zhou, Hongyu Wang, Yiming Du, Yan Wang, Mingzhen Li, Siyu Hu, Xiangyu Zhang, Weijian Liu, Chen Wang, Zhuoqiang Guo, Long Wang, Jingde Bu, Yutong Lu, Guangming Tan, Weile Jia

分类: cs.DC, cs.LG

发布日期: 2026-04-17

备注: 11 pages, 8 figures

💡 一句话要点

提出MatRIS-MoE和Janus框架，加速十亿参数通用机器学习原子间势模型的训练。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 原子间势 机器学习 混合专家模型 分布式训练 Exascale计算 AI4S 材料科学

📋 核心要点

  1. 通用机器学习原子间势模型训练面临二阶导数计算和十亿参数规模带来的计算与通信挑战。
  2. 论文提出MatRIS-MoE模型和Janus分布式训练框架，利用混合专家模型和硬件感知优化加速训练。
  3. 实验表明，该方法在Exascale超算上实现了1.2/1.0 EFLOPS的峰值性能，并将训练时间从数周缩短至数小时。

📝 摘要（中文）

通用机器学习原子间势（uMLIPs）是在包含整个元素周期表的无机材料和有机分子的大规模多样化数据集上预训练的，作为量子精确物理模拟的基础模型。然而，uMLIP训练需要二阶导数，而目前缺乏相应的并行训练框架；此外，扩展到十亿参数级别会导致计算和通信开销的爆炸性增长，使其训练成为一项巨大的挑战。我们引入了MatRIS-MoE，这是一个建立在不变架构之上的十亿参数混合专家模型，以及{Janus}，这是一个开创性的高维分布式训练框架，用于具有硬件感知优化的uMLIPs。在两台Exascale超级计算机上部署后，我们的代码在单精度下达到了1.2/1.0 EFLOPS的峰值性能（理论峰值的24％/ {35.5％}），并行效率超过90％，将十亿参数uMLIP的训练从几周压缩到几小时。这项工作为Exascale上的AI-for-Science（AI4S）基础模型建立了一个新的高水位线，并为快速科学发现提供了重要的基础设施。

🔬 方法详解

问题定义：现有通用机器学习原子间势（uMLIPs）的训练，尤其是扩展到十亿参数级别时，面临着计算和通信开销的巨大挑战。主要痛点在于缺乏高效的并行训练框架来处理二阶导数计算，以及大规模模型带来的通信瓶颈。这限制了uMLIPs在实际科学研究中的应用。

核心思路：论文的核心思路是结合混合专家模型（MoE）和硬件感知的分布式训练框架，以降低计算复杂度并提高并行效率。MoE模型通过将计算分配给不同的专家网络，减少了每个设备的计算负担。硬件感知的分布式训练框架则通过优化数据分布和通信策略，最大限度地利用底层硬件的性能。

技术框架：整体框架包含两个主要部分：MatRIS-MoE模型和Janus分布式训练框架。MatRIS-MoE是一个基于不变架构的十亿参数MoE模型，用于学习原子间势。Janus是一个高维分布式训练框架，负责将训练任务分配到多个计算节点，并优化数据传输和梯度聚合过程。该框架针对Exascale超级计算机进行了优化。

关键创新：最重要的技术创新点在于Janus框架的硬件感知优化。该框架能够根据底层硬件的特性（如CPU、GPU、网络带宽等）动态调整数据分布和通信策略，从而最大限度地提高并行效率。此外，MatRIS-MoE模型采用不变架构，能够更好地处理原子间相互作用的对称性。

关键设计：MatRIS-MoE模型采用了混合专家网络结构，每个专家网络学习不同的原子间相互作用模式。Janus框架的关键设计包括：1) 一种高效的数据划分策略，能够将大规模数据集均匀地分配到各个计算节点；2) 一种优化的梯度聚合算法，能够减少通信开销；3) 一种动态负载均衡机制，能够根据各个节点的计算能力调整任务分配。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

该研究在两台Exascale超级计算机上实现了显著的性能提升。在单精度下，代码达到了1.2/1.0 EFLOPS的峰值性能，分别占理论峰值的24％和35.5％。并行效率超过90％，成功将十亿参数uMLIP的训练时间从数周压缩到数小时。这些结果表明，该方法能够有效解决大规模原子间势模型训练的计算瓶颈。

🎯 应用场景

该研究成果可广泛应用于材料科学、化学、生物物理等领域，加速新材料的发现和分子动力学模拟。通过高效训练通用原子间势模型，研究人员可以更快速、更准确地预测材料的性质和行为，从而推动相关领域的科学研究和技术创新。未来，该框架有望扩展到其他科学计算领域，例如气候模拟和药物发现。

📄 摘要（原文）

Universal Machine Learning Interatomic Potentials (uMLIPs), pre-trained on massively diverse datasets encompassing inorganic materials and organic molecules across the entire periodic table, serve as foundational models for quantum-accurate physical simulations. However, uMLIP training requires second-order derivatives, which lack corresponding parallel training frameworks; moreover, scaling to the billion-parameter regime causes explosive growth in computation and communication overhead, making its training a tremendous challenge. We introduce MatRIS-MoE, a billion-parameter Mixture-of-Experts model built upon invariant architecture, and {Janus}, a pioneering high-dimensional distributed training framework for uMLIPs with hardware-aware optimizations. Deployed across two Exascale supercomputers, our code attains a peak performance of 1.2/1.0 EFLOPS (24\%/{35.5\%} of theoretical peak) in single precision at over 90\% parallel efficiency, compressing the training of billion-parameter uMLIPs from weeks to hours. This work establishes a new high-water mark for AI-for-Science (AI4S) foundation models at Exascale and provides essential infrastructure for rapid scientific discovery.