Efficient Parallelization of Message Passing Neural Network Potentials for Large-scale Molecular Dynamics

TL;DR

This paper introduces an efficient parallel algorithm for message passing neural network potentials, enabling large-scale molecular dynamics simulations with improved scalability and minimal data communication.

Contribution

The authors develop a scalable parallelization method for MPNN models that reduces communication overhead and allows simulations on over 100 million atoms.

Findings

Achieves excellent strong and weak scaling in benchmark systems.

Enables large-scale MD simulations with MPNN models comparable in speed to local models.

Extends MPNN applicability to unprecedented system sizes.

Abstract

Machine learning potentials have achieved great success in accelerating atomistic simulations. Many of them relying on atom-centered local descriptors are natural for parallelization. More recent message passing neural network (MPNN) models have demonstrated their superior accuracy and become increasingly popular. However, efficiently parallelizing MPNN models across multiple nodes remains challenging, limiting their practical applications in large-scale simulations. Here, we propose an efficient parallel algorithm for MPNN models, in which additional data communication is minimized among local atoms only in each MP layer without redundant computation, thus scaling linearly with the layer number. Integrated with our recursively embedded atom neural network model, this algorithm demonstrates excellent strong scaling and weak scaling behaviors in several benchmark systems. This approach enables massive molecular dynamics simulations on MPNN models as fast as on strictly local models for over 100 million atoms, vastly extending the applicability of the MPNN potential to an unprecedented scale. This general parallelization framework can empower various MPNN models to efficiently simulate very large and complex systems.