PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yutack Park, Jaesun Kim, Seungwoo Hwang, Seungwu Han

arXiv:2402.03789v1cond-mat.mtrl-sci

TL;DR

GNN-IPs are difficult to parallelize with spatial decomposition because message passing extends communication beyond local cutoffs, limiting scalable MD use. This paper proposes a compatible decomposition scheme and implements it in the NequIP-based SevenNet package. The modified scheme provides efficient scaling across message-passing depths, while GPU underutilization remains a practical limitation.

  • Problem

    GNN-IPs complicate spatial-decomposition parallelization because message passing extends communication beyond the cutoff radius, creating redundant computation and limiting available parallel implementations.

  • Method

    The paper develops a GNN-IP-compatible spatial-decomposition scheme that retains the original communication range while exchanging positions, node features, and gradient information, implemented in SevenNet.

  • Results

    The modified spatial decomposition offers efficient scaling performance regardless of the number of message-passing layers and supports large-scale amorphous Si3N4 simulation with a pretrained SevenNet model.

  • Takeaways & Limitations

    SevenNet provides a scalable route for applying equivariant GNN interatomic potentials to large-scale molecular dynamics simulations.

  • Takeaways & Limitations

    Strong-scaling speedup is limited by underutilized GPUs, especially for models with fewer channels.

Abstract

from arXiv · show

