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

TL;DR

FastCHGNet is an optimized, GPU-efficient version of CHGNet that drastically reduces training time from over 8 days to under 2 hours using 32 GPUs, enabling rapid development of universal interatomic potentials.

Contribution

The paper introduces innovative modules and GPU optimizations that significantly accelerate CHGNet training without losing accuracy, supporting multi-GPU scaling.

Findings

Training time reduced to 1.53 hours on 32 GPUs

Memory footprint decreased by a factor of 3.59

Maintains accuracy while enabling faster training

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.