POSTER: TENSORMD: Molecular Dynamics Simulation with Ab Initio Accuracy of 50 Billion Atoms
AI has been coupled into HPC in many scientific fields and demonstrated significantly enhanced performances. Among them, molecular dynamics simulation emerges as an important area that HPC+AI helps to investigate the physical properties at the atomic scale, with machine-learning interatomic potentials (MLIPs) being used. Until now, general-purpose machine-learning (ML) tools (e.g., TensorFlow) have been leveraged in MLIPs, but they are not perfectly matched with each other, since many optimization opportunities in MLIPs have been missed by ML tools. This inefficiency arises from the fact that HPC+AI applications work with far more computational complexity and range compared with pure AI scenarios, which general-purpose ML tools target. This paper has developed an MLIP, named TensorMD, independently from any ML tool, to enable a set of optimizations to be flexibly performed, which are otherwise impossible to be applied in ML tools. TensorMD has been evaluated on two supercomputers and scaled to $51.8$ billion atoms, i.e., $\sim 3\times$ compared with state-of-the-art. We further examine the reasons why optimizations successfully applied in TensorMD are beyond the capabilities of ML tools, providing insights into compiler optimization design for HPC+AI applications.