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

Abstract

Machine learning potentials have achieved great success in accelerating atomistic simulations, among which message passing neural networks (MPNNs) have become increasingly prevalent thanks to their superior accuracy. However, MPNN potentials are difficult to be efficiently parallelized due to their semilocal architecture. Here, we propose an efficient scheme for parallelizing MPNN potentials, in which additional data communication is minimized among local atoms only in each layer without redundant computation, thus scaling linearly with the message-passing depth. This scheme is first tested on several bulk systemsincluding silver, liquid water, and high-entropy alloysdemonstrating excellent strong- and weak-scaling performance up to over 100 million atoms, extending MPNN applicability to an unprecedented scale. Moreover, we develop a universal potential trained on a comprehensive data set of C, H, O, and N, and leverage this parallel algorithm to perform efficient reactive MD simulations of graphene formation via acetylene–oxygen detonation, a central process in carbon nanomaterial synthesis. With atomic resolution, our simulations reveal how oxygen actively mediates the reaction network and identify the O₂/C₂H₂ ratio that optimally favors graphene-precursor formation, offering mechanistic insights into the experimentally observed necessity of oxygen for controlling graphene quality and yield. The proposed parallelization framework can be readily extended to other MPNN potentials.

Introduction

Recent years have witnessed the revolutionary development of machine learning (ML) methods in representing potential energy surfaces (PESs) of complex systems based on ab initio calculations. − These machine-learned potentials (MLPs) have achieved much higher flexibility than empirical force fields with much lower computational costs than on-the-fly ab initio molecular dynamics (AIMD) simulations. One particularly successful family of MLPs is based on an atom-wise representation of the PES. The total energy of the system is expressed as a sum of individual atomic energies, each of which is determined by many-body features within an atom-centered environment given by a cutoff radius (r _c ). Consequently, such atom-wise MLPs exhibit linear scaling with system size, ensuring high scalability to extended systems. Their local feature-based nature also allows for straightforward parallelization. , The energy evaluation can be efficiently distributed across multiple processors in an atom-wise fashion, eliminating the need for costly interatomic communication.

More recently, message passing neural network (MPNN) based approaches have become increasingly popular for learning PESs in an end-to-end manner. ,,,− ,,,− In this framework, the system is represented as an atomistic graph constructed by nodes (atoms) and edges (internode connections), in which atomic features are formed by iteratively propagating the geometric message from neighboring nodes via their edge features. After several message-passing (MP) iterations (T, often 2 ≤ T ≤ 6), MPNNs effectively learn both local high-order many-body interactions and certain semilocal information beyond the original r _c, thus generally achieving superior accuracy compared to conventional local MLP models. Initial MPNN models propagated symmetry-invariant functions of interatomic distances or angles. ,, More advanced versions now incorporate equivariant features that include directional information, enabling them to capture orientation-dependent atomic interactions. ,− ,− These developments greatly improve the structural discrimination, data efficiency, and transferability of MPNNs. Consequently, MPNN-based potentials have achieved great success in recent applications, − including protein-scale systems comprising tens to hundreds of thousands of atoms. −

Despite these advantages, MPNN models still face inherent limitations in parallel scalability, which restricts their application to large-scale simulations involving millions of atoms and beyond. This is because message communication occurs over T sequential steps, effectively expanding the neighborhood radius to Tr _c . Consequently, all atoms within this effective T-dependent cutoff r _c,T contribute to the final state of a central atom. It is commonly argued that the parallel computational cost of MPNNs scales unfavorably with r _c,T or equivalently T, due to the cubically growing number of interacting atoms. ,− To circumvent this difficulty, Kozinsky et al. proposed Allegro, a strictly local equivariant architecture that confines information exchange within the original cutoff radius. This simplifies its parallelization to be the same as other local MLPs. Alternatively, the MACE model proposed by Csanyi et al. reduces T to a maximum of two by utilizing high body-order equivariant messages to limit the parallel cost. The parallelization schemes of Allegro and MACE largely rely on the conventional domain decomposition, which is not optimal for MPNN architectures with semilocal message passing. Park et al. proposed SevenNet, a spatial decomposition algorithm based on the NequIP architecture, which restricts the communication range between neighboring processors up to r _c to limit redundant computations. However, this restriction prevents it from using an additional skin distance to reduce the frequency of updating neighbor lists, which could lower the efficiency of MD simulations in highly movable systems.

