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 (18)

• Crystal Graph Convolutional Neural Networks for an Accurate and Interpretable Prediction of Material Properties
• 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
• Graph Networks as a Universal Machine Learning Framework for Molecules and Crystals
• Quantum-Chemical Insights from Deep Tensor Neural Networks
• A Spectral Analysis Method for Automated Generation of Quantum-Accurate Interatomic Potentials
• PhysNet: A Neural Network for Predicting Energies, Forces, Dipole Moments and Partial Charges
• 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
• Hierarchical modeling of molecular energies using a deep neural network
• Embedded Atom Neural Network Potentials: Efficient and Accurate Machine Learning with a Physically Inspired Representation
• Equivariant message passing for the prediction of tensorial properties and molecular spectra
• A Survey on Compiler Autotuning using Machine Learning
• 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