FastCHGNet: Training one Universal Interatomic Potential to 1.5 Hours with 32 GPUs

Abstract

Graph neural network universal interatomic potentials (GNN-UIPs) have demonstrated remarkable generalization and transfer capabilities in material discovery and property prediction. These models can accelerate molecular dynamics (MD) simulation by several orders of magnitude while maintaining \textit{ab initio} accuracy, making them a promising new paradigm in material simulations. One notable example is Crystal Hamiltonian Graph Neural Network (CHGNet), pretrained on the energies, forces, stresses, and magnetic moments from the MPtrj dataset, representing a state-of-the-art GNN-UIP model for charge-informed MD simulations. However, training the CHGNet model is time-consuming(8.3 days on one A100 GPU) for three reasons: (i) requiring multi-layer propagation to reach more distant atom information, (ii) requiring second-order derivatives calculation to finish weights updating and (iii) the implementation of reference CHGNet does not fully leverage the computational capabilities. This paper introduces FastCHGNet, an optimized CHGNet, with three contributions: Firstly, we design innovative Force/Stress Readout modules to decompose Force/Stress prediction. Secondly, we adopt massive optimizations such as kernel fusion, redundancy bypass, etc, to exploit GPU computation power sufficiently. Finally, we extend CHGNet to support multiple GPUs and propose a load-balancing technique to enhance GPU utilization. Numerical results show that FastCHGNet reduces memory footprint by a factor of 3.59. The final training time of FastCHGNet can be decreased to \textbf{1.53 hours} on 32 GPUs without sacrificing model accuracy.

References in corpus (8)

• ANI-1: An extensible neural network potential with DFT accuracy at force field computational cost
• A Spectral Analysis Method for Automated Generation of Quantum-Accurate Interatomic Potentials
• A Universal Graph Deep Learning Interatomic Potential for the Periodic Table
• Atomistic Line Graph Neural Network for Improved Materials Property Predictions
• SchNet: A continuous-filter convolutional neural network for modeling quantum interactions
• SpookyNet: Learning Force Fields with Electronic Degrees of Freedom and Nonlocal Effects
• Scalable Parallel Algorithm for Graph Neural Network Interatomic Potentials in Molecular Dynamics Simulations
• GPTFF: A high-accuracy out-of-the-box universal AI force field for arbitrary inorganic materials