In this work, we propose a conceptually novel, general and efficient parallelization scheme for MPNN models. A key concept is that the additional communication occurs only among local neighboring atoms in each MP layer, achieving a linear scaling rather than a cubic scaling with respect to T. This scheme is implemented with a simple and physically inspired MPNN model, namely recursively embedded atom neural network (REANN) method. Numerical tests on the various periodic systems validate the high efficiency of this scheme, including silver, liquid water, and high-entropy alloys. Our implementation shows strong scaling with the number of CPUs/GPUs and weak scaling with the number of atoms, and enables parallel molecular dynamics simulations of semilocal MPNNs up to ∼100 million of atoms for the Ag system. Leveraging this methodologic advance, we have additionally developed a universal potential for C, H, O, N elements. Its successful application to highly parallelized large-scale reactive MD simulations of graphene formation via acetylene–oxygen detonation demonstrates the potential of applying MPNN potentials to model large systems.

Methods

REANN Architecture

As a simple MP adaptation of a local descriptor-based embedded atom neural network (EANN) framework, the REANN modelillustrated in Figure serves as an ideal case study for comparative analysis of parallelization strategies of MPNN and strictly local MLP models. Specifically, we start by applying physically motivated embedded atom density (EAD) features to characterize local atomic environments. An EAD feature of a central atom i is defined by the square of the linear summation of orbital functions located at all neighboring atoms within r _c ,

$${\rho }_i=\sum _{l=0}^L\sum _{l_x,l_y,l_z}^{l_x+l_y+l_z=l}\frac{l!}{l_x!l_y!l_z!}{[\sum _{j\ne i}^{N_{\mathrm{c}}}{c}_j{G}_j({\mathbf{r}}_{ij})]}^2$$ 1

where N _c is the number of atoms inside the cutoff sphere, G _j (r _ij ) is an orbital function of the jth neighbor atom obtained by the contraction of N _φ primitive Gaussian-type orbitals (see the mathematical details of the REANN model section in Supporting Information), each of which is characterized by the hyperparameters of its center, width, and angular momenta (l = l _x + l _y + l _z and 0 ≤ l ≤ L), and c _j is the corresponding element-wise orbital coefficient. Varying these hyperparameters forms a vector of EAD features (ρ _i ), which encodes three-body interactions implicitly and serves as the input of atomic NNs outputting atomic energies in the EANN model. Interestingly, the element-dependent orbital coefficient c _i itself should inherently rely on the neighborhood of atom i, which can be also expressed as the output of an atomic NN analogous to the atomic energy. Inspired by the concept of iterative message-passing, this relationship can be written in a vectorized form, for the tth iteration

$${\mathbf{c}}_i^t=M_i^{t-1}({\mathit{\rho }}_i^{t-1}(\{{\mathbf{c}}_{j\ne i}^{t-1}\},\{{\mathbf{G}}_{j\ne i}\})),1\le t\le T-1$$ 2

where ρ _i ({c _{j≠ i} }, {G _j≠i }) is simplified as a function of the collection of orbital functions ({G _j≠i }) and coefficients ({c _{j≠ i} }) of neighbor atoms (see eq ), which serves as the input of the atomic NN, M _i , to yield the central atom’s orbital coefficient (c _i ) that can be used to update ρ _i for the next cycle. All vectors here have the same dimension of N _φ(N _φ + 1)­(L + 1)/2. Consequently, ρ _i and c _i are iteratively updated via eqs and in an end-to-end MP manner, corresponding to node features and messages in the MPNN terminology. After T layers of MP, the last atomic NN layer outputs the atomic energy, whose sum yields the total energy E. In this sense, REANN is a three-body invariant feature-based MPNN model, which reduces to a strictly local EANN model when T = 1.

1.

Schematic diagram of the REANN framework. (a) Input and MP layers, (b) each MP layer where ⊗ denotes the tensor product, and (c) the output layer after T MP iterations.

Parallelization Algorithm

