Scalable Parallel Algorithm for Graph Neural Network Interatomic Potentials in Molecular Dynamics Simulations

Area of Science:

• Computational Materials Science
• Machine Learning for Physics

Background:

• Message-passing graph neural network interatomic potentials (GNN-IPs), like NequIP, offer data efficiency and accuracy.
• Parallelizing GNN-IPs for molecular dynamics (MD) is challenging due to complex data communication in spatial decomposition.

Purpose of the Study:

• To propose an efficient parallelization scheme for GNN-IPs.
• To develop a scalable package, SevenNet, for large-scale MD simulations.
• To enable accurate and efficient exploration of complex material systems.

Main Methods:

• Developed SevenNet, a package based on the NequIP architecture, for GNN-IPs.
• Integrated SevenNet with the LAMMPS MD package.
• Conducted benchmark tests on a 32-GPU cluster using SiO2 and amorphous Si3N4 systems.

Main Results:

• SevenNet achieved over 80% parallel efficiency in weak-scaling scenarios.
• Demonstrated nearly ideal strong-scaling performance with full GPU utilization.
• Observed performance decline with suboptimal GPU utilization, especially for small systems or lightweight models.

Conclusions:

• SevenNet provides a scalable solution for GNN-IPs in large-scale MD simulations.
• The package facilitates the application of advanced machine learning models in materials research.
• SevenNet empowers researchers to study complex materials with high accuracy and efficiency.