Scaling Molecular Dynamics with ab initio Accuracy to 149 Nanoseconds per Day

Abstract

Physical phenomena such as chemical reactions, bond breaking, and phase transition require molecular dynamics (MD) simulation with ab initio accuracy ranging from milliseconds to microseconds. However, previous state-of-the-art neural network based MD packages such as DeePMD-kit can only reach 4.7 nanoseconds per day on the Fugaku supercomputer. In this paper, we present a novel node-based parallelization scheme to reduce communication by 81%, then optimize the computationally intensive kernels with sve-gemm and mixed precision. Finally, we implement intra-node load balance to further improve the scalability. Numerical results on the Fugaku supercomputer show that our work has significantly improved the time-to-solution of the DeePMD-kit by a factor of 31.7x, reaching 149 nanoseconds per day on 12,000 computing nodes. This work has opened the door for millisecond simulation with ab initio accuracy within one week for the first time.

11 pages, 11 figures, 3 tables, SC'24

References in corpus (12)

• Deep Potential Molecular Dynamics: a scalable model with the accuracy of quantum mechanics
• DeePMD-kit: A deep learning package for many-body potential energy representation and molecular dynamics
• ANI-1: An extensible neural network potential with DFT accuracy at force field computational cost
• E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials
• A Spectral Analysis Method for Automated Generation of Quantum-Accurate Interatomic Potentials
• DeePMD-kit v2: A software package for Deep Potential models
• Ab initio theory and modeling of water
• SchNet: A continuous-filter convolutional neural network for modeling quantum interactions
• SpookyNet: Learning Force Fields with Electronic Degrees of Freedom and Nonlocal Effects
• The Individual and Collective Effects of Exact Exchange and Dispersion Interactions on the Ab Initio Structure of Liquid Water
• Isotope Effects in Liquid Water via Deep Potential Molecular Dynamics
• Extending the limit of molecular dynamics with ab initio accuracy to 10 billion atoms