Session:

Machine-learned potentials have emerged as a preferred method for addressing the scale limitations inherent in ab initio molecular dynamics (MD) simulations. Early-stage machine-learned potentials utilized invariant, atom-centered descriptors to encode the atomic environment. Recently, however, attention has been drawn to graph-based neural network interatomic potentials (GNNs), especially those of the equivariant type. These GNNs generally surpass descriptor-based models in accuracy and can efficiently manage multi-component systems. However, the message propagation in GNNs inherently makes them challenging to integrate with spatial decomposition, a conventional parallel algorithm used in MD. To the best of our knowledge, no GNN package currently supports distributed-memory parallelism in MD. Consequently, the length scale that GNNs can study is restricted by the memory size of a single processor.
In implementing spatial decomposition of GNNs, two problems arise when two atoms connected by a graph edge belong to different processors. The first problem occurs when inferencing the atomic energy using the GNN. The message-passing layers in the GNN require atomic features of connected atoms, which become inaccessible due to spatial decomposition. The second problem arises when calculating the atomic forces by automatic differentiation derivatives of the energy (for instance, Autograd module in PyTorch). Force calculation requires gradient information of neighbor atoms, necessitating non-trivial evaluation and communication of partial derivatives for disconnected edges. To address these issues, we present SEVENN (Scalable EquiVariance-Enabled Neural Network), an implementation of a parallel algorithm that integrates GNNs with spatial decomposition. The baseline GNN model is the same as NequIP, and we employed the LAMMPS package for spatial decomposition. To optimize the parallel implementation, we have included a communication routine for GNNs with multiple GPUs, leveraging NVIDIA GPUDirect^® technology for fast inter-GPU communication.
To assess the parallelization efficiency of our algorithm, we conducted MD simulations for SiO₂ up to 100,000 atoms on A100 GPU clusters. We tested up to 5 message passing layers and 32 GPUs, where intensive communication of atomic features and gradient information between processors occurs in each message passing layer. SEVENN demonstrated over 80%-90% parallel efficiency when scaling the number of atoms proportionally with the number of GPUs. This result indicates that the simulation using SEVENN can be linearly scaled up to accommodate a large number of atoms. We believe that this study paves the way for the application of GNNs in large-scale simulations.

Co-Author(s): Jaesun Kim (Seoul National University), Seungwu Han (Seoul National University)

Yutack Park

Appearances