Breaking the Training Barrier of Billion-Parameter Universal Machine Learning Interatomic Potentials 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∗ and Weile Jia∗ ∗ State Key Lab of Processors, Institute of Computing Technology, Chinese Academy of Sciences
arXiv:2604.15821v1 [cs.DC] 17 Apr 2026
† University of Chinese Academy of Sciences ‡ Independent Researcher § Sun Yat-Sen University
III. OVERVIEW O F T HE P ROBLEM
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 quantumaccurate 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-ofExperts 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. Index Terms—AI for Science, Machine Learning Interatomic Potential, Exascale, High Performance Computing
A. Molecular Dynamics and MLIPs Molecular dynamics (MD) simulations with quantum accuracy are an indispensable workhorse for modern scientific discovery, enabling breakthroughs in advanced materials [1], clean energy [2], and pharmaceuticals [3] by unveiling atomiclevel mechanisms. Driven by AI for Science (AI4S), the field is now shifting from computationally expensive Kohn-Sham Density Functional Theory (KS-DFT) to highly efficient surrogate modeling. Recent advancements in specialized MLIPs [4], [5], such as DeePMD [6], have addressed the computational bottlenecks by training on ab initio data to reach unprecedented spatio-temporal scales. Nevertheless, such specialized MLIPs often suffer from poor generalization and an inability to represent the elemental diversity of multi-component physical systems. This limitation stems from their restricted model capacity, which is typically limited to tens of thousands of parameters [4], [5]. The growing complexity of scientific problems—such as high-throughput screening of solid-state electrolytes and iron-based catalysts—has driven the transition to universal MLIPs [7], [8]. These challenges involve heterogeneous interfaces and intricate surface reactions governed by highdimensional potential energy surfaces (PES) that span diverse chemical elements and chemical domains. To capture such complexity while preserving physical symmetries, the stateof-the-art uMLIPs are predominantly formulated as Graph Neural Networks (GNNs) [7], [9], [10]. For each atom i, the local environment within a cutoff radius is represented as a graph, with atoms as nodes and bonds as edges. Iterative message passing enables local representations to incorporate information from increasingly distant atoms, capturing multibody correlations and effective long-range interactions while maintaining linear scaling with system size. Driven by the explosive growth of large-scale open datasets [8], [11]–[16], such as OMat24 [17] and OMol25 [18], uMLIPs have rapidly scaled from millions to the billionparameter regime (Table I). Consequently, the primary computational bottleneck has shifted entirely to the training phase, making the optimization of billion-parameter uMLIPs a formidable HPC challenge. Three fundamental factors exacerbate this difficulty: a). Force-matching via second-
I. J USTIFICATION F OR P RIZE Record training of billion-parameter uMLIPs with 473 million atomic configurations. Single-precision performance reaches 1.2/1.0 EFLOPS, maintaining 90%–93% parallel efficiency on two Exascale systems. For the 11.5B model, normalized throughput is 653–3201× compared to the current stateof-the-art, establishing a new performance high watermark for training atomic AI4S models on Exascale systems. II. P ERFORMANCE ATTRIBUTES Performance attribute
Our submission
Category of achievement Type of method used Results reported on basis of Precision reported System scale Measurements
Scalability, Throughput Other: training of uMLIP Whole application with IO Single precision Measured on full system Timers, FLOP count
# These authors contributed equally to this work. Corresponding authors: Weile Jia ([email protected])
1
order derivatives: Unlike the first-order optimizations used in Large Language Models (LLMs), uMLIP training requires strict force-matching (Fi = −∂E/∂Xi ), necessitating doublebackward (second-order automatic differentiation). This doubles both the computational costs and memory footprints from exploding intermediate activations. b). High-precision (FP32) requirement: While LLMs routinely exploit low-precision formats (e.g., FP16 or FP8), uMLIPs must use full singleprecision (FP32) arithmetic to preserve the quantum-accurate fidelity essential for stable molecular dynamics. This constraint doubles the memory footprint and severely limits tensor-core utilization on many-core architectures. c). Extreme edge-token throughput demand: Message passing along edges dominates the computational cost, making the edges the fundamental “tokens” of uMLIP training. The 473 million configurations used in this work yield approximately 3.6 trillion interacting edges, imposing unprecedented demands on system-level throughput.
(a)
(b)
Crystal
Atomic Embed.
Add & Norm. Graph Refinement
𝑎!𝑗𝑖𝑘
Add & Norm.
Interaction Block (×𝑻)
Separable Attention
Concat MLP
⊙
Agg.
Agg.
Concat MLP
(d)
Concat MLP
𝑘 𝑖
𝑣𝑖! 𝐺(𝑣!𝑖 , 𝑒!𝑖𝑗 , 𝑎!𝑗𝑖𝑘 )
𝑣𝑡𝑗
Linear
⊙
𝑟𝑐𝑢𝑡
𝑗
𝑣𝑡𝑖
Dim-wise Softmax
Linear
Feat. Embed. Pair-wise Three-body Embed. Embed.
Periodic Repetition
𝑒!𝑖𝑗
𝑒𝑡𝑖𝑗
(c)
Envelope
Readout Block
Agg.
E E, 𝑥𝑖 , 𝜖
M Autograd
𝐹! , 𝜎
MLP
MLP
𝑣𝑡!" 𝑖
𝑒𝑡!" 𝑖𝑗
Fig. 1. Overview of the MatRIS architecture. (a) Graph construction under periodic repetition condition. (b) Overall workflow, including feature embedding, interaction blocks, and readout. (c) Graph separable attention for invariant message passing. (d) Graph Refinement module.
B. MatRIS
to enable multi-task learning across heterogeneous chemical domains and datasets (Sec. V-A).
Materials Representation and Interaction Simulation (MatRIS) is an invariant, attention-based uMLIP designed for highly efficient atomistic modeling. As illustrated in Fig. 1(a), ➊ an atomistic system is represented as G(Z, X, L), where Z, X, and L denote atomic numbers, Cartesian coordinates, and lattice vectors, respectively. Under periodic boundary conditions, the structure is first periodically repeated and then converted by Graph Generation Module into the graph structure G(vi0 , e0ij , a0ijk ). Fig. 1(b) shows the architecture of MatRIS. ➋ In the Embedding Module, atomic numbers, pairwise distances, and angles are further embedded as initial node, edge, and angular features. ➌ MatRIS adopts a message-passing mechanism with L Interaction Blocks (IBs). In IBt , {vit , etij , atijk } are iteratively transformed to t+1 {vit+1 , et+1 ij , aijk }, where t = 1, . . . , T − 1, T is the number of IBs. Each IB consists of a Graph Separable Attention module (Fig. 1(c)) and a Refinement module (Fig. 1(d)). In the Separable Attention module, two linear projections of eij are normalized by dim-wise softmax over the target and source neighborhoods to obtain the score of target and source atoms, and then the weighted messages are aggregated ′ separately and concatenated to update the node feature vi . By replacing dense pairwise attention with source- and targetaware neighborhood aggregation, separable attention reduces the attention complexity from O(N 2 ) to O(N ). Finally, ➍ the Readout Block predicts total energy E and magnetic moments M . Forces and stress are obtained by automatic differentiation, ∂E and σ = V1 ∂E Fi = − ∂X ∂ϵ , where V denotes volume. i Extensive benchmarking shows that the 10-millionparameter MatRIS establishes a new Pareto frontier in both accuracy and training efficiency. On the Matbench-Discovery benchmark, it achieves a record F1 score of 0.847 and an RMSD of 0.0717 while delivering a 13× training speedup over state-of-the-art uMLIPs [10]. The model further demonstrates strong zero-shot generalization across diverse downstream tasks [10]. In this work, we extend MatRIS to MatRIS-MoE by incorporating the Mixture-of-Experts (MoE) architecture
IV. C URRENT S TATE O F T HE A RT A. Scaling of the uMLIPs The ultimate goal of uMLIPs is to enable zero-shot, ab initio-level accurate simulations across the entire periodic table, thereby unifying traditionally isolated domains such as organic molecules, crystalline solids, and catalytic interfaces. However, enforcing 3D spatial symmetry—a fundamental physical requirement—imposes significant architectural tradeoffs. Equivariant models (e.g., eqV2 [24], MACE [25]) natively satisfy these symmetries via high-degree irreducible representations; yet their reliance on Clebsch-Gordan tensor products introduces substantial computational complexity and memory overhead. Conversely, unconstrained models (e.g., ORB [21], PET [22]) implicitly learn symmetries through extensive data augmentation, circumventing complex tensor operations. While this approach enhances inference efficiency, it introduces potential thermodynamic inconsistencies, such as non-conservative forces during long-timescale molecular dynamics simulations. Invariant models (e.g., CHGNet [9], DPA3 [26], and MatRIS [10]) ensure rotational and translational invariance by utilizing distances and angles. Although computationally efficient, invariant models have not yet been scaled to the billion-parameter regime using MoE architectures. Matching the fidelity of billion-parameter equivariant models in capturing complex, high-order multi-body interactions therefore requires a significantly larger parameter space. To bridge this accuracy-efficiency gap and support diverse chemical domains, the community has rapidly increased both dataset size and model capacity, driving an exponential rise in computational demands. As detailed in Table I, the hardware configurations required to train representative uMLIPs have escalated from a single GPU to large-scale clusters. In 2023, training the lightweight CHGNet (0.41M parameters) required 8.3 days on a single A100 GPU. By 2024, training the 86M-parameter eqV2 model necessitated 6 days across a
2
TABLE I C OMPARISON OF STATE - OF - THE - ART UNIVERSAL M ACHINE L EARNING I NTERATOMIC P OTENTIALS ( U MLIP S ) PUBLISHED BETWEEN 2022 AND 2026. Work
Year
Category
Multi-task
Dataset
M3GNet [19] CHGNet [9] GNoME [20] eqV2 [17] ORB [21] MatRIS-L [10] PET [22] UMA [7] MACE-mh [23] SevenNet-Omni [8] This work (M) This work (L) This work (M) This work (L)
2022 2023 2023 2024 2025 2026 2026 2025 2025 2026 2026 2026 2026 2026
Invariant Invariant Equivariant Equivariant Unconstrained Invariant Unconstrained Equivariant Equivariant Equivariant Invariant Invariant Invariant Invariant
No No No No No No No Yes Yes Yes Yes Yes Yes Yes
0.19M 1.58M 89.0M 113M 133M 113M 113M 459M 117M 243M 473M 473M 473M 473M
#Total Params 0.23M 0.41M 16.2M 86M 25.5M 10.4M 730M 1.4B 6.4M 54.9M 2.47B 11.50B 2.47B 11.50B
#Active Params 0.23M 0.41M 16.2M 86M 25.5M 10.4M 730M 50M 6.4M 54.9M 0.56B 2.89B 0.56B 2.89B
Hardware RTX3090 A100 TPUv3 A100 A100 A800 H100 H200 H100 H200 GPGPU GPGPU ARMv9 ARMv9
#GPUs/ #Cores 1 1 4 64 8 128 512 256 48 8 45K 45K 12.4M 12.4M
Peak
FLOPS
– – – – – – – – – – 25.44% 35.52% 17.53% 24.41%
– – – – – – – – – – 750.9P 1048P 861.6P 1200P
Norm. Throughput [UMA=1] – 0.0022 – 3.09 – 0.565 – 1.00 – 0.706 653.4 2795.9 749.9 3201.8
* This table details model capacity (total and active parameters), multi-task learning capability, dataset size, training hardware configuration, and core
performance metrics. Peak performance (as a percentage of the theoretical peak) and FLOPS are reported for the training of the uMLIP models. Normalized Throughput is calculated via #Active Params×(Dataset Size×Epochs)/Training Days. The throughput of UMA is 1,050, which is normalized to 1.0. The symbol “–” denotes data not reported in the respective literature. This work reports the training of the medium (M, 2.47B) and large (L, 11.50B) variants of the MatRIS-MoE.
utilization of Tensor Cores on modern accelerators, necessitating a shift toward more compute-dense operators such as standard self-attention. (2) Lack of second-order parallel training framework: As active parameters reach the billion scale, a single GPU can no longer accommodate the uMLIP model. Consequently, pure data parallelism is no longer applicable, necessitating the development of a second-order parallel training framework. (3) Bottlenecks on large-scale platforms: Training billion-parameter uMLIPs on large-scale many-core supercomputers (e.g., GPU-based and ARM-based Exascale systems) significantly increases cross-node communication volume and exposes the inefficiency of second-order operators, limiting overall parallel scalability.
cluster of 64 A100 GPUs. In 2025, the 1.4B-parameter UMA model utilized 256 H200 GPUs for 21 continuous days. Even with recent architectural advancements, hardware requirements remain massive; for instance, the 730M-parameter PET model requires the parallel scheduling of 512 H100 GPUs. Despite this rapid development, scaling uMLIPs into the billion-parameter regime confronts formidable barriers. First, current uMLIP training relies predominantly on pure data parallelism. As active parameters reach the billion scale, the memory footprint inherently exceeds the capacity of a single GPU. Moreover, state-of-the-art frameworks optimized for first-order LLMs are fundamentally incompatible with the second-order computational graphs required for uMLIP force-matching. This lack of native model parallelism effectively stalls the development of larger models. Second, a tenfold increase in active model parameters will dramatically extend training time on conventional clusters of only hundreds of GPUs, rendering brute-force scaling economically unaffordable. Finally, while algorithmic innovations have advanced rapidly, system-level hardware–software co-design has significantly lagged behind, resulting in poor utilization of modern HPC and severely limiting overall training efficiency.
V. I NNOVATIONS R EALIZED A. Algorithmic Innovation Extending the single-task MatRIS to a multi-task model that jointly supports isolated molecules, periodic crystals, catalytic systems, MOFs, and periodic molecules presents three core challenges: resolving dataset inconsistencies (e.g., varying DFT functionals), balancing massive model capacity with computational feasibility, and fully utilizing GPU hardware for modelling complex many-body interactions. We address these challenges through the following architectural and training innovations. 1) Task-aware feature embedding: To align the systems that are calculated under different density functionals (e.g., PBE, R2SCAN, ωB97M) within a shared representation space, MatRIS-MoE augments the feature embedding with three additional components: a task embedding that injects dataset embeddings, charge and spin embeddings, and a global feature embedding to distinguish element composition. The computational cost of the embedding module remains strictly linear in the graph size. 2) Sparse mixture-of-experts (MoE) routing: MatRIS-MoE inserts sparse expert modules both before and after the selfattention layer. A message-update MoE to specialize in mes-
B. Scaling MatRIS to Billion-parameter Regime Building on the 10-million-parameter MatRIS—an invariant SOTA uMLIP originally trained with pure data parallelism on 128 A800 GPUs—we scale the architecture to the billionparameter regime to create a true multi-task foundation model. This extension enables joint learning across heterogeneous domains (molecules, materials, catalysis, MOFs, and direct air capture) on Exascale systems. The scaling process has three major challenges: (1) Architectural adaptation for multi-task learning: Expanding the parameter space requires integrating a Mixture-of-Experts (MoE) architecture to handle heterogeneous chemical rules across tasks. In addition, although the original O(N ) separable attention is theoretically efficient, its fragmented and memory-bound operations result in low
3
(a) MatRIS-MoE Molecular
Three-body Convolution
Three-body Embed.
" 𝑎&!'
Linear
Pair-wise Embed.
" 𝑒!&
Linear
Crystal
Atomic Embed.
Molecular Crystal
"($ 𝑎&!'
Linear
MLP
Linear
Expert-1
Router (top-k)
Global Feat. Embed.
⨂ ⨁
Expert-2 …
Linear
Charge&Spin Embed. Catalysis
Expert-S
Linear
Linear
𝑣!"
Task Embed.
Multi-head Attention Layer
MoE Layer
Agg. Linear Matmul
Linear
⨁
Matmul Linear Trans.
⨂
Expert-N
MoE Layer
Feat. Embed. MOFs
Agg.
Softmax
"($ 𝑒!&
𝑣!"($
Energy Autograd
𝜃!"#
Force Stress
Interaction Block (×𝑻)
e3
(c) FSEP Dist Planner
"
!
Attn
"
Rank 0 MoE
Planner MoE 0
Rank 1
Planner MoE 1
Planner MoE 2
Planner MoE 3
All-Reduce
All-Gather Params
(e) Pipelined Gradient Synchronization and Parameter Update ··· Double Backward FW1 FW2 1
2
3 4 1 2 Intra-Rack Broadcast 1 Parameter Update
5 3 2 1
Step per Moe Execution Plan
All-to-All Active Experts
6 ··· Intra-Rack Aggregation 4 5 6 ··· Inter-Rack Sync 3 4 5 6 ··· 2 3 4 5 6 ···
(f) Atom-Aware Comm Compression T 1
T 1
T 2
T 1
T 2
T 3
T 2
T 3
T 1
T 1
T 1
T 2
T 2
T 2
T 3
T 3
Wide-Range
Balanced
(g) HBM Optimization SDMA Engine
1
Fine-Scale
Sub-Graph
𝑊$
𝑉!" GP Feat. Comm
𝑊%
Output
Load A, B
Recovery Parameters
Graph Partition
𝑊#
…
Set Prefetch(order, active experts)
Router Pack
!
Store C
"
Tri
Atomic Graph
×
Shard on Init !
(d) FSGP
𝑉!"
Input
Router Result
DDR Memory
3
A B C Prefetch the Next
Sub-Graph Feat. Embed. All-Gather
Interaction Block
(h) Parallelization Grid HBM
B
A
Pre-Allocated
C
2 Compute CPU Cores
FS-3D
FS-3D …
e2
…
e1
pli ca
e0
Re
MoE
Attn
GP
Tri
N*E
(b) FSDP
FS-3D
… FS-3D
DP Replica
Fig. 2. Overview of our work. (a) Model architecture. (b)–(d) Framework-level optimizations, including FSDP, FSEP, and FSGP. (e)–(g) Supercomputer-level optimizations for communication and memory efficiency. (h) Parallel strategy of DP and GP replica with FS-3D
the algorithmic cost, the compute-intensive multi-head selfattention maps more efficiently onto dense matrix operations, enabling superior utilization of modern GPU hardware for both training and inference. 4) Multi-task training strategies: MatRIS-MoE adopts a conservative training manner, where forces and virial stresses are obtained strictly by automatic differentiation of the predicted energy. To enhance the optimization stability in the multi-task setting, MatRIS-MoE PN adopts a multi-task robust training loss: Lrobust = N1 i=1 W(zi , τ )2 · Li Here, N denotes the number of structures in a batch, and Li is the base loss of the i-th sample (e.g., L1 or Huber loss), defined as a weighted combination of errors in system energy (e), atomic forces (f ), stress (s), and magnetic moments (m). For each forward pass, we compute the mean µ and standard deviation σ of the batch losses and use them to form the standardized score zi = (Li − µ)/σ. The weighting function W(zi , τ ) is a smooth monotone soft-thresholding function. Samples in the normal range (zi ≤ τ , with τ = 2.0 in this work) retain weights close to 1, whereas extreme outliers are progressively down-weighted toward 0. To avoid interference among heterogeneous tasks, the loss statistics used to evaluate zi are computed independently for each task.
sage construction and a feature-update MoE to specialize in post-attention refinement. In each MoE layer, routing is determined by element types. Each element type activates its own top-K experts. Because element types are timeindependent, the activated expert set remains stable during simulation, thereby maintaining a continuous and smooth PES. In multi-task modeling, this sparse element-wise routing allows the model to handle a wider variety of tasks. Different experts can specialize in the chemistry associated with different elements, thereby improving expressiveness compared to the dense parameterization setting. During training, expert parallelism enables efficient distributed execution across largescale clusters. Once the model is trained, during inference, only the top-K routed experts for each element are activated, so the effective parameter number and FLOPs are much smaller than those of the full MatRIS-MoE model. When the number of activated experts is K, and the hidden width is d, the total parameter number of one MoE layer grows as O(K · d2 ). 3) Multi-head self-attention: MatRIS employs separable attention to reduce the complexity of dense pairwise attention from O(N 2 ) to O(N ) by decoupling source- and target-aware aggregation. However, multi-task learning across heterogeneous chemical domains requires greater expressive power to capture complex, task-dependent many-body interactions. To address this limitation, MatRIS-MoE replaces the original separable attention with multi-head self-attention. Under a fixed cutoff radius, the computational complexity of multi-head selfP attention is O(m · d i |N (i)|2 ), where m is the number of attention heads, d is the hidden dimension, and |N (i)| denotes the number of neighbors of atom i. Although this quadratic scaling with respect to the number of neighbors increases
B. Janus: the uMLIP Training Framework Janus is designed to support agile development of MatRISMoE while delivering high efficiency at exascale, with two primary goals: minimizing code intrusion to preserve rapid model iteration, and supporting sharded execution across both double-backward training and MoE layers. Janus adopts FS3D (Fully Sharded 3 Dimensions) as its basic execution unit, and replicates this unit along the data-parallel (DP) and graph-
4
parallel (GP) dimensions (see Fig. 2(h)). Within each DPreplica, the global batch is first divided into mini-batches. Each mini-batch is then partitioned across GP-replicas, and the resulting subgraphs are further partitioned along the intraFS-3D graph dimension. To the best of our knowledge, Janus is the first hybrid-parallel training framework for uMLIPs. The key contributions are as follows. 1) FS-3D, a unified hybrid-parallel runtime for training MatRIS-MoE: FS-3D integrates FSDP (fully sharded data parallelism), FSGP (fully sharded graph parallelism), and FSEP (fully sharded expert parallelism) into a unified execution unit for large-scale training. As illustrated in Fig. 2(b)(d), FS-3D initializes sharded interaction blocks with FSDP, partitions the atomic graph across ranks with FSGP, and then executes sharded MoE layers with expert parallelism (FSEP). Specifically, ➊ FSDP reduces static memory usage, including model parameters, gradients, and optimizer states. Each GPU stores only a local parameter shard. Before computation, FS3D restores the full parameters of non-MoE layers through allgather, whereas MoE expert parameters are restored through all-to-all communication. After computation, the restored parameters are re-sharded to recover the memory savings of sharded training. ➋ FSGP reduces activation memory usage for large graph batches. Built on GP [27], it partitions a large atomic graph (batch) across multiple GPUs, so that each GPU stores and computes only its local partition. ➌ FSEP improves load balance across devices. Inspired by LAER-MoE [28], FSEP partitions expert parameters across devices and restores only the active experts via all-to-all communication. We further design a just-in-time sparse expert planning mechanism to achieve finer-grained load balancing and eliminate redundant communication. We also reorder collective communication operators to ensure both correctness and performance (see Fig. 3). 2) Just-in-time sparse expert planning: Building on FS3D, we further address two core inefficiencies in MoE execution: (i) each training step activates only a small subset of experts, yet conventional FSEP still keeps many unused expert parameters; (ii) the active expert set and its token load can vary sharply across steps in heterogeneous material workloads, making prior approaches (e.g., static expert-to-rank binding, and history-based re-layout) brittle. Because MatRISMoE routes experts based on global atomic embeddings and neighborhood information, all MoE layers can be routed before entering the interaction blocks. To leverage this feature, we adopt a replica-free expert execution approach driven by justin-time (JIT) planning (see Fig. 2(c)). ➊ Plan Generation. At the beginning of each training step, we perform batched routing for all MoE layers and collect per-layer token counts. On each device, we run a local planner based on token counts, which inserts experts into a heap by load, repeatedly pops the heaviest expert, and deterministically assigns it to the leastloaded rank. The global plan, containing the expert prefetch order and token dispatch splits required by the FS-3D runtime, is derived by merging the local plans. Therefore, active experts can be placed with both balanced load and a reproducible
S0
F_TriL S1 P_ATTNL S2
F_MoE0L P_MoE1L GP
F_AttnL P_TriL+1
A2A
F_MoE1L P_MoE0L+1
A2A
A2A
A2A
(a) Forward timeline S0 S1
B_MoE1L P_MoE0L
B_AttnL P_TriL
S2 A2A
A2A
B_MoE0L P_MoE1L-1 A2A
B_TriL P_ATTNL-1 A2A
GP
(b) First backward timeline S0
DB_TriL S1 P_ATTNL
DB_MoE0L P_MoE1L
S2
GP
S0
LB_MoE1L P_MoE0L
S1
DB_AttnL P_TriL+1 A2A
DB_MoE1L P_MoE0L+1 A2A
A2A
(c) Double-backward timeline
S2 A2A S3
A2A LB_AttnL P_TriL A2A
Syn_TriL+1
LB_MoE0L P_MoE1L-1 A2A
Syn_MoE1L
LB_TriL P_ATTNL-1 A2A
GP
Syn_ATTNL
Syn_MoE0L
(d) Final backward with deferred gradient synchronization
Fig. 3. Execution timeline of our framework in MatRIS-MoE. Each interaction block contains Tri, MoE0, Attn, and MoE1. Blue blocks denote operator computation, pink blocks denote parameter prefetch, green blocks denote GP communication or token-routing All-to-All (A2A), and purple blocks denote deferred gradient synchronization.
owner layout. ➋ Expert Sparsification. Based on this plan, FS3D avoids all-gathering the full expert set in the forward pass. Instead, only active experts are materialized through parameter all-to-all communication along the shard dimension, and each active expert is recovered only on its owner rank. ➌ Neighbor Feature Dispatch. FS-3D then dispatches the routed neighboratom features atomicj to the ranks that host the corresponding active experts and gathers the expert outputs through all-to-all communication. ➍ Gradient Combination. In the backward pass, sparse expert gradients are returned through all-to-all gradient exchange along the shard dimension, while replica synchronization is restricted to the union of active experts in the global batch. To preserve consistency under asymmetric expert activation, ranks explicitly zero-fill missing gradients before sparse synchronization along the replica dimension. 3) Runtime for multi-phase double-backward execution: For MatRIS-MoE, each iteration consists of four phases: forward for energy computation, first backward for force computation, double backward for force-loss propagation, and final backward for energy-loss propagation. FS-3D restores parameters before each phase and re-shards them after use through operator-specific communication. It also records the execution order in the forward phase and reuses it in later phases to enable prefetch and overlap. We introduce a doublebackward-aware lifecycle that restores parameters on demand and defers gradient synchronization until the final backward phase. Fig. 3 shows the execution timeline of an FS-3D unit. 4) Load balancing: To mitigate load imbalance caused by variable atomic graph sizes, we adopt a deterministic greedy batching strategy: samples are inserted into a max-heap by atom count, and iteratively assigned to the currently leastloaded mini-batch. This load balancing is performed only within each global batch and does not change the parameter update granularity or training semantics. C. Training Optimizations on CNIS Supercomputer We train MatRIS-MoE on China’s New-generation Intelligent Supercomputer (CNIS), which is equipped with SIMT-
5
based GPGPU accelerators with 64 GB of HBM2e memory each. Each node contains 8 GPGPUs, and the training leverages the Janus framework. 1) Pipelined gradient synchronization and parameter update: To reduce the parameter update overhead of synchronous training, we design a pipelined scheduling mechanism across multiple gradient buckets, coordinating the hierarchical communication and parameter updates. As shown in Fig. 2(e), for each bucket, we preserve the dependency chain of intrarack aggregation → inter-rack synchronization → intra-rack broadcast → parameter update. We leverage the fact that current optimizers introduce no data dependency across gradient buckets. Therefore, our scheduler parallelizes the execution of these operations across different buckets and effectively fills the pipeline while minimizing bubbles. In addition, our scheduler ensures that computation for the next iteration’s forward pass corresponding to bucketi can begin immediately after the parameter update for bucketi is completed. 2) Atom-type-aware communication compression: All-toall communication in MoE layers scales with token count and feature dimension. At the scale of MatRIS-MoE, a single all-to-all incurs a large communication volume per layer. We reduce the communication cost using an atom-type-aware FP16 compression scheme (see Fig. 2(f)). After routing, tokens in the local dispatch buffer are grouped by atom type, quantized with per-type scaling factors, transferred in FP16, and dequantized after communication, before restoring the original order. We exploit the fact that tokens from the same chemical element have similar activation ranges, achieving lower quantization error than coarse-grained per-tensor scaling while maintaining convergence. 3) Kernel engineering for critical operators: MatRIS-MoE operates on atomic graphs with irregular neighbor structures, making several kernels severely memory-bound. We optimize four major performance-critical kernels: neighbor gather, edge aggregation, multi-head attention, and MoE dispatch/combine. The optimizations include fused indexing for coalesced memory access, GEMM-based reformulation of reductions, batched GEMM invocation, and contention-aware scatter. These kernel-level optimizations substantially improve single-accelerator throughput and remain effective at scale in distributed training.
a software-defined asynchronous MPI runtime. Recognizing that the LX2 is a core-rich processor architecture, we physically partition CPU resources. We partition CPU resources as follows: ➊ computation cores: 32 cores are pinned for OpenMP computation, which is sufficient to saturate memory bandwidth; ➋ communication cores: 5 dedicated cores are reserved exclusively for the MPI communication thread pool; and ➌ 1 core reserved for the operating system. By offloading collective operations to communication cores via taskflow task graphs, we successfully emulate CUDA-like stream behavior, thereby avoiding explicit inter-NUMA SDMA calls and keeping the computation cores unblocked. This enables robust overlap between communication and computation, allowing Janus to capitalize on the advanced scheduling optimizations described above. 2) HBM memory optimization via SDMA engine: To mitigate DDR contention, we design an explicit memory hierarchy management strategy that utilizes the on-package SDMA engines. We customize the HBM data movement directly at the memory-bound operator level rather than relying on compilerdriven graph optimizations. We use the optimized GEMM operator as an example. As illustrated in Fig. 2(g), it contains three phases: ➊ upon invocation, the operator synchronously dispatches SDMA instructions to bulk-transfer both the input activations A and weights B from DDR to HBM buffers; ➋ the GEMM computation C=AB is performed using pre-allocated HBM buffers; ➌ once the computation is completed, SDMA is invoked again to write the result C back to DDR memory. We implement a double-buffered SDMA pipeline: while the CPU computes layer k in HBM, the SDMA engine asynchronously prefetches layer k+1 from DDR, effectively hiding the memory transfer latency. By offloading large transfers (e.g., tensors exceeding 1 MB) to the SDMA, we achieve up to a 1.4× bandwidth improvement over CPU-driven transfers under full core utilization, while simultaneously freeing valuable CPU cycles for subsequent instruction decoding. VI. H OW P ERFORMANCE WAS M EASURED A. Physical System Used to Measure Performance To achieve universal generalizability, we construct a large multi-domain dataset comprising over 473 million atomic configurations spanning isolated molecules, periodic crystals, catalytic surfaces, molecular crystals, and metal-organic frameworks (MOFs) [8], [11]–[18]. Ground-truth labels for system energy, atomic forces, and virial stresses are obtained from high-fidelity DFT calculations, with each configuration represented as a unique spatial graph. The MatRIS-MoE architecture exhibits strict linear scaling, O(N ), where N denotes the total number of edges in the spatial graphs. Notably, our 473M dataset contains approximately 3.6 trillion edges, posing significant challenges for system throughput. We evaluate three MatRIS-MoE variants: small (S), medium (M), and large (L). Due to time and computational constraints, only the small model is fully trained; the (M) and (L) variants are used exclusively for weak and strong scaling studies.
D. Training Optimizations on LineShine Supercomputer In this section, we describe the training optimizations on the LineShine platform. The system has an asymmetric memory topology: each LX2 CPU socket contains two compute dies, with each die integrating four NUMA domains (38 ARMv9 cores and 4 GB of HBM per domain) and a dedicated SDMA engine. 1) Async MPI via software-defined streams: FS-3D orchestrates the ordering and dependencies of collective communications, and relies on CUDA streams for asynchronous communication and computation overlap. Given that PyTorch’s CPU backend lacks the asynchronous CUDA stream semantics available on GPUs, we address this limitation by developing
6
All three configurations share the same backbone architecture, consisting of six interaction layers with a MoE routing mechanism and multi-head self-attention (Fig. 2(a)). Detailed hyperparameters are summarized in Table II.
The LineShine supercomputer, developed by the National Supercomputing Center in Shenzhen (NSCC-SZ), is an exascale system consisting of 20,480 computing nodes, denoted the full machine. Each node is equipped with two ARMv9based LX2 processors. Each LX2 integrates two compute dies (304 cores total) and eight on-package HBM stacks (32 GB, 4 TB/s aggregate bandwidth). Each compute die contains 152 cores and 128 GB of off-package DDR memory organized into four NUMA domains. A dedicated SDMA engine handles data movement between DDR and HBM. The LX2 supports FP64/FP32/FP16/INT8 via SME and SVE units, delivering up to 60.3/120.6 TFLOPS in FP64/FP32. Nodes are interconnected via the LingQi high-speed network with a dualplane multi-rail fat-tree topology, offering 1.6 Tb/s bandwidth per node. On CNIS, experiments run on Anolis OS 8.9 with a ROCmcompatible environment, including GCC 8.5.0, rocBLAS, and PyTorch 2.7.1. Inter-node communication uses OpenMPI 5.0.3 and an RCCL-compatible collective library, with a strict oneto-one mapping of 8 RCCL ranks per node, each bound to a single GPGPU. On LineShine, the environment consists of PyTorch 2.10.0, KML 25.2.1 BLAS, and OpenMPI 4.1.6rc4. To maximize the many-core efficiency of the LX2, we employ a hybrid MPI+OpenMP configuration with 16 MPI processes per node, each using multi-threading across 38 physical cores.
TABLE II C ONFIGURATIONS OF M AT RIS-M O E (S), (M) AND (L) MODELS . Hyperparameter Total parameters Active parameters Number of experts Top-k routing Number of layers Node feature dimension Edge feature dimension MLP hidden dimension Number of attention heads Attention dimension
(M)
(L)
1.09 B 0.19 B 72 4 6 2560 384 2560 8 384
2.47 B 0.56 B 40 8 6 1536 1536 1536 4 512
11.5 B 2.89 B 72 16 6 1920 1920 1920 4 512
BS=1024 ODAC
Opoly
1
2
3
4
BS=2048
Norm. F MAE
Norm. E MAE
BS=512
1 0.5 0 1 0.5 0
(S)
5 Epoch
1 0.5 0 1 0.5 0
BS=4096 ODAC
Opoly
1
2
3
4
5
Fig. 4. Convergence behavior of MatRIS-MoE under different batch sizes on the ODAC25 [13] and Opoly26 [11] validation sets.
Figure 4 examines the convergence behavior under varying batch sizes using a 1% sample of the entire dataset. The left and right panels present the energy and force learning curves of MatRIS-MoE, respectively. By scaling the learning rate by the square root of the batch size multiplier, we successfully stabilize the optimization at extreme scales. Quantitatively, increasing the batch size from 512 to 4,096 (4K) results in less than a 10% increase in the final energy loss. Notably, in our Exascale evaluation, the maximum global batch sizes reach 4K and 15K on the CNIS and LineShine systems, respectively.
C. Measurement Methodology
B. HPC Systems and Software
•
The total floating-point operations (FLOPs) are measured in FP32 using the PyTorch Profiler. Since all compute-intensive operations are offloaded to accelerators, the profiled GPGPU FLOPs accurately represent the overall workload. Over 1,000 training iterations processing 30 million atomistic samples, the cumulative workload reaches 2974 EFLOPs for the 11.50Bparameter MatRIS-MoE (L) and 1765 EFLOPs for the 2.47Bparameter MatRIS-MoE (M). The following criteria are used to measure the performance of our code: 7,610×global batch size Average throughput, defined as average time per iteration . Each atomistic configuration has 7610 edges on average. total FLOPs • Peak performance, defined as MatRIS-MoE training loop time . total FLOPs • Sustained performance, defined as total wall clock time . The “total wall clock time” includes the whole application running time (including IO, MPI initialization and finalization).
All performance measurements are conducted on two distinct Exascale supercomputers: the China New-generation Intelligent Supercomputer (CNIS) and the LineShine system. The CNIS is a heterogeneous Exascale supercomputer consisting of 5, 632 computing nodes, denoted the full machine. Each node is equipped with two 64-bit CISC-based server processors and eight SIMT-based GPGPU accelerators. The host processor runs at 2.4 GHz with 64 cores in a NUMA architecture, connected to 8-channel DDR5-6400 memory and PCIe Gen5 interfaces, delivering 64 GB/s host-to-device bandwidth. Each GPGPU provides 32.7 TFLOPS (FP64), 65.5 TFLOPS (FP32), and 470 TFLOPS (FP16) peak performance, with 64 GB HBM (1.8 TB/s bandwidth), 320 SIMD units, 768 KB registers, 64 KB LDS, and 8 MB L2 cache. The accelerators are interconnected via high-speed chip-to-chip links, while nodes are connected through a proprietary InfiniBand-like RDMA network with a three-layer Clos dual-plane topology, providing 4 × 400 Gbps per node.
VII. P ERFORMANCE R ESULTS A. Accuracy Results Although MatRIS-MoE (S) has only about 1 billion parameters, it achieves state-of-the-art (SOTA) or near-SOTA accuracy across multiple benchmarks. These tasks comprise highly complex, real-world downstream applications and effectively demonstrate the generalization capacity of MatRIS-MoE. As shown in Fig. 5, the evaluation spans a broad spectrum of systems, including materials (Matcalc [16], reporting structural similarity dE, bulk moduli K and shear moduli G), molecules (Wiggle150 [29] with energy RMSE and MAE;
7
SevenNet- MACEOmni mh-1
𝑑𝐸 Matcalc
Molecular Crystal
TABLE III OVERALL IMPROVEMENT ON CNIS AND L INE S HINE .
1.0
MAE small-system large-system
GMTKN55 barrier-height (kcal/mol)
X23 (kJ/mol)
E_formation
DMC Water Ice (kJ/mol)
E_relative 𝐸𝑎𝑑𝑠
OC20NEB (eV)
𝐸𝑟𝑥𝑛
ASD (eV)
𝐸𝑟𝑥𝑛
Noble metal ads (eV)
MOFSim
Machine
Model
Baseline
Optimized
Speedup
CNIS
MatRIS-MoE (M) MatRIS-MoE (L)
5.96s 7.71s
2.21s 2.66s
2.7× 2.9×
LineShine
MatRIS-MoE (M) MatRIS-MoE (L)
33.1s 40.7s
8.08s 8.14s
4.1× 5.0×
RMSE
Inter. NCls
Catalysis
MatRISMoE
𝐾 (GPa)
Intra. NCls
MOFs
UMAm-1p1
𝐺 (GPa) Wiggle150 (kcal/mol)
Molecule
Metric
Normalized score (higher is better)
Material
Domain Benchmark
Although MatRIS-MoE (M) has fewer parameters than MatRIS-MoE (L) (2.47B vs. 11.5B), their per-step training times remain comparable. This is because MatRIS-MoE (M) processes twice as many atoms per FS-3D unit (128 vs. 64), resulting in a higher per-rank computational workload. Consequently, the theoretical per-step FLOP count of MatRISMoE (M) reaches 59.36% of that of MatRIS-MoE (L). The observed performance gains come from three systemlevel optimizations. First, we apply an asynchronous optimizer that overlaps gradient synchronization with parameter updates, thereby reducing synchronization stalls during double backward. Second, we apply the atom-type-aware compression that reduces the MoE routing all-to-all communication volume by 50% without sacrificing numerical precision. Third, we design several high-performance kernels including graph operators, attention, and MoE dispatch/combine, thereby improving single-accelerator performance. Together with the above communication optimizations, the up to 5.0× speedup on LineShine is further enabled by a platform-specific optimization: SDMA-based data movement between DDR and on-package HBM. This addresses the limited HBM utilization of native PyTorch operators on the LX2 processor, allowing operators to achieve up to 1.4× higher memory bandwidth. We evaluate the strong scaling of MatRIS-MoE training from 1/8 to full-machine scale while keeping the global batch size fixed for 1000 iterations. This reduces the workload per FS-3D unit through atomic-graph partitioning, decreasing the number of atoms per FS-3D unit from 128 to 64, 32, and 16 for MatRIS-MoE (M), and from 64 to 32, 16, and 8 for MatRIS-MoE (L). As shown in Fig. 6, both models scale up to full-machine scale.
𝐸𝑎𝑑𝑠 E_interaction (eV) F MAE (eV/ Å) volume RMSE
CV RMSE (J/K/g)
0.0
Fig. 5. The out-of-the-box accuracy results of MatRIS-MoE on cross-domain benchmarks. The representative tasks span molecules, materials, catalysis, molecular crystals, and MOFs. “-” indicates that the corresponding results are not evaluated because the training dataset of MACE-mh-1 does not include MOF configurations. Lighter colors indicate higher accuracy.
GMTKN55 [30] with small-system, large-system, barrierheight, and intra-/intermolecular NCI subsets), molecular crystals (X23 [31] formation energies and DMC Ice [32] relative energies), catalysis (OC20NEB-OOD [33] with adsorption energies Eads and reaction energies Erxn , ASD Erxn , and noble-metal adsorption [8] with Eads ), and MOFs (MOFSimBench [34] with interaction-energy MAE, force MAE, volume RMSE and Heat capacity CV ). Lighter background colors indicate higher normalized accuracy. These results highlight MatRIS-MoE’s strong generalization capability and representational power (e.g., achieving ωB97M-V-level accuracy on Wiggle150 [29] while delivering an approximately threeorders-of-magnitude speedup), enabling a single model to deliver robust, out-of-the-box predictions across highly diverse chemical systems, including materials, molecules, catalytic surfaces, molecular crystals, and metal-organic frameworks.
C. Strong Scaling Taking 1/8-scale performance as baseline, MatRIS-MoE (L) achieves parallel efficiencies of 53.93% on CNIS and 50.60% on LineShine at full scale, reaching peak performances of 602.9 PFLOPS and 672.5 PFLOPS, respectively. The average throughput scales from 0.80 million to 1.63 million edges/second on CNIS and from 0.90 million to 1.77 million edges/second on LineShine. For MatRIS-MoE (M), the efficiencies are 51.87% and 49.9%, with peak performances of 414.7 PFLOPS and 471.3 PFLOPS. The average throughput scales from 1.86 million to 3.92 million edges/second on CNIS and from 2.12 million to 4.24 million edges/second on LineShine.
B. Performance Improvement In this section, we measure the reduction in per-step training time achieved by our system-level optimizations. All experiments are performed at 1/8 full-machine scale. For MatRISMoE, both the 2.47B-parameter (M) and 11.5B-parameter (L) models use the same parallel configuration, with an FS-3D unit size of 8 and a GP-replica size of 8, yielding 128 and 64 atoms per FS-3D unit, respectively. Both configurations scale to the full system through the DP-replica dimension. Table III summarizes the per-step training time before and after optimization, together with the resulting speedup.
8
400 166.1P 152.4P (27.02%) (37.88%) 200
602.9P (20.43%)
259.9P 245.6P (21.15%) (30.51%)
672.5P (13.68%)
397.5P 382.1P (16.18%) (23.74%)
264.1P (16.4%)
277.9P (11.3%)
414.7P 471.3P (14.05%) (9.59%)
1/4
1/2
Fraction of Entire Machine
400 152.4P
166.1P (27.02%)
100
1
118P 109P (19.21%) (27.09%)
1/8
628.7P 587.4P(25.58%) (36.49%)
325.7P 300.5P (26.51%) (37.34%)
200 (37.88%)
173.1P 180.2P 118P (21.5%) (14.66%) 109P (19.21%) (27.09%)
1/8
CNIS - L Lineshine - L CNIS - M Lineshine - M
800
PFLOPS
PFLOPS
800
100
1600
CNIS - L Lineshine - L CNIS - M Lineshine - M
1200.2P (24.41%) 1048.3P (35.52%)
447P (18.19%) 421.2P (26.17%)
861.6P (17.53%) 750.9P (25.44%)
231.9P (18.87%) 215.3P (26.75%)
1/4
1/2
Fraction of Entire Machine
1
Fig. 6. Strong scaling of MatRIS-MoE training on CNIS and LineShine. MatRIS-MoE (L) (solid markers) and MatRIS-MoE (M) (hollow markers) are scaled from 1/8 to full-machine scale with a fixed global problem size. The corresponding training throughput in PFLOPS and relative parallel efficiency are presented. Both models maintain over 50% relative efficiency at fullmachine scale.
Fig. 7. Weak scaling of MatRIS-MoE training on CNIS and LineShine. The problem size is increased proportionally with the machine scale from 1/8 to full-machine scale. The corresponding training throughput in PFLOPS is shown. Performance scales near-linearly, reaching 1.05 EFLOPS on CNIS and 1.20 EFLOPS on LineShine for MatRIS-MoE (L), with over 90% parallel efficiency on both systems.
These consistent efficiencies across GPGPU-accelerated and many-core ARM platforms validate the portability of our Janus framework. The primary scaling bottlenecks arise from reduced GEMM sizes (due to fewer atoms per FS-3D unit), which lower arithmetic intensity and accelerator utilization, particularly for graph-based operators, as well as increased gradient All-Reduce communication overhead across more DP replicas, especially during second-order backpropagation.
imports trigger additional HBM and PyTorch runtime initialization (531.8 seconds). Measured over 1, 000 training steps on the full machine with the MatRIS-MoE (L), sustained performance reaches 762.3 PFLOPS (25.84% of peak) on CNIS and 1,033.3 PFLOPS (21.02% of peak) on LineShine. The gap between sustained and peak performance arises primarily from initialization overheads, which become increasingly pronounced at Exascale scale.
D. Weak Scaling VIII. I MPLICATIONS
In the weak scaling tests, the global batch size and the number of nodes are scaled proportionally from 1/8 to full-machine scale while maintaining the same intra-group configuration as in the strong-scaling experiments. The number of atoms per FS-3D unit was kept constant at 128 for (M) and 64 for (L) across all scales. Both models exhibit near-linear scaling. For MatRIS-MoE (L), peak performance increases from 152.4 PFLOPS to 1,048.3 PFLOPS on CNIS (93.78% efficiency) and reaches 1,200.2 PFLOPS (1.2 EFLOPS) on LineShine (90.3% efficiency). For MatRIS-MoE (M), throughput scales from 109.0 PFLOPS to 750.9 PFLOPS on CNIS (93.91% efficiency) and reaches 861.6 PFLOPS on LineShine (91.3% efficiency). The only scaling overhead is gradient All-Reduce across expanding DP replica groups, which our optimized collectives handle efficiently. The sustained efficiency reaches 72.73% on CNIS and 86.61% on LineShine, demonstrating the effectiveness of our system–algorithm co-design.
A. Implication for Applications Figure 8 highlights the broad applicability and robustness of MatRIS-MoE across representative molecular, catalytic, and materials systems. Across a range of downstream tasks, the model remains in close agreement with DFT. In practical polymorph energy-ranking tasks [35], MatRIS-MoE predictions are very close to the DFT references (Fig. 8(a)). MatRISMoE can accurately reproduce the radial distribution functions of an Aspirin configuration in rMD17 (Fig. 8(b)). In the structural relaxation [36] of Ca20 O60 Se20 (Fig. 8(c)), MatRISMoE efficiently drives the system toward a stable local minimum, with both the energy difference (∆E) and Root Mean Square Displacement (RMSD) approaching the DFTrelaxed reference. Continuous NVT molecular dynamics simulations [34] of a complex metal-organic framework (Fig. 8(d)) further show that the model preserves structural integrity and thermodynamic stability as the temperature increases. MatRIS-MoE can also capture the relative energy profile along a representative catalytic reaction pathway [37] (Fig. 8(e)). Beyond the representative examples shown in Fig. 8, MatRISMoE can also support a wide range of practical applications, such as the prediction of mechanical properties, phonons, and phase diagrams. MatRIS-MoE bridges quantum-level fidelity with enhanced representational power for heterogeneous, multi-component, and multi-task systems, and therefore serves as a powerful and
E. Sustained Performance While the peak performance in Sec. VII-D represents steady-state throughput, sustained performance accounts for the end-to-end wall-clock time, encompassing overheads such as initialization and I/O. On CNIS, setup requires 591.3 seconds for RCCL initialization and 472.2 seconds for firststep warm-up. On LineShine, MPI initialization takes 923.5 seconds due to the larger node count, and Python module
9
versatile simulation engine for critical scientific challenges, including ion transport in solid-state electrolytes, active-site dynamics in single-atom catalysis, and the thermodynamic stability of complex doping configurations. Moreover, compared to traditional DFT calculations, MatRIS-MoE is two orders of magnitude faster. Due to its high out-of-the-box predictive accuracy and strong generalization, MatRIS-MoE has the potential to further expand the horizon of MD simulations with ab initio accuracy. (a)
data movement, reducing memory footprints during secondorder differentiation, and compressing inter-node communication overheads are imperative to sustain parallel efficiency. Conquering these system bottlenecks will be the definitive key to successful model scaling in the future. Furthermore, the rapid evolution of AI4S signals a fundamental computing paradigm shift: AI workloads and traditional HPC simulation workflows are becoming equally important. With their full-precision capabilities (FP64/FP32/FP16), massive high-bandwidth memory, and tightly coupled interconnects, Exascale supercomputers are unequivocally the ideal platforms for the next-generation workloads. To fully harness these architectures, the community must develop more flexible and extensible hybrid programming frameworks that seamlessly integrate conventional scientific computing tools (e.g., molecular dynamics engines and DFT solvers) with large-scale distributed neural architectures. Ultimately, as atomic foundation models continue to scale in both capacity and efficiency, their enhanced representational power will unlock the ability to tackle increasingly complex and high-impact scientific problems. By accurately capturing complex multi-body interactions across the periodic table, large-scale uMLIPs will bridge the long-standing gap between theoretical simulation and practical application, catalyzing transformative breakthroughs in clean energy, advanced materials, and next-generation pharmaceuticals.
(b)
(c)
∆𝐸: 18.139 eV RMSD: 0.144 Å
∆𝐸: 0.065 eV RMSD: 0.011 Å
(d)
ACKNOWLEDGMENT Numerical calculations are performed on the LineShine and the China New-generation Intelligent Supercomputer (CNIS) supercomputers. We thank NSCC-SZ for technical support throughout this project. This work was partially supported by the National Science Foundation of China(92270206, T2125013, 62372435).
(e) ∆E: 0.061 eV ∆E: 0.011 eV
∆E: 0.037 eV
R EFERENCES
Fig. 8. Applicability of MatRIS-MoE across (a) energy ranking, (b) molecular distribution functions, (c) structural relaxation, (d) molecular dynamics, and (e) catalytic reaction profiling.
[1] T. Li and G. Galli, “Electronic Properties of MoS2 Nanoparticles,” The Journal of Physical Chemistry C, vol. 111, no. 44, pp. 16 192–16 196, Nov. 2007. [2] G. Ceder, Y.-M. Chiang, D. R. Sadoway, M. K. Aydinol, Y.-I. Jang, and B. Huang, “Identification of cathode materials for lithium batteries guided by first-principles calculations,” Nature, vol. 392, no. 6677, pp. 694–696, Apr. 1998. [3] M. W. Tibbitt, C. B. Rodell, J. A. Burdick, and K. S. Anseth, “Progress in material design for biomedical applications,” Proceedings of the National Academy of Sciences, vol. 112, no. 47, pp. 14 444–14 451, 2015. [4] Z. Fan, Z. Zeng, C. Zhang, Y. Wang, K. Song, H. Dong, Y. Chen, and T. Ala-Nissila, “Neuroevolution machine learning potentials: Combining high accuracy and low cost in atomistic simulations and application to heat transport,” Physical Review B, vol. 104, no. 10, p. 104309, 2021. [5] S. Batzner, A. Musaelian, L. Sun, M. Geiger, J. P. Mailoa, M. Kornbluth, N. Molinari, T. E. Smidt, and B. Kozinsky, “E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials,” Nature Communications, vol. 13, no. 1, May 2022. [6] W. Jia, H. Wang, M. Chen, D. Lu, L. Lin, R. Car, W. E, and L. Zhang, “Pushing the limit of molecular dynamics with ab initio accuracy to 100 million atoms with machine learning,” in Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, ser. SC ’20. IEEE Press, 2020.
B. Implication for HPC Systems This paper establishes billion-parameter universal MLIP training as a new Exascale AI-for-Science workload, and demonstrates that only through model–framework–system codesign can such workloads achieve efficient, end-to-end training at Exascale across many-core Exascale platforms. As the volume and diversity of ab initio atomistic datasets continue to grow, pushing the boundaries of universal MLIPs will inevitably require scaling models beyond the 10-billionparameter regime explored in this work. This trajectory, however, indicates that substantial future efforts must be directed toward extreme-scale system optimization. The effective utilization of dense matrix-multiply units (e.g., Tensor Cores and SME) will remain the cornerstone of arithmetic efficiency. Simultaneously, aggressively mitigating
10
for fast and accurate force fields,” in Advances in Neural Information Processing Systems, 2022. [26] D. Zhang, A. Peng, C. Cai, W. Li, Y. Zhou, J. Zeng, M. Guo, C. Zhang, B. Li, H. Jiang, T. Zhu, W. Jia, L. Zhang, and H. Wang, “A graph neural network for the era of large atomistic models,” 2026. [27] A. Sriram, A. Das, B. M. Wood, and C. L. Zitnick, “Towards training billion parameter graph neural networks for atomic simulations,” in International Conference on Learning Representations, 2022. [28] X. Liu, Y. Wang, F. Fu, X. Xiao, H. Li, J. Li, and B. Cui, “Laer-moe: Load-adaptive expert re-layout for efficient mixture-of-experts training,” in ASPLOS ’26, Volume 2, NY, USA, 2026, p. 1055–1072. [29] R. R. Brew, I. A. Nelson, M. Binayeva, A. S. Nayak, W. J. Simmons, J. J. Gair, and C. C. Wagen, “Wiggle150: Benchmarking density functionals and neural network potentials on highly strained conformers,” Journal of Chemical Theory and Computation, vol. 21, no. 8, pp. 3922–3929, 2025, pMID: 40211427. [30] L. Goerigk, A. Hansen, C. Bauer, S. Ehrlich, A. Najibi, and S. Grimme, “A look at the density functional theory zoo with the advanced gmtkn55 database for general main group thermochemistry, kinetics and noncovalent interactions,” Phys. Chem. Chem. Phys., vol. 19, pp. 32 184–32 215, 2017. [31] A. M. Reilly and A. Tkatchenko, “Understanding the role of vibrations, exact exchange, and many-body van der waals interactions in the cohesive properties of molecular crystals,” The Journal of Chemical Physics, vol. 139, no. 2, p. 024705, 07 2013. [32] F. Della Pia, A. Zen, D. Alfè, and A. Michaelides, “Dmc-ice13: Ambient and high pressure polymorphs of ice from diffusion monte carlo and density functional theory,” The Journal of Chemical Physics, vol. 157, no. 13, p. 134701, 10 2022. [33] B. Wander, M. Shuaibi, J. R. Kitchin, Z. W. Ulissi, and C. L. Zitnick, “Cattsunami: Accelerating transition state energy calculations with pretrained graph neural networks,” 2024. [34] H. Kraß, J. Huang, and S. M. Moosavi, “Mofsimbench: evaluating universal machine learning interatomic potentials in metal-organic framework molecular modeling,” npj Computational Materials, vol. 12, no. 1, p. 4, Dec 2025. [35] L. S. Price, M. Paloni, M. Salvalaglio, and S. L. Price, “One size fits all? development of the cposs209 data set of experimental and hypothetical polymorphs for testing computational modeling methods,” Crystal Growth & Design, vol. 25, no. 9, pp. 3186–3209, 2025. [36] J. Riebesell, R. E. A. Goodall, P. Benner, Y. Chiang, B. Deng, G. Ceder, M. Asta, A. A. Lee, A. Jain, and K. A. Persson, “A framework to evaluate machine learning crystal stability predictions,” Nature Machine Intelligence, vol. 7, no. 6, pp. 836–847, Jun 2025. [37] J. M. Bowman, C. Qu, R. Conte, A. Nandi, P. L. Houston, and Q. Yu, “The md17 datasets from the perspective of datasets for gas-phase “small” molecule potentials,” The Journal of Chemical Physics, vol. 156, no. 24, p. 240901, 06 2022.
[7] B. M. Wood, M. Dzamba, X. Fu, M. Gao, M. Shuaibi, L. BarrosoLuque, K. Abdelmaqsoud, V. Gharakhanyan, J. R. Kitchin, D. S. Levine, K. Michel, A. Sriram, T. Cohen, A. Das, S. J. Sahoo, A. Rizvi, Z. W. Ulissi, and C. L. Zitnick, “UMA: A family of universal models for atoms,” in The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. [8] J. Kim, J. You, Y. Park, Y. Lim, Y. Kang, J. Kim, H. Jeon, S. Ju, D. Hong, S. Y. Lee, S. Choi, Y. Kim, J. W. Lee, and S. Han, “Optimizing crossdomain transfer for universal machine learning interatomic potentials,” Nature Communications, Mar. 2026. [9] B. Deng, P. Zhong, K. Jun, J. Riebesell, K. Han, C. J. Bartel, and G. Ceder, “Chgnet as a pretrained universal neural network potential for charge-informed atomistic modelling,” Nature Machine Intelligence, p. 1–11, 2023. [10] Y. Zhou, S. Hu, X. Zhang, H. Wang, G. Tan, and W. Jia, “MatRIS: Toward reliable and efficient pretrained machine learning interatomic potentials,” in The Fourteenth International Conference on Learning Representations, 2026. [11] D. S. Levine, N. Liesen, L. Chua, J. Diffenderfer, H. Ingolfsson, M. P. Kroonblawd, N. Kumar, A. Maiti, S. S. Mohottalalage, M. Shuaibi, B. V. Essen, B. M. Wood, C. L. Zitnick, S. M. Blau, and E. R. Antoniuk, “The open polymers 2026 (opoly26) dataset and evaluations,” 2026. [12] V. Gharakhanyan, L. Barroso-Luque, Y. Yang, M. Shuaibi, K. Michel, D. S. Levine, M. Dzamba, X. Fu, M. Gao, X. Liu, H. Ni, K. Noori, B. M. Wood, M. Uyttendaele, A. Boromand, C. L. Zitnick, N. Marom, Z. W. Ulissi, and A. Sriram, “Open molecular crystals 2025 (omc25) dataset and models,” 2025. [13] A. Sriram, L. M. Brabson, X. Yu, S. Choi, K. Abdelmaqsoud, E. Moubarak, P. de Haan, S. Löwe, J. Brehmer, J. R. Kitchin, M. Welling, C. L. Zitnick, Z. Ulissi, A. J. Medford, and D. S. Sholl, “The open dac 2025 dataset for sorbent discovery in direct air capture,” 2025. [14] S. J. Sahoo, M. Maraschin, D. S. Levine, Z. Ulissi, C. L. Zitnick, J. B. Varley, J. A. Gauthier, N. Govindarajan, and M. Shuaibi, “The open catalyst 2025 (oc25) dataset and models for solid-liquid interfaces,” arXiv preprint arXiv:2509.17862, 2025. [15] L. Chanussot, A. Das, S. Goyal, T. Lavril, M. Shuaibi, M. Riviere, K. Tran, J. Heras-Domingo, C. Ho, W. Hu et al., “Open catalyst 2020 (oc20) dataset and community challenges,” Acs Catalysis, vol. 11, no. 10, pp. 6059–6072, 2021. [16] A. D. Kaplan, R. Liu, J. Qi, T. W. Ko, B. Deng, J. Riebesell, G. Ceder, K. A. Persson, and S. P. Ong, “A foundational potential energy surface dataset for materials,” 2025. [17] L. Barroso-Luque, M. Shuaibi, X. Fu, B. M. Wood, M. Dzamba, M. Gao, A. Rizvi, C. L. Zitnick, and Z. W. Ulissi, “Open materials 2024 (omat24) inorganic materials dataset and models,” arXiv preprint arXiv:2410.12771, 2024. [18] D. S. Levine, M. Shuaibi, E. W. C. Spotte-Smith, M. G. Taylor, M. R. Hasyim, K. Michel, I. Batatia, G. Csányi, M. Dzamba, P. Eastman, N. C. Frey, X. Fu, V. Gharakhanyan, A. S. Krishnapriyan, J. A. Rackers, S. Raja, A. Rizvi, A. S. Rosen, Z. Ulissi, S. Vargas, C. L. Zitnick, S. M. Blau, and B. M. Wood, “The open molecules 2025 (omol25) dataset, evaluations, and models,” 2026. [19] C. Chen and S. P. Ong, “A universal graph deep learning interatomic potential for the periodic table,” Nature Computational Science, vol. 2, no. 11, pp. 718–728, Nov. 2022. [20] A. Merchant, S. Batzner, S. S. Schoenholz, M. Aykol, G. Cheon, and E. D. Cubuk, “Scaling deep learning for materials discovery,” Nature, vol. 624, no. 7990, pp. 80–85, 2023. [21] B. Rhodes, S. Vandenhaute, V. Šimkus, J. Gin, J. Godwin, T. Duignan, and M. Neumann, “Orb-v3: atomistic simulation at scale,” 2025. [22] F. Bigi, P. Pegolo, A. Mazitov, and M. Ceriotti, “Pushing the limits of unconstrained machine-learned interatomic potentials,” 2026. [23] I. Batatia, C. Lin, J. Hart, E. Kasoar, A. M. Elena, S. W. Norwood, T. Wolf, and G. Csányi, “Cross learning between electronic structure theories for unifying molecular, surface, and inorganic crystal foundation force fields,” 2025. [24] Y.-L. Liao, B. M. Wood, A. Das, and T. Smidt, “Equiformerv2: Improved equivariant transformer for scaling to higher-degree representations,” in The Twelfth International Conference on Learning Representations, 2024. [25] I. Batatia, D. P. Kovacs, G. N. C. Simm, C. Ortner, and G. Csanyi, “MACE: Higher order equivariant message passing neural networks
11