AtomFlow: Accelerating GNN-Based Molecular Dynamics with ab initio Accuracy on Fugaku Supercomputer
Molecular dynamics (MD) simulations with ab initio accuracy have become indispensable for studying microscopic phenomena. Although graph neural network (GNN)-based MD models provide higher accuracy than deep neural network (DNN)-based models, their complex message-passing architectures introduce substantial computational and communication bottlenecks for large-scale inference. In this paper, we propose AtomFlow, a holistic acceleration for a representative GNN-based MD model on the Fugaku supercomputer. Our key optimizations include: (1) we design a fused domain-specific EC kernel that avoids materializing large edge-related intermediate tensors; (2) we design an inter-process communication operator that leverages Fugaku’s uTofu interconnect to accelerate atomic-feature exchange; and (3) we redesign the inference workflow with a framework-free implementation to avoid the overhead of the PyTorch framework on Fugaku. Experimental results show that AtomFlow achieves up to a 9.77 × speedup over the baseline implementation, reaching up to 26.14 nanoseconds per day for a 576,000-atom system with 12,000 nodes on Fugaku.