Message-passing graph neural network interatomic potentials (GNN-IPs), particularly those with equivariant representations such as NequIP, are attracting significant attention due to their data efficiency and high accuracy. However, parallelizing GNN-IPs poses challenges because multiple message-passing layers complicate data communication within the spatial decomposition method, which is preferred by many molecular dynamics (MD) packages. In this article, we propose an efficient parallelization scheme compatible with GNN-IPs and develop a package, SevenNet (Scalable EquiVariance-Enabled Neural NETwork), based on the NequIP architecture. For MD simulations, SevenNet interfaces with the LAMMPS package. Through benchmark tests on a 32-GPU cluster with examples of SiO$_2$, SevenNet achieves over 80% parallel efficiency in weak-scaling scenarios and exhibits nearly ideal strong-scaling performance as long as GPUs are fully utilized. However, the strong-scaling performance significantly declines with suboptimal GPU utilization, particularly affecting parallel efficiency in cases involving lightweight models or simulations with small numbers of atoms. We also pre-train SevenNet with a vast dataset from the Materials Project (dubbed `SevenNet-0') and assess its performance on generating amorphous Si$_3$N$_4$ containing more than 100,000 atoms. By developing scalable GNN-IPs, this work aims to bridge the gap between advanced machine learning models and large-scale MD simulations, offering researchers a powerful tool to explore complex material systems with high accuracy and efficiency.

1 INTRODUCTION

GNN-IPs offer accurate, data-efficient modeling of many-body interactions, but their extended message-passing neighborhoods complicate spatial-decomposition parallelization. The paper introduces SevenNet, a NequIP-based implementation designed for scalable MD simulations.

  • Motivation: MLPs infer energies and forces from DFT-labeled data and can scale linearly with atom count under a locality assumption.Large MLP simulations have reached billions of atoms with accuracy often nearing DFT calculations.
  • GNN-IPs: GNN-IPs represent atoms as graph nodes and nearby bonds as edges, using message passing to capture many-body and medium- or long-range ordering.Message passing extends geometric information beyond the cutoff radius.
  • GNN-IPs: Equivariant GNN-IPs change node features systematically under symmetry operations and can improve accuracy and data efficiency.NequIP uses spherical harmonics transformed by Wigner D-matrices and tensor products to capture many-body information without precomputed angles or dihedrals.
  • Parallelization challenge: Spatial decomposition is difficult for GNN-IPs because message passing requires communication out to the cutoff radius multiplied by the number of steps.Simply expanding communication creates redundant computation because neighboring processors share substantial subgraphs.
  • Contribution: The paper introduces a GNN-IP-compatible spatial-decomposition scheme and SevenNet, which exchanges positions, node features, and gradients while retaining the original communication range.SevenNet adopts the NequIP architecture and interfaces with LAMMPS.
  • Results: SevenNet achieves over 80% weak-scaling parallel efficiency and nearly ideal strong scaling when GPUs are fully utilized.Strong-scaling performance declines significantly with suboptimal GPU utilization.

2 GRAPH NEURAL NETWORK INTERATOMIC PO-

GNN interatomic potentials represent atoms and their local environments as graphs, using message passing to predict energies and forces. Their expanding receptive fields complicate spatial-decomposition parallelization, while equivariant architectures such as NequIP preserve symmetry-aware features.

  • GNN-IPs predict energy and atomic forces from atomic numbers and positions, with total energy expressed as a sum of atomic energies.Atomic forces are obtained as negative gradients of the energy with respect to atomic positions.
  • Nodes connect when interatomic distances fall below a predefined cutoff, and atomic positions define distance and displacement-vector edge features.Atomic numbers initialize node features in the embedding layer.
  • Message-Passing Layer: The receptive field expands with every message-passing layer, making spatial decomposition inefficient even though it works well for short-range force fields.This creates the central parallelization challenge for multi-layer GNN-IPs.
  • Message-Passing Layer: Each message-passing layer updates node features by combining messages gathered from neighboring nodes and edges, followed by a readout predicting atomic energies.The architecture stacks T message-passing layers, with message and update functions propagating features along connected edges.
  • NequIP Structure: NequIP is an E(3)-equivariant GNN-IP whose symmetry-aware features provide the foundation for SevenNet.SevenNet retains the original NequIP model outputs while adding communication routines for parallelization.

3 PARALLELIZATION of GNN-IP

The proposed spatial-decomposition algorithm distributes atoms across processors and communicates ghost-atom data through forward and reverse paths. SevenNet implements this scheme for LAMMPS while preserving the original NequIP architecture and outputs.

  • Spatial Decomposition: Spatial decomposition distributes smaller subdomains among processors, with each processor evaluating atomic energies and forces for its assigned atoms.Neighboring atoms on other processors are represented as ghost atoms.
  • Forward Communication: Forward communication first transfers ghost-atom positions and atomic numbers, then repeatedly transfers updated ghost-node features for later message-passing layers.The first layer requires no additional feature communication because initial ghost-node features come from atomic-number embeddings.
  • Forward Path: Atomic energies are globally pooled after the message-passing sequence to calculate the potential energy.The forward path contains T message-passing steps and T −1 forward communications of node features.
  • Reverse Path: Force computation reverses the network through energy gradients, using reverse communication for ghost-atom gradient and partial-force terms.The reverse path pairs T gradient-computing steps with message-passing steps and includes T −1 reverse communications of ∇hE.
  • SevenNet Implementation: SevenNet implements the parallel algorithm for LAMMPS while retaining the NequIP architecture and preserving the original model outputs.Its implementation uses PyTorch, e3nn, TorchScript, and separate TorchScript files for each message-passing layer.
  • SevenNet Implementation: SevenNet also supports multi-GPU training, serial stress computation, and fine-tuning when a pre-trained model lacks sufficient downstream accuracy.The analysis focuses on GPU-cluster parallel efficiency because neural-network computations are more efficient on GPU clusters.

4 BENCHMARK TESTS

The benchmarks evaluate SevenNet parallel performance under scaled-size and fixed-size conditions while varying communication size and frequency through model channels and message-passing layers. Weak scaling remains robust, whereas strong scaling is limited when GPU utilization becomes suboptimal.

  • Scaled-size test: Weak-scaling tests increase atom counts proportionally with GPU count, keeping 4,608 atoms per GPU and measuring efficiency as t(1)/t(n).The tests use up to 32 GPUs and a total of 147,456 atoms at the largest configuration.
  • Benchmark design: Twenty trained potentials vary channel count and message-passing depth to assess communication size, communication frequency, and parallel performance in α-quartz SiO2 MD.Channels control exchanged data size, while additional layers add forward and reverse communications.
  • Scaled-size test: 0.67–0.83 parallel efficiency with 32 GPUs is observed for 2-layer models, compared with 0.79–0.84 for 5-layer models.With 2 GPUs, the corresponding ranges are 0.83–0.98 and 0.95–1.0, respectively.
  • Scaled-size test: Parallel efficiency increases slightly with more channels or message-passing layers because greater computation makes communication costs relatively less significant.The authors attribute the effect to increases in both serial and parallel computation time as model complexity grows.
  • Fixed-size test: Strong-scaling performance saturates earlier for lightweight models, while more channels delay saturation; some 64-channel, 3–5-layer tests exceed A100 80 GB memory.The fixed-size tests use 12,960 atoms, nearly exhausting one A100 80 GB GPU for a 32-channel, 4-layer model.

5 GENERAL-PURPOSE GNN-IP

SevenNet-0 is a general-purpose pretrained GNN interatomic potential based on the NequIP architecture, with architectural changes that reduce parameters and mitigate overfitting. It achieves reported energy and force accuracy improvements over M3GNet and supports large-scale amorphous Si3N4 melt-quench simulations.

  • Model and pretraining: SevenNet-0 is developed as a general-purpose GNN-IP for multi-GPU MD without additional training and is pretrained using a large, chemically diverse dataset.The model is trained on the M3GNet dataset, which covers 89 elements.
  • Architecture: Replacing a NequIP self-connection tensor product with a linear layer reduces validation error and increases training error, suggesting reduced overfitting.Redundant tensor-product paths are also removed from the final interaction block, retaining scalar-producing paths.
  • Architecture: 16.24 million to 0.84 million parameters is the reduction reported when comparing GNoME with SevenNet-0.The reduction follows the architectural modifications to tensor-product connections and paths.
  • Accuracy: 24 meV/atom, 0.067 eV/Å, and 0.65 GPa are SevenNet-0’s test-set MAEs for energy, force, and stress, respectively.Compared with M3GNet, SevenNet-0 improves energy and force accuracies.
  • Structural validation: SevenNet-0 reproduces key amorphous Si3N4 structural features despite training solely on related crystal structures, demonstrating reported generalization beyond its training domain.The authors associate this capability with training on a large dataset containing diverse chemistry.

6 DISCUSSIONS

The proposed communication approach applies to GNN interatomic potentials whose message passing uses local neighborhoods, including equivariant models such as NequIP and MACE. The discussion also indicates that related architectures with explicit angular or edge features could use similar adaptations.

  • Applicability: Node-feature communication is required because each update uses the features of a node and its neighbors connected within the cutoff radius.This locality defines the information exchanged between processors.
  • Applicability: SchNet, PhysNet, PaiNN, NequIP, and MACE can be parallelized with the approach by adapting atomic-force computations to each architecture.The cited models use two-body message-passing functions, while equivariant models can still learn higher-than-two-body terms.
  • Related architectures: ALIGNN and CHGNet incorporate explicit three-body or angle features through edge-connected triplets and auxiliary line graphs.Their more complex message-passing schemes retain locality through a cutoff radius.
  • Related architectures: Similar parallelization could extend to these architectures by communicating edge features and gradients derived from them during forward and reverse passes.This is presented as an example of adapting the communication scheme to models with richer message passing.

7 CONCLUSION

The work proposes and implements a spatial-decomposition scheme compatible with GNN-IPs, with benchmarks showing efficient weak scaling and strong scaling dependent on GPU utilization. It also introduces SevenNet-0 and demonstrates its use in large-scale amorphous Si3N4 simulations.

  • The proposed spatial-decomposition scheme enables parallel molecular-dynamics simulations with GNN-IPs.
  • The modified spatial decomposition offers efficient weak-scaling performance regardless of message-passing layers and channel count.
  • Strong-scaling speed-up is limited by underutilized GPUs, especially for models with fewer channels.
  • SevenNet-0 is introduced as a general-purpose GNN-IP model, with large-scale results for generating amorphous Si3N4.
  • The SevenNet package is intended to facilitate GNN-IP use in a broad spectrum of large-scale molecular-dynamics simulations.

Corresponding Author

The supplied passage identifies Seungwu Han with the Department of Materials Science and Engineering at Seoul National University.

  • Seungwu Han is affiliated with Seoul National University’s Department of Materials Science and Engineering.

Authors

The supplied passages list the paper’s authors and affiliations, then describe how atom count and channel count affect single-GPU utilization.

  • Authors: Yutack Park, Jaesun Kim, and Seungwoo Hwang are listed among the paper’s authors.
  • Authors: The authors are affiliated with Seoul National University and the Korea Institute for Advanced Study.
  • GPU UTILIZATION: Figure S1 plots single-GPU utilization for four message-passing layers and channel counts of 4, 8, 16, 32, and 64.
  • GPU UTILIZATION: The utilization metric is the number of atoms multiplied by timesteps per second, representing work produced per second.
  • GPU UTILIZATION: A plateau indicates full GPU utilization, whereas lower-atom ballistic regions indicate suboptimal utilization.
  • GPU UTILIZATION: As channel count decreases, more atoms are needed to reach full GPU utilization.
Loading 2402.03789v1…