Let us first address the scaling of MPNN models with respect to the number of MP layers, which was not clearly discussed in literature. It is well-known that the evaluation cost of local MLPs of condensed-phase systems scales cubically with respect to r _c due to the cubic growth of neighbors in the cutoff sphere. Because the effective cutoff radius in MPNNs incrementally increases with the layer number of MP (see Figure a), it was argued that the feature evaluation cost would scale with t ³ r _c ³ in the tth MP layer, resulting in an actual accumulative scaling of $\sim \mathcal{O}(T^4)$ and significant demanding of communication overhead across multiple processes in parallel computation. In practice, however, the MP procedure is sequential and in each MP layer interatomic interactions remain confined to an atom-centered sphere within r _c. Critically, the neighbor list for each atom is calculated only oncein the first layer. Once computed, it can be reused across subsequent MP layers without modification, as the configuration does not change during message passing. No need to maintain an extensive neighbor list as T increases. As a result, the overall computational cost and communication overhead of the MPNN model can scale as favorably as with $\sim \mathcal{O}(T)$ . This is the key prerequisite for efficient parallelization of MPNN models.

2.

Illustration of the growth of neighbor atoms via message passing. (a) False impression considering all neighbor atoms of a fixed central atom as the effective cutoff radius increases. (b) Actual growth of the number of neighbor atoms in which each neighbor atom in the previous layer becomes a central atom of the current layer with the same cutoff radius. Supposing r _c is the original cutoff radius and T = 2.

Figure a illustrates schematically the conventional spatial decomposition scheme employed in the mainstream MD softwares , for parallel simulations based on atomistic force fields. In each MD step, the simulation box is partitioned into nonoverlapping subdomains, with each subdomain assigned to a unique process. Each process stores (i) positions/velocities and a full set of neighbor lists of its assigned local atoms (LAs), e.g., atoms 1 and 2 (3 and 4) in process P ₁ (P ₂); (ii) positions/velocities for a group of ghost atoms (GAs), which lie within the cutoff sphere of its LAs but belong to adjacent subdomains, e.g., atom 3 (2) in P ₁ (P ₂). Atoms are reassigned to new processes as they move through the physical domain. This design ensures each process to compute interactions which involve assigned LAs, but have accessibility to all necessary neighbors within the specified cutoff radius through GAs, enabling correct energy and force calculations across subdomain boundaries. Figure b shows clearly how this scheme works for the local EANN model via a computational graph, where downward and upward arrows represent forward (energy) and backward (force) workflows, respectively. Taking the forward inference as an example, all atomic coordinates (and coefficients) associated with GAs have been preassigned in each process in the neighbor list of LAs, enabling independent computation of atomic features of LAs and thus atomic energies.

3.

Failure of the commonly used parallel algorithm in MPNN models. (a) Schematic diagram for domain decomposition in a toy system where local atoms and ghost atoms are respectively located in white and gray areas, with neighboring atoms connected. Computational graphs of (b) local models and (c) MPNN models (T = 2) using the commonly used parallel algorithm with two processes, P ₁ and P ₂. Forward (backward) evaluations between nodes (variables) are depicted by downward (upward) arrows. Nodes in the forward evaluation represent the calculation of variables themselves, while in the backward evaluation, they represent a set of derivatives of the previous nodes with respect to them. The red color indicates incorrect values due to the lack of the correct data transfer.

However, as illustrated in Figure c, nonlocal MPNN models cannot directly adopt this parallel algorithm. This is because orbital coefficients of GAs used for the feature generation of LAs in the current layer are virtually evaluated in the previous MP layer and distributed across different processes. They are messages that need to be passed from other processes rather than being directly calculated in the process where the neighbor lists of GAs are incomplete. For example, the incomplete neighbor list of the atom 3 as a GA in process P ₁ would lead to a false feature vector, ρ ₃ , and then cause continuous errors in relevant variables (marked in red). A similar issue arises in the backward evaluation for atomic forces, and further exacerbated by the chain broadcasting of erroneous tensor differentiation.

To address this challenge, we propose an efficient parallel algorithm for MPNN models. As illustrated in Figure a, for the forward evaluation of the tth MP in each process, the messages of GAs are synchronized from the processes where they are classified as LAs and their features are correctly calculated. Meanwhile, the features of GAs in each MP layer are stored in relevant processes to keep the computational subgraph intact, enabling data communication and gradient calculations in the backward evaluation. As shown in Figure b, atomic forces are calculated based on the chain rule of tensor differentiation in the backward evaluation. For atom k, the contributions to its force F _k can be evaluated in parallel within N _p processes

