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

TL;DR

This paper presents MatRIS-MoE and Janus, enabling efficient billion-parameter uMLIP training on supercomputers, significantly reducing training time and setting new standards for AI foundation models in scientific research.

Contribution

Introduction of MatRIS-MoE and Janus frameworks that facilitate scalable, hardware-aware training of billion-parameter uMLIPs on Exascale supercomputers.

Findings

Achieved 1.2/1.0 EFLOPS peak performance with over 90% efficiency.

Reduced training time of billion-parameter uMLIPs from weeks to hours.

Set new benchmarks for AI4S foundation models at Exascale.

Abstract

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.