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

• URL: http://arxiv.org/abs/2604.15821v1
• Date: Fri, 17 Apr 2026 08:20:32 GMT
• Title: Breaking the Training Barrier of Billion-Parameter Universal Machine Learning Interatomic Potentials
• Abstract summary: We introduce MatRIS-MoE and Janus, a pioneering high-dimensional distributed training framework for uMLIPs with hardware-aware optimizations.<n>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.<n>This work establishes a new high-water mark for AI-for-Science (AI4S) foundation models at Exascale.
• Score: 25.963136125890298
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• SLAI T-Rex: Full-Parameter Post-training of the DeepSeek-V4 Family on Ascend SuperPOD [91.71050339479261]
  This report presents an end-to-end optimization practice on the Ascend NPU SuperPOD.<n>We develop a hierarchical optimization framework spanning model-level parallelism, computation-communication orchestration, and low-level kernel execution.<n>The resulting system achieves 34.22% Model FLOPs Utilization (MFU) with a 2.93x improvement over the open-source baseline recipe.
  arXiv  Detail & Related papers  (2026-07-22T13:49:17Z)
• MALOQ: Massively Accelerated Learning of Operators for Quantum Transport [71.9141151643915]
  MALOQ is an application built to train on and predict electronic-structure matrices for systems made of few to 100k atoms.<n>It reduces time-per-epoch by over 30% compared to a molecule-wise-distributed framework.<n>We demonstrate scalable training and inference for 3,000-12,000 atoms on the Alps supercomputer.
  arXiv  Detail & Related papers  (2026-06-27T13:34:52Z)
• Scalable Training of Mixture-of-Experts Models with Megatron Core [26.9162079065285]
  Scaling Mixture-of-Experts (MoE) training introduces systems challenges absent in dense models.<n>Because each token activates only a subset of experts, this sparsity allows total parameters to grow much faster than per-token computation.<n>We address these challenges for MoE training through integrated optimizations spanning memory (fine-dimensional recomputation, offloading), communication, and computation.
  arXiv  Detail & Related papers  (2026-03-08T15:42:43Z)
• MagicAgent: Towards Generalized Agent Planning [73.21129030631421]
  We present textbfMagicAgent, a series of foundation models specifically designed for generalized agent planning.<n>We introduce a lightweight and scalable synthetic data framework that generates high-quality trajectories across diverse planning tasks.<n>We show that MagicAgent-32B and MagicAgent-30B-A3B achieve superior performance across diverse open-source benchmarks.
  arXiv  Detail & Related papers  (2026-02-22T01:39:16Z)
• Youtu-LLM: Unlocking the Native Agentic Potential for Lightweight Large Language Models [78.73992315826035]
  We introduce Youtu-LLM, a lightweight language model that harmonizes high computational efficiency with native agentic intelligence.<n>Youtu-LLM is pre-trained from scratch to systematically cultivate reasoning and planning capabilities.
  arXiv  Detail & Related papers  (2025-12-31T04:25:11Z)
• Iterative Pretraining Framework for Interatomic Potentials [46.53683458224917]
  We propose Iterative Pretraining for Interatomic Potentials (IPIP) to improve predictive performance of MLIP models.<n>IPIP incorporates a forgetting mechanism to prevent iterative training from converging to suboptimal local minima.<n>Compared to general-purpose force fields, this approach achieves over 80% reduction in prediction error and up to 4x speedup in the challenging Mo-S-O system.
  arXiv  Detail & Related papers  (2025-07-27T03:59:41Z)
• AutoHete: An Automatic and Efficient Heterogeneous Training System for LLMs [68.99086112477565]
  Transformer-based large language models (LLMs) have demonstrated exceptional capabilities in sequence modeling and text generation.<n>Existing heterogeneous training methods significantly expand the scale of trainable models but introduce substantial communication overheads and CPU workloads.<n>We propose AutoHete, an automatic and efficient heterogeneous training system compatible with both single- GPU and multi- GPU environments.
  arXiv  Detail & Related papers  (2025-02-27T14:46:22Z)
• EPS-MoE: Expert Pipeline Scheduler for Cost-Efficient MoE Inference [49.94169109038806]
  This paper introduces EPS-MoE, a novel expert pipeline scheduler for MoE that surpasses the existing parallelism schemes.<n>Our results demonstrate at most 52.4% improvement in prefill throughput compared to existing parallel inference methods.
  arXiv  Detail & Related papers  (2024-10-16T05:17:49Z)
• Physics-Informed Weakly Supervised Learning for Interatomic Potentials [17.165117198519248]
  We introduce a physics-informed, weakly supervised approach for training machine-learned interatomic potentials (MLIPs)<n>We demonstrate reduced energy and force errors -- often lower by a factor of two -- for various baseline models and benchmark data sets.<n>Our approach improves the fine-tuning of foundation models on sparse, highly accurate ab initio data.
  arXiv  Detail & Related papers  (2024-07-23T12:49:04Z)
• A Multi-Grained Symmetric Differential Equation Model for Learning Protein-Ligand Binding Dynamics [73.35846234413611]
  In drug discovery, molecular dynamics (MD) simulation provides a powerful tool for predicting binding affinities, estimating transport properties, and exploring pocket sites. We propose NeuralMD, the first machine learning (ML) surrogate that can facilitate numerical MD and provide accurate simulations in protein-ligand binding dynamics. We demonstrate the efficiency and effectiveness of NeuralMD, achieving over 1K$times$ speedup compared to standard numerical MD simulations.
  arXiv  Detail & Related papers  (2024-01-26T09:35:17Z)
• MAD Max Beyond Single-Node: Enabling Large Machine Learning Model Acceleration on Distributed Systems [6.8519529064678375]
  Training and deploying large-scale machine learning models is time-consuming, requires significant distributed computing infrastructures, and incurs high operational costs. To minimize this outstanding communication latency and other inherent at-scale inefficiencies, we introduce an agile performance modeling framework, MAD-Max. This framework is designed to optimize parallelization strategies and facilitate hardware-software co-design opportunities.
  arXiv  Detail & Related papers  (2023-10-04T13:00:53Z)
• h-analysis and data-parallel physics-informed neural networks [0.7614628596146599]
  We explore the data-parallel acceleration of machine learning schemes with a focus on physics-informed neural networks (PINNs) We detail a novel protocol based on $h$-analysis and data-parallel acceleration through the Horovod training framework. We show that the acceleration is straightforward to implement, does not compromise training, and proves to be highly efficient and controllable.
  arXiv  Detail & Related papers  (2023-02-17T12:15:18Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.