$${\mathbf{F}}_k=-\sum _p^{N_{\mathrm{p}}}\sum _{i\in L_{\mathrm{p}}}\frac{\partial E_{\mathrm{p}}}{\partial {\mathit{\rho }}_i^{T-1}}\frac{\partial {\mathit{\rho }}_i^{T-1}}{\partial {\mathbf{r}}_k}$$ 3

in which L _l,p is the set of local atoms in process p and E _p is the sum of their energies. Specifically, for the tth MP iteration in process p, we have

$$\frac{\partial {\mathit{\rho }}_i^t}{\partial {\mathbf{r}}_k}=\sum _{j\ne i}^{N_c}(\frac{\partial {\mathit{\rho }}_i^t}{\partial {\mathbf{c}}_j^t}\frac{\partial {\mathbf{c}}_j^t}{\partial {\mathbf{r}}_k}+\frac{\partial {\mathit{\rho }}_i^t}{\partial {\mathbf{G}}_j}\frac{\partial {\mathbf{G}}_j}{\partial {\mathbf{r}}_k})$$ 4

$$\frac{\partial {\mathbf{c}}_j^t}{\partial {\mathbf{r}}_k}=\frac{\partial {\mathbf{c}}_j^t}{\partial {\mathit{\rho }}_j^{t-1}}\frac{\partial {\mathit{\rho }}_j^{t-1}}{\partial {\mathbf{r}}_k}$$ 5

where $\frac{\partial {\mathit{\rho }}_i^t}{\partial {\mathbf{G}}_j}$ and $\frac{\partial {\mathbf{G}}_j}{\partial {\mathbf{r}}_k}$ can be calculated locally in their respective processes, while only the gradient components associated with neighbor processes need data transfer, e.g., $\frac{\partial {\mathit{\rho }}_2^t}{\partial {\mathbf{c}}_3^t}$ and $\frac{\partial {\mathbf{c}}_2^t}{\partial {\mathit{\rho }}_2^{t-1}}$ from P ₁ to P ₂ (note that when t = 0, $\frac{\partial {\mathbf{c}}_j^0}{\partial {\mathbf{r}}_k}=0$ ). There is no need to synchronize the entire gradient tensor or the related computational subgraphs. In our implementation, both computational and communication overhead of gradients are further reduced by employing the vector-Jacobian product instead of explicitly computing the full Jacobian matrix. For example, in this way, the size of $\frac{\partial {\mathit{\rho }}_2^t}{\partial {\mathbf{c}}_3^t}$ can be reduced by N _φ(N _φ + 1)­(L + 1)/2 times. This scheme can also be extended to communicate gradient data associated with complex tensors in equivariant MPNN models, significantly lowering the communication bandwidth requirement. This parallel algorithm is done via the Message Passing Interface (MPI) based on the REANN architecture, referred to as REANN–MPI hereafter (see the Software section in Supporting Information for details).

4.

Computational graphs of the proposed MPI parallel algorithm. (a) Forward and (b) backward evaluations of MPNN models shown to, using with REANN as an example for illustration. Curved arrows indicate the interprocess data communications, solid circles denote variables that require data transfer, dotted circles indicate variables preserved for data communication of gradients in the backward evaluation of the computational subgraph.

The proposed algorithm offers a general framework for parallelizing MPNN models. It is not limited to specific feature representations or MP schemes, requiring only the communication of GA-related data, either features or messages. In this sense, any MPNN models with messages aggregated over neighboring atoms, either based on invariant (e.g., SchNet and AIMNet) or equivariant (e.g., NequIP and MACE) features, can adopt a similar data communication scheme for their GA-related data in forward and backward evaluations. In comparison, the parallel scheme in SevenNet is based on the communication of node features. For MPNN models where node features are updated by message pooling from neighboring atoms like REANN, SO3KRATES, or SEGNN, this scheme would incur repeated computation of GA message-related terms in both forward and backward passes across processes, resulting in noticeable increase of the parallelization overhead, even if the communication range is limited within r _c . Interestingly, to enable parallelization via conventional spatial decomposition, MACE treats GAs as independent nodes, storing their complete neighbor-related information in each process. This allows GAs to evaluate both feature- and message-related terms, supporting the evaluation of LAs. Conceptually similar to the parallel scheme in MACE, both chemtrain-deploy and FeNNix-Bio1 adopt domain decomposition strategies. Specifically, to enable complete message-passing evaluation for all LAs, chemtrain-deploy includes all GAs within a communication radius R _comm ≥ Tr_c in each process, whereas FeNNix-Bio1 recomputes messages between atoms in neighboring domains on each GPU. While these schemes alleviate data communication, they would introduce substantial repetitive computation for GAs and increase the parallelization overhead.

