Ex-NNQMD: Extreme-Scale Neural Network Quantum Molecular Dynamics

Deep learning is revolutionizing countless scientific and engineering fields. In particular, SC20 Gordon Bell award represented a breakthrough in molecular simulation, i.e., 100-million-atom simulation with quantum-mechanical accuracy on the Summit supercomputer at ORNL, using deep potential molecular dynamics (MD). Moving forward, while these simulations were performed only in gentle equilibrium conditions, far-from-equilibrium MD simulation involving light-induced electronic excited states finds numerous scientific and engineering applications. However, it remains a challenge to perform such far-from-equilibrium simulations at larger spatiotemporal scales, where growing number of unphysical predictions of interatomic force prohibits simulations involving larger numbers of atoms for longer times. In this paper, we propose a physically-based inductive bias, maximally-preserved Maxwell-Boltzmann (MPMB), to overcome this fidelity-scaling problem. Along with hybrid divide-and-conquer parallelization and single-node level optimization using multithreading and data parallel SIMD, the resulting Ex-NNQMD (extreme-scale neural network quantum molecular dynamics) algorithm has achieved unprecedented scales of far-from-equilibrium simulations: 1) 5.1-billion atom system with a parallel efficiency of 0.94, and 2) a sustained performance of 6.4 nanoseconds/day for 10-million atom system both on 262,144 cores of the Theta supercomputer at Argonne Leadership Computing Facility. Extended fidelity scaling and efficient parallelization have allowed us for the first time to study light-induced ferroelectric switching under extreme electronic excitation at experimentally relevant spatiotemporal scales with accuracy.