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

1 INTRODUCTION

• Conventional DFT scales as O(N^3), whereas locality-based machine-learning potentials scale as O(N) and can simulate systems of up to billions of atoms with accuracy often nearing DFT.These data-driven potentials infer energies and forces from configurations using DFT-labeled training data.Source: §1 · p. 2
• Equivariant GNN-IPs are noted for accuracy and data efficiency; tensor-product message passing captures many-body information without precomputing geometric features such as angles and dihedral angles.Message passing also lets GNN-IPs account for medium- or long-range ordering.Source: §1 · p. 3
• Message passing expands GNN-IP receptive fields beyond the graph cutoff, requiring communication up to the cutoff times the number of steps; wider regions create redundant computation.These requirements make GNN-IPs difficult to parallelize with spatial decomposition.Source: §1 · p. 4
• SevenNet introduces a GNN-IP spatial-decomposition algorithm that keeps neighbor communication at the original cutoff, exchanges positions, node features, and gradients, and uses NequIP with LAMMPS.The package name expands to Scalable EquiVariance-Enabled Neural NETwork.Source: §1 · p. 4
• SevenNet exceeds 80% weak-scaling parallel efficiency and achieves nearly ideal strong scaling when GPUs are fully utilized, but performance declines significantly with suboptimal utilization.The reported results are from GPU-cluster benchmarks.Source: §1 · p. 4

TENTIAL

• GNN-IPs map atomic numbers and positions to energies and forces, summing atomic energies and defining forces as negative energy gradients.The atomic-energy decomposition is identified as key to linear scalability and to the validity of spatial decomposition for parallelization.Source: TENTIAL · p. 5

2.1 Message-Passing Layer

• Message-passing layers combine neighboring node and edge information, and each layer expands the receptive field, making spatial decomposition inefficient.The final readout predicts atomic energies from the last node features.Source: §2.1 · pp. 6–7

2.2 NequIP Structure

• SevenNet uses NequIP as its E(3)-equivariant foundation, adding communication routines without changing model outputs and omitting them from the first interaction block.The communication routines support parallelization while maintaining the integrity of the original NequIP model.Source: §2.2 · pp. 7, 9

3.1 Spatial Decomposition

• The cell is split among processors, which compute atom energies and forces while ghost atoms supply within-cutoff neighbor positions and features; forward communication sends updated node features between layers.Initial atomic positions and numbers are also communicated, as shown in Figure 3.Source: §3.1 · p. 9, Figure 3 · p. 10
• Force evaluation reverses through message-passing layers, communicating energy gradients for ghost nodes and partial forces; Figure 3 specifies T - 1 reverse gradient communications paired with feature forwards.Ghost-related terms are transferred by reverse communication to complete atomic force computation.Source: Figure 3 · p. 10, §3.1 · pp. 11–12

3.2 Implementation Details

• SevenNet packages this algorithm for LAMMPS using the NequIP architecture, PyTorch and e3nn equivariance, and TorchScript files paired with message-passing layers and LAMMPS communication routines.Source: §3.2 · p. 12
• SevenNet also supports serial stress calculations, multi-GPU training, and fine-tuning pretrained models when downstream accuracy is insufficient.Source: §3.2 · p. 12

4 BENCHMARK TESTS

• The experiments vary channel count and message-passing layer count; in each layer, processors exchange node features or energy gradients, with channel count setting communication-data size and layer count affecting communication frequency.Source: §4 · p. 13

4.1 Scaled-Size Test

• With 4,608 atoms per GPU, weak-scaling efficiency ranges from 0.67–0.83 at 32 GPUs for 2-layer models and 0.79–0.84 for 5-layer models.At 2 GPUs, the respective ranges are 0.83–0.98 and 0.95–1.0; efficiency is maintained despite increased communication demands.Source: Figure 4 · p. 15, §4.1 · p. 16

4.2 Fixed-Size Test

• With 32 GPUs, models configured with 32 channels and 2–5 message-passing layers achieve gains of 13.1–15.4x; even the best-performing model reaches only half the ideal 32x speed-up.Source: §4.2 · p. 17
• Strong-scaling gains saturate when GPUs are underutilized, because adding GPUs then does not reduce computation time.The weak-scaling tests used homogeneous bulk SiO2, which minimizes synchronization waits; more heterogeneous systems or more GPUs could increase synchronization costs.Source: §4.1 · p. 16, §4.2 · p. 17
• Channel count affects GPU utilization because it correlates with node-feature dimensionality, while adding message-passing layers requires only a small additional memory allocation.Source: §4.2 · p. 19
• Performance saturation occurs later with more channels in the fixed-size tests, as shown in Figure 5.Figure 5 covers 12,960 atoms and models with 4, 8, 16, 32, or 64 channels and 2–5 message-passing layers; some single-GPU times for 64-channel models were unavailable due to limited GPU memory.Source: Figure 5 · p. 18, §4.2 · p. 19
• A GNN-IP designed to use GPU resources efficiently with limited atoms is expected to improve speed gains when deployed across multiple GPUs.Source: §4.2 · p. 19

