TensorMD: Scalable Tensor-Diagram based Machine Learning Interatomic Potential on Heterogeneous Many-Core Processors

TL;DR

TensorMD is a novel machine learning interatomic potential that leverages tensor diagrams for efficient, scalable molecular dynamics simulations on heterogeneous many-core processors, enabling unprecedented large-scale atomic simulations.

Contribution

The paper introduces TensorMD, combining physical principles with tensor diagrams, and develops optimized implementations for the Sunway supercomputer, achieving record-breaking simulation scale and performance.

Findings

Simulates up to 52 billion atoms.

Achieves a time-to-solution of 31 ps/step/atom.

Sets new records for HPC + AI + MD performance.

Abstract

Molecular dynamics simulations have emerged as a potent tool for investigating the physical properties and kinetic behaviors of materials at the atomic scale, particularly in extreme conditions. Ab initio accuracy is now achievable with machine learning based interatomic potentials. With recent advancements in high-performance computing, highly accurate and large-scale simulations become feasible. This study introduces TensorMD, a new machine learning interatomic potential (MLIP) model that integrates physical principles and tensor diagrams. The tensor formalism provides a more efficient computation and greater flexibility for use with other scientific codes. Additionally, we proposed several portable optimization strategies and developed a highly optimized version for the new Sunway supercomputer. Our optimized TensorMD can achieve unprecedented performance on the new Sunway, enabling simulations of up to 52 billion atoms with a time-to-solution of 31 ps/step/atom, setting new records for HPC + AI + MD.