Results and Discussion

Linear Scaling with Message Passing

We first validate the proposed parallel scheme by extensive MD simulations for solid Ag and liquid H₂O systems, which have been frequently used to benchmark the performance of various MLP models. , The Ag data set consists of 1159 configurations of 71 Ag atoms with energies and atomic forces computed by density functional theory (DFT) using the PBE functional, while the H₂O data set comprises 1593 configurations of 64 H₂O molecules calculated by DFT at the revPBE0-D3 level (see the systems and training details section in Supporting Information (SI)). Figure shows the relative cost of the REANN–MPI as a function of the number of MP layers for these two benchmark systems. To assess the actual computational cost, including communication overhead, the simulation box in each case is spatially decomposed into two equal-sized subdomains along each dimension, with each subdomain processed by a single CPU core. Figure offers ambiguous numerical evidence that a properly parallelized MPNN model scales linearly with T regardless of the system size or composition. The observed scaling behavior slightly deviates from the ideal relationship for systems comprising 1 K atoms or fewer, where communication overhead becomes dominant. These findings overthrow the previous argument on the unfavorable scaling of $\sim \mathcal{O}(T^3)$ for most MPNN architectures ( $\sim \mathcal{O}(T^4)$ considering the accumulative cost).

5.

Validating the model scaling with respect to message passing. Comparison of relative computational costs of MD simulations for (a) Ag and (b) H₂O using the REANN–MPI model, as a function of T, relative to the baseline cost of T = 1. The dashed lines indicate several possible scalings with respect to T, namely $\sim \mathcal{O}(T)$ , $\sim \mathcal{O}(T^3)$ and $\sim \mathcal{O}(T^4)$ . The number of atoms ranges from ∼1 thousand (K) to ∼1 million (M).

Strong Scaling

Figure a demonstrates the strong scaling behavior of REANN–MPI, namely the computational speed versus the number of CPU cores (ranging from 1 to 1000) in MD simulations of liquid water (1.92 million atoms). Note that every CPU node in the test contains two Intel Xeon 9242 CPUs and 96 cores in total. Also shown are the results of the native just-in-time (JIT) implementation of REANN (REANN–JIT), which is a widely adopted parallel algorithm intrinsically supported in most Python-based platforms. It should be noted that REANN–JIT allows the internode MPI parallelization only for local models (T = 1). In such a case, both REANN–MPI and REANN–JIT demonstrate comparable parallel performance, achieving near-linear speedups scaling to 1000 CPU cores. This implies that our new parallel algorithm is compatible to local models while introducing negligible extra expense. When T > 1, REANN–JIT is only compatible with the multithreaded computation in a single node using pen multi-processing (OpenMP), which suffers from a very low parallel efficiency with many cores. By contrast, the REANN–MPI model retains a strong scalability and high efficiency, e.g. ∼90% with 500 cores and ∼80% with 1000 cores, as shown in Figure S1 in Supporting Information. This clearly demonstrates that the superiority of REANN–MPI originates from our newly designed parallel scheme for MPNN, rather than from refactoring the code with C++. We note in passing that energies and forces evaluated by REANN–MPI and REANN–JIT are identical, validated by the same radial distribution functions of liquid water, which agrees well with experiments (Figure S2). Figure S3 compares the total energy conservation of REANN–JIT and REANN–MPI implementations with various T, obtained from classical NVE MD simulations of 64 H₂O molecules. It shows that REANN–MPI preserves the total energy comparably well as the native REANN–JIT during the NVE MD simulation, regardless of the number of MP layers, validating the correct implementation of the parallelization scheme.

6.

