Extending the limit of molecular dynamics with ab initio accuracy to 10 billion atoms

TL;DR

This paper presents a comprehensive set of algorithmic and system innovations that significantly extend the scale and speed of ab initio molecular dynamics simulations to over 17 billion atoms on supercomputers.

Contribution

It introduces new optimization techniques and system-level strategies that enable large-scale ab initio molecular dynamics simulations with unprecedented system sizes and efficiency.

Findings

Simulations scaled up to 17 billion atoms on Fugaku and Summit.

Achieved 11.2 nanoseconds per day for a 13.5-million atom system.

Code can utilize entire supercomputers with 7 times faster performance.

Abstract

High-performance computing, together with a neural network model trained from data generated with first-principles methods, has greatly boosted applications of \textit{ab initio} molecular dynamics in terms of spatial and temporal scales on modern supercomputers. Previous state-of-the-art can achieve $1-2$ nanoseconds molecular dynamics simulation per day for 100-million atoms on the entire Summit supercomputer. In this paper, we have significantly reduced the memory footprint and computational time by a comprehensive approach with both algorithmic and system innovations. The neural network model is compressed by model tabulation, kernel fusion, and redundancy removal. Then optimizations such as acceleration of customized kernel, tabulation of activation function, MPI+OpenMP parallelization are implemented on GPU and ARM architectures. Testing results of the copper system show that the optimized code can scale up to the entire machine of both Fugaku and Summit, and the corresponding system size can be extended by a factor of $134$ to an unprecedented $17$ billion atoms. The strong scaling of a $13.5$-million atom copper system shows that the time-to-solution can be 7 times faster, reaching $11.2$ nanoseconds per day. This work opens the door for unprecedentedly large-scale molecular dynamics simulations based on {\it ab initio} accuracy and can be potentially utilized in studying more realistic applications such as mechanical properties of metals, semiconductor devices, batteries, etc. The optimization techniques detailed in this paper also provide insight for relevant high-performance computing applications.