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

arXiv: https://arxiv.org/abs/2412.20796 · HTML: https://mlsanity.com/paper/2412.20796/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 40 references that are in the map.
- First posted on arXiv 30 Dec 2024.

## Builds on (most cited first, 18 of 18)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [SchNet: A continuous-filter convolutional neural network for modeling quantum interactions](https://mlsanity.com/paper/1706.08566/) (2017)
- [Directional Message Passing for Molecular Graphs](https://mlsanity.com/paper/2003.03123/) (2020)
- [MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields](https://mlsanity.com/paper/2206.07697/) (2022)
- [Equivariant message passing for the prediction of tensorial properties and molecular spectra](https://mlsanity.com/paper/2102.03150/) (2021)
- [GemNet: Universal Directional Graph Neural Networks for Molecules](https://mlsanity.com/paper/2106.08903/) (2021)
- [Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules](https://mlsanity.com/paper/2011.14115/) (2020)
- [Ansor: Generating High-Performance Tensor Programs for Deep Learning](https://mlsanity.com/paper/2006.06762/) (2020)
- [SpookyNet: Learning Force Fields with Electronic Degrees of Freedom and Nonlocal Effects](https://mlsanity.com/paper/2105.00304/) (2021)
- [Degree-Quant: Quantization-Aware Training for Graph Neural Networks](https://mlsanity.com/paper/2008.05000/) (2020)
- [DeePMD-kit: A deep learning package for many-body potential energy representation and molecular dynamics](https://mlsanity.com/paper/1712.03641/) (2017)
- [A Survey on Compiler Autotuning using Machine Learning](https://mlsanity.com/paper/1801.04405/) (2018)
- [VQ-GNN: A Universal Framework to Scale up Graph Neural Networks using Vector Quantization](https://mlsanity.com/paper/2110.14363/) (2021)
- [Hierarchical modeling of molecular energies using a deep neural network](https://mlsanity.com/paper/1710.00017/) (2017)
- [SGQuant: Squeezing the Last Bit on Graph Neural Networks with Specialized Quantization](https://mlsanity.com/paper/2007.05100/) (2020)
- [QGTC: Accelerating Quantized Graph Neural Networks via GPU Tensor Core](https://mlsanity.com/paper/2111.09547/) (2021)
- [Operator Fusion in XLA: Analysis and Evaluation](https://mlsanity.com/paper/2301.13062/) (2023)

## Built on by (most cited first, 2 of 2)

- [A practical guide to machine learning interatomic potentials -- Status and future](https://mlsanity.com/paper/2503.09814/) (2025)
- [Facet: highly efficient E(3)-equivariant networks for interatomic potentials](https://mlsanity.com/paper/2509.08418/) (2025)

## Abstract (arXiv)

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 textitab 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 1.53 hours on 32 GPUs without sacrificing model accuracy.