Strong scaling and weak scaling. (a) Comparison of speed (T = 1 and T = 2) of REANN–MPI and REANN–JIT models as a function of CPU cores when simulating liquid water (1.92 million atoms). (b) Weak scaling performance of REANN–MPI in simulating both bulk Ag and liquid H₂O, where the number of atoms per node remains constant while the number of nodes increases from 1 to 32.

Weak Scaling

Figure b further illustrates the scalability of the REANN–MPI model for both solid Ag and liquid H₂O systems, where the number of nodes increases up to 32 while the number of atoms distributed in each node keeps fixed roughly from one million to three million. In each set of calculations, the actual speeds are nearly constant with the increasing number of nodes, demonstrating the excellent weak scaling behavior of the REANN–MPI model. Its absolute speed is also very fast, e.g., reaching 7.4 × 10^–8 s/step/atom for over 100 million Ag atoms using 32 nodes with roughly 3 million atoms per node. This system size significantly surpasses the maximum scales previously achieved with semilocal MPNN architectures, , matching the computational tractability of state-of-the-art strictly local MPNN models such as Allegro. , In qualitative comparison with a well-recognized efficient local NN model, moreover, REANN–MPI is faster than the deep potential (DP) model under similar conditions in the simulation of liquid water, e.g., 4.0 × 10^–7 s/step/atom for ∼12.9 million of atoms vs 5.8 × 10^–7 s/step/atom for ∼12.6 million of atoms, even though the latter utilizes ∼4 times more CPU cores (768 vs 3360). These results highlight the superior efficiency of our model and our parallel algorithm.

Performance on GPUs

The new parallel scheme naturally applies to GPU, since the implementation is based on LibTorch which supports both CPU and GPU. Figure a compares the absolute speeds of the REANN–MPI and Allegro models for roughly one million Ag atoms, with the number of GPUs increasing from 1 to 8. Note that the GPU node in the test equips 8 NVIDIA A100 GPUs, each with 80 GB memory. The speed of REANN–MPI on a single GPU (5.2 × 10^–1 steps/s) is significantly faster than that on a single CPU core (1.1 × 10^–2 steps/s) for simulations of 1 million Ag atoms. Impressively, our REANN–MPI model is in general 2 orders of magnitude faster the strictly local Allegro model. As shown in Figure b, both models exhibit a similar strong scaling relationship, with linear speedup with respect to the number of GPUs up to 8. REANN–MPI shows a slightly lower speedup than Allegro, given that the latter is essentially a local MPNN model relying on conventional spatial decomposition and thus does not require the message-related data communication as semilocal MPNNs. The slight decline in the speedup of REANN–MPI with more GPUs may also be partially attributed to the increasing cost of data transfer between the GPU and CPU, which can be migrated by direct data communication between GPUs in the future.

7.

Comparison of the (a) absolute speeds and (b) relative speedups between REANN–MPI and Allegro in the simulations of Ag (1,022,400 atoms) with the increasing number of NVIDIA A100 GPUs and comparison of the (c) absolute speeds and (d) parallel efficiencies between REANN–MPI, SevenNet, MACE with its native serial (MACE-s) and parallel (MACE-p) implementations, in the simulations of HEA with 576 atoms/GPU, with the increasing number of NVIDIA H200 GPUs. The speed data of Allegro were taken from the ref . Speed of SevenNet and MACE were evaluated using our retrained models based on their settings by default with comparable model size.