5.1 Model Training

• SevenNet-0 was trained on the M3GNet dataset of three relaxation steps for Materials Project crystals spanning 89 elements.Source: §5.1 · p. 19
• Replacing the self-connection tensor product with a linear layer reduced validation error while increasing training error, suggesting less overfitting; last-block path pruning cut parameters from 16.24 million in GNoME to 0.84 million.The retained paths in the last interaction block produce scalar outputs.Source: §5.1 · p. 20
• 24 meV/atom, 0.067 eV/Å, and 0.65 GPa were SevenNet-0's test-set energy, force, and stress MAEs, versus M3GNet's 35 meV/atom, 0.072 eV/Å, and 0.41 GPa.The reported comparison indicates better energy and force accuracy for SevenNet-0, but a higher stress MAE.Source: §5.1 · p. 21

5.2 Large-Scale Si3N4 Melt-Quench Simulation

• 0.94 parallel efficiency was achieved on 8 GPUs in the 14,000-atom weak-scaling run; a 112,000-atom, 60-ps simulation took 12.7 hours.The 14,000-atom run lasted 12.0 hours; the 112,000-atom simulation was infeasible on one A100 because of memory limits, making multi-GPU execution necessary.Source: §5.2 · p. 22
• Despite training only on crystal structures of the relevant composition, SevenNet-0 reproduced key amorphous Si3N4 features from melt-quench and predicted energies reasonably outside its training domain.The authors attribute this generalization capability to training on a large dataset with diverse chemistry.Source: §5.2 · pp. 22, 24
• Figure 6 presents amorphous Si3N4 with 112,000 atoms; its structure maintains fourfold coordination, and SevenNet-0 radial and angular distributions agree well with prior DFT reference data, including reproduced 90° shoulder peaks.Source: §5.2 · p. 22, Figure 6 · p. 23

6 DISCUSSIONS

• The general GNN-IP framework supports parallelizing two-body message-passing models—including SchNet, PhysNet, PaiNN, NequIP, and MACE—by adapting force computations to each architecture.The required communication arises because each node update uses features of the node and its neighbors within the cutoff radius.Source: §6 · p. 24
• Except for SchNet and PhysNet, the listed two-body models are equivariant and can learn more than two-body terms despite using simple pairwise message functions.This notes that the message function's pairwise form does not limit the interactions the equivariant models can learn.Source: §6 · p. 24
• Invariant models including DimeNet, ALIGNN, CHGNet, and M3GNet depart from the two-body framework by incorporating explicit three-body terms.ALIGNN and CHGNet handle angle (triplet) features connecting two edges with an auxiliary line graph.Source: §6 · p. 24
• Because cutoff radii preserve locality in these more complex schemes, similar parallelization could communicate edge features and their gradients in forward and reverse passes.The discussion proposes this as a possible adaptation for models such as ALIGNN and CHGNet.Source: §6 · p. 24

7 CONCLUSION

• The proposed spatial decomposition scheme enables large-scale MD simulations by adapting the original method to NequIP with additional forward and reverse communications.Source: §7 · pp. 24–25
• Weak scaling remains efficient regardless of the number of message-passing layers and channels.The conclusion identifies weak scaling as measuring performance as system size increases.Source: §7 · p. 25
• Strong-scaling speed-up is limited by GPU underutilization, especially for models with fewer channels.The strong-scaling benchmark increases processor count while keeping total problem size constant.Source: §7 · p. 25
• SevenNet-0 is introduced as a general-purpose GNN-IP model and demonstrated in large-scale generation of amorphous Si3N4.The pretrained model can be fine-tuned for downstream tasks to achieve higher accuracy.Source: §7 · p. 25
• The authors aim for SevenNet to facilitate GNN-IP use across large-scale MD applications.Source: §7 · p. 25

• The passage names an algorithm for graph neural networks.Source: Supporting Information for “Scalable… · p. 35

GPU UTILIZATION

• Figure S1 shows throughput plateauing at full GPU utilization with enough atoms, while lower-count ballistic behavior indicates suboptimal utilization.Source: Figure S1 · p. 36, GPU UTILIZATION · p. 36
• Throughput is measured as atoms multiplied by timesteps per second, representing GPU work produced per unit time.Source: GPU UTILIZATION · p. 36
• As channel count decreases, more atoms are needed to transition from suboptimal to full GPU utilization.Source: GPU UTILIZATION · p. 36