Figure c demonstrates comparison between REANN–MPI, SevenNet and, the native serial (MACE-s) and parallel (MACE-p) implementations of MACE with comparable model sizes in the simulation of a more challenging high entropy alloy (HEA) system consisting of five transition metal elements, with the number of GPUs increasing from 1 to 32. The corresponding HEA data set comprises 25,500 structures containing 3–25 transition metal elements, with energies and forces computed by DFT using the PBEsol functional (see the systems and training details section in Supporting Information). Note that the GPU node in this test equips 8 NVIDIA H200 GPUs, each with 141 GB memory. When loading 576 atoms/GPU, REANN–MPI is 2.6–4.7 times as fast as SevenNet and 14.8–19.3 times the MACE model, exhibiting an excellent weak scaling behavior from 1 to 32 GPUs. This is due at least partially to the faster intrinsic inference speed of REANN than the other two models with a comparable number of parameters. To isolate the comparison of parallelization scheme used in each architecture, we further evaluated their relative parallel efficiencies. Figure d shows that all models show decreases of parallel efficiency with increasing number of GPUs. Specifically, REANN–MPI retains the highest parallel efficiency across the tested range, with a modest decrease from 0.94 (2 GPUs) to 0.69 (32 GPUs). In contrast, SevenNet shows a more pronounced efficiency drop, falling to only 0.38 at 32 GPUs. Although MACE-p appears to achieve a parallel efficiency comparable to REANN–MPI for more than 8 GPUs, this is largely because its single-GPU baseline performance is considerably lower than that of MACE-s, which artificially inflates the apparent efficiency. In fact, the parallel version of MACE is significantly slower in its single-GPU evaluation, which makes the parallel overhead negligible and results in an almost constant relative efficiency as the number of GPUs increases. However, when referenced to the native serial baseline, the parallel efficiency of MACE drops substantially to 0.34 (32 GPUs)comparable to that of SevenNet but considerably lower than that of REANN–MPI. Notably, achieving high parallel efficiency is considerably more challenging for intrinsically faster models like REANN, as any overhead from the parallelization scheme would constitute a larger fraction of the total runtime, leading to a more pronounced reduction relative to the serial baseline. These results indicate that the proposed approach delivers superior scalability in multi-GPU parallelization.

Application to Reactive MD Simulations of Graphene Formation

Graphene formation through the detonation of gas-phase hydrocarbon mixtures with oxidizing agents at high temperature is known to be a high-yield route for producing graphene nanosheets and has attracted considerable interest for potential industrial applications. − Modeling this process and elucidating its mechanism are challenging due to the simultaneous involvement of chemical reactions, molecular clustering, and phase transformations. To address this problem, we have trained a highly parallelized universal REANN model and performed extensive MD simulations accelerated by the proposed scheme for the graphene formation from acetylene detonation with varying O₂/C₂H₂ ratios (see systems and training details section in Supporting Information). This is motivated by the recent work of Zhang et al., who studied this process using a local descriptor-based ANI-1xnr model trained with a comprehensive data set of C, H, O, N elements. This data set consists of 26,560 configurations generated by an active learning process starting from periodic cells containing randomly placed and oriented small molecules with a randomly selected composition of C, H, N and O. Energies and forces are computed by DFT using the BLYP functional, TZV2P basis set, GTH pseudopotentials, and D3 dispersion correction with zero damping (see the systems and training details section in Supporting Information). Based on the same data set, the accuracy of our REANN model converges roughly at T = 3, yielding considerably lower test root-mean-square errors (RMSEs) for both atomic energies (6.6 meV/atom) and forces (355.2 meV/Å), compared with those of the ANI-1xnr model (18.6 meV/atom and 448.4 meV/Å).

Figure a–d displays some key snapshots during this process modeled by REANN–MPI (T = 3). First, the collision-induced dissociation of C₂H₂ leads to fragments like hydrogen atoms, ethynyl and other free radicals (Figure b). This is followed by the formation of the intermediate long carbon chains (Figure c), some of which branch, and subsequently form graphene precursor composed of multiple fused carbon rings (Figure d). In addition, Figure e–l shows in more detail the evolution of multimembered rings up to 10 ns for eight different O₂/C₂H₂ ratios, using the REANN–MPI model with T = 3 including 27 K atoms. At low oxygen concentrations, six-membered (6-m) rings are the dominant product, followed by the subsequent 5-m and 7- m rings. As the O₂ concentration increases, the number of all carbon rings formed decreases, consistent with the experiment and previous simulations. − Interestingly, as seen in Figures S4 and S5, the production rates of graphene rings at different O₂/C₂H₂ ratios predicted by a local model (T = 1) are much higher than that by a deep MPNN model with T = 3, while the rates are relatively insensitive to the system size. Similar trend was also observed by Zhang et al. using the local ANI-1xnr model. This discrepancy is possibly related to the semilocal interaction of the long carbon chain intermediates, which is better described by semilocal MPNN-based models than local descriptor-based MLPs. , An explicit inclusion of long-range electrostatic interactions could further affect the dynamics of graphene ring formation. However, a detailed analysis of this effect lies beyond the scope of the present work and is therefore deferred to future studies.

8.

Evolution of graphene formation from C₂H₂ and O₂. (a–d) Representative snapshots illustrating key stages of graphene formation from C₂H₂ and O₂. Hydrogen, carbon, and oxygen atoms are shown in white, gray, and red, respectively. Green circles highlight hydrogen atoms, ethynyl species, and other radicals; red circles denote various long carbon chains; and purple circles identify the emerging graphene precursor. (e–l) Comparison of the formation of 3- to 10- membered (m) rings under different O₂/C₂H₂ ratios ranging from 0.00 to 1.33, obtained using the REANN–MPI model with T = 3.

The practical objective of this synthesis is the formation of extended graphene nanosheets, rather than isolated multimembered rings. Figure demonstrates the effect of simulation box on the maximum size of multiple fused carbon ringspotential precursors of graphene nanosheetsgenerated with different O₂/C₂H₂ ratios. When using a small periodic box with 1 K atoms, the maximum size of the generated graphene precursor decreases monotonically as the O₂/C₂H₂ ratio increases. When each length of the simulation box is doubled (8 K atoms) or tripled (27 K atoms), the maximum precursor size uniformly increases and peaks at the O₂/C₂H₂ ratio of 0.08 before decreasing at higher oxygen concentrations. In such cases, even with an O₂/C₂H₂ ratio of 0.86, one can still obtain the maximum graphene precursor size of ∼20 carbon atoms. This is in contrast with previous simulations with smaller simulation cells, e.g. ANI-1xnr (1000 atoms) and DFTB (up to 240 atoms), which identified very few graphene precursors beyond an O₂/C₂H₂ ratio of 0.50 and 0.20, respectively. Our results better agree with the experimental observation that graphene nanosheets form even at an O₂/C₂H₂ ratio as high as 0.8. These findings provide a mechanistic rationale for the experimental necessity of oxygen in this graphene synthesis.

9.

Size effect on graphene precursor formation. Variation in the maximum size (S _max) of the graphene precursor as a function of supercell size.

This case study indicates that a deeply semilocal MPNN model with a large simulation box enables a more reliable description of the graphene precursor formation process. While not infeasible, typical MPNN-based MD simulations in this case are less efficient and would require extensive computational resources. However, with the proposed parallel algorithm, simulations can be easily done with limited resources. For example, in the case of 27 K atoms with T = 3, the multiprocess REANN–MPI simulation with only 27 CPU cores distributing 1000 atoms in each core can greatly speed up the serial simulation by 24.8 times. This application demonstrates the practical benefit of our parallelization scheme for efficient large-scale MPNN simulations of real-world systems where local MLPs are less reliable.

Conclusions

In summary, we propose an efficient scheme for massive parallelization of MPNN potentials. A core concept of this scheme is that the data communication is minimized among local neighboring atoms in each MP layer, warranting its superior efficiency. Numerical tests on our REANN model demonstrate that this algorithm allows a linear scaling with the MP layer instead of a power exponential scaling as previously argued for such models. The resultant REANN–MPI model exhibits strong scaling with the number of CPU cores and GPUs, as well as excellent weak scaling with the number of atoms across various benchmark systems. This model enables fast parallel MD simulations involving over 100 million atoms, significantly extending the applicability of the MPNN potential. The successful application to graphene synthesis using REANN–MPI demonstrates the necessity of both multiple message-passing and large simulation size for an adequate description of the graphene precursor formation, highlighting the advantages of the proposed scalable MPNN parallelization scheme in simulating realistic systems. Our algorithm is not subject to a specific MP form and can be easily adapted for the parallelization of other MPNN potentials without changing their original structure. Further acceleration of the implementation of REANN–MPI can be achieved by enabling direct GPU-to-GPU data transfer. We anticipate that this scheme will empower MPNN models to efficiently simulate very large and complicated systems that were only accessible by strictly local models.

All data sets can be found in the corresponding original publications. Code availability: Code of REANN is available at https://github.com/bjiangch/REANN. The implementation of the MPI parallelization with REANN and the interface with LAMMPS is available at https://github.com/bjiangch/REANN-MPI.

The Supporting Information is available free of charge at https://pubs.acs.org/doi/10.1021/jacsau.6c00210.

• Mathematical details of the REANN model; a pseudocode for the REANN–MPI interface; systems and training details; software details; additional results on speed comparisons and MD simulations (PDF)

The authors declare no competing financial interest.