Year: 2021

Abstract:
Recent progress in Graph Neural Networks (GNNs) for modeling atomic simulations has the potential to revolutionize catalyst discovery, which is a key step in making progress towards the energy breakthroughs needed to combat climate change. However, the GNNs that have proven most effective for this task are memory intensive as they model higher-order interactions in the graphs such as those between triplets or quadruplets of atoms, making it challenging to scale these models. In this paper, we introduce Graph Parallelism, a method to distribute input graphs across multiple GPUs, enabling us to train very large GNNs with hundreds of millions or billions of parameters. We empirically evaluate our method by scaling up the number of parameters of the recently proposed DimeNet++ and GemNet models by over an order of magnitude. On the large-scale Open Catalyst 2020 (OC20) dataset, these graph-parallelized models lead to relative improvements of 1) 15% on the force MAE metric on the S2EF task and 2) 21% on the AFbT metric on the IS2RS task, establishing new state-of-the-art results.

heading: 1 INTRODUCTION
text: Graph Neural Networks (GNNs) (Gori et al., 2005; Zhou et al., 2020) have emerged as the standard architecture of choice for modeling atomic systems, with a wide range of applications from protein structure prediction to catalyst discovery and drug design (Schütt et al., 2017b; Gilmer et al., 2017; Jørgensen et al., 2018; Schütt et al., 2017a; 2018; Xie & Grossman, 2018; Qiao et al., 2020; Klicpera et al., 2020b). These models operate on graph-structured inputs, where nodes of the graph represent atoms, and edges represent bonds or atomic neighbors within a cutoff distance. Despite their widespread success and the availability of large molecular datasets, training massive GNNs (with up to billions of parameters) is an important but under-explored area. Success of similarly large models in computer vision, natural language processing, and speech recognition (Shoeybi et al., 2020; Huang et al., 2019; Brown et al., 2020; Zhai et al., 2021) suggests that scaling up GNNs could yield large performance gains.
Most previous approaches to scaling GNNs has focused on scaling small models (with up to a few million parameters) to massive graphs. These methods generally assume that we are working with a large, fixed graph, leading to the development of methods like neighborhood sampling (Jangda et al., 2021; Zheng et al., 2021) or efficient partitioning that depend on the graph structure being provided during training (Jia et al., 2020; Ma et al., 2019; Tripathy et al., 2020). These methods do not apply to atomic simulation datasets which contain millions of smaller graphs where it is necessary to consider the entire graph for prediction. In this paper, we focus on this complementary problem of scaling very large models to a dataset of many moderately-sized graphs (∼ 1k nodes, 500k edges).
Another limitation of existing methods is that they focus on scaling simple GNN architectures such as Graph Convolutional Networks (GCNs) that only represent lower-order interactions, such as representations for nodes and edges, that are then updated by passing messages between neighboring nodes. In practice, the most successful GNNs used for atomic systems also model higher-order interactions between atoms, such as the interactions between triplets or quadruplets of atoms (Klicpera et al., 2020a; 2021; Liu et al., 2021). These interactions are necessary to capture the geometry of the underlying system, which is critical in making accurate predictions. Scaling up such GNN architectures is challenging because even moderately-sized graphs can contain a large number of
higher-order interactions. For example, a single graph with 1000 nodes could contain several million combinations of triplets of atoms. In this paper, we introduce Graph Parallelism, an approach to scale up training of such GNNs with billions of parameters, by splitting up the graph across multiple GPUs.
We benchmark our approach by scaling up two recent GNN architectures – DimeNet++ (Klicpera et al., 2020a) and GemNet-T (Klicpera et al., 2021) – on the Open Catalyst (OC20) dataset (Chanussot et al., 2021). The OC20 dataset, aimed at discovering new catalyst materials for renewable energy storage, consists of 134M training examples spanning a wide range of adsorbates and catalyst materials. A GNN that can accurately predict per-atom forces and system energies on OC20 has the potential to replace the computationally expensive quantum chemistry calculations based on Density Functional Theory (DFT) that are currently the bottleneck in computational catalysis. Our large-scale, graph-parallelized models lead to relative improvements of 1) 15% for predicting forces on the force MAE metric (S2EF task), and 2) 21% on the AFbT metric for predicting relaxed structures (IS2RS task), establishing new state-of-the-art results on this dataset.

heading: 2 GRAPH PARALLELISM

heading: 2.1 EXTENDED GRAPH NETS
text: Battaglia et al. (2018) introduced a framework called Graph Network (GN) that provides a general abstraction for many popular Graph Neural Networks (GNNs) operating on edge and node representations of graphs. We build on their work and define the Extended Graph Network (EGN) framework to include GNNs that also operate on higher order terms like triplets or quadruplets of nodes.
In the Graph Network (GN) framework, a graph is defined as a 3-tuple G = (u, V, E), where u represents global attributes about the entire graph; V = {vi}i=1:Nv is the set of all nodes, with vi representing the attributes of node i; and E = {(ek, rk, sk)}k=1:Ne is the set of edges in the graph where ek represents the edge attributes for the edge from node sk to rk. A Graph Net is a neural network that contains a series of GN blocks that iteratively operate on the input graph, updating the various representations.
In our Extended Graph Network (EGN) framework, a graph is defined as a 4-tuple G = (u, V, E, T ), where u, V, and E are defined as in the Graph Network, and T = {(tm, em1 , em2 , . . .)} is the set of higher-order interaction terms involving edges indexed by m1,m2, . . ..
As a concrete example, consider an atomic system represented as a graph in this framework with the nodes representing the atoms and the edges representing atomic neighbors. The node attributes vi and edge attributes ek could represent the atom’s atomic numbers and distances between atoms respectively. The higher order interactions could represent triplets of atoms, i.e. pairs of neighboring edges with tm representing the bond angle, which is the angle between edges that share a common node. Finally, the global attribute u can represent the energy of the system. For clarity of exposition, we will limit our discussion to triplets in the rest of the paper, but higher order interactions can be handled in a similar manner. We denote these triplets as (tm, em1 , em2).
In the EGN framework, neural networks contain a series of EGN blocks that iteratively update the graph. Each EGN block consists of several update and aggregation functions that are applied to transform an input graph (u, V, E, T ) into an output graph (u′, V ′, E′, T ′). Each EGN block starts by updating the highest order interactions, which are then aggregated before updating the next highest order interaction. For instance, the triplet representations are first updated (using TripletUpdate function) and then aggregated (TripletAggr) at each edge. These aggregated representations are then used to update the edge representation (EdgeUpdate). Next, the edges going into a node are aggregated and used to update the node representations, and so on. This is illustrated in figure 1a.
Many popular GNNs can be implemented in this framework by making appropriate choices for the various update and aggregation functions.

heading: 2.2 GRAPH PARALLELISM FOR EXTENDED GRAPH NETS
text: Training large EGNs can be challenging even on moderately sized graphs because of the large memory footprint required in storing and updating the representation for each triplet, edge, and node. In
Triplet Update
Triplet Aggr.
Edge Update Node Update Global Update
Edge Aggr. Node Aggr.vi t′m
ek t′m
vi
e′k e ′ k t′m
v′i
e′k t′m
v′i
(a) Forward computation of an EGN block. First, each triplet representation is updated (TripletUpdate function), followed by an aggregation of these updated representations (TripletAggr). These aggregated values are then used to update the edge representations (EdgeUpdate), followed by an edge aggregation (EdgeAggr) and so on until the node and global representations are also updated.
All Reduce
Processor 0 Processor 1 Processor 2 Processor 3
All Reduce All Reduce
Triplet Update and Local Triplet Aggregation
Global Triplet Aggregation
Edge Update and Local Edge Aggregation
Global Edge Aggregation Node Update and Local
Node Aggregation
Global Node Aggregation
Global Update
(b) Distributed computation of an EGN block. The graph is split up among the different GPUs such that processing unit p contains its subset of triplets T (p) in memory, along with all edges E and nodes V . The triplet attributes are updated in parallel, followed by a triplet aggregation to locally update the edge attributes. Next, an allreduce operation is performed to update all edge attributes globally. We continue this process to update the node and global attributes.
many applications, the number of edges is one or two orders of magnitude larger than the number of nodes, while the number of triplets is one or two orders of magnitude larger than the number of edges. Thus, storing and updating the triplet representations is often the bottleneck in terms of GPU memory and compute. Many recent methods such as (Klicpera et al., 2020a; 2021) overcome this problem by using a very low-dimensional representation for the triplets. However, we found this to significantly reduce model capacity leading to underfitting for some applications with large training datasets. It is necessary to overcome these memory limitations for better generalization.
One way to avoid the memory limits is to distribute the computation across multiple GPUs. For a graph neural network, a natural choice to distribute the computation is by splitting the graph. The update functions in an EGN are easy to apply in parallel since they are applied independently for each triplet, edge, or node. It is substantially more challenging to apply the aggregations in parallel. To simplify parallel aggregation, we restrict ourselves to aggregation functions that are commutative and associative. This is not limiting in practice since most popular GNN architectures use sum or mean aggregations which can be implemented in this manner.
We will now describe the implementation of the distributed EGN block. Suppose we have access to P processing units that we wish to split the graph computation over. Each unit p is responsible for computing the updates to a subset of triplets, edges and nodes, that we denote by T (p), E(p), and V (p), respectively. At the beginning of computation, we split the graph so that the processing
unit p contains its subset of triplets T (p) in memory, along with the entire set of edges E and nodes V . During the forward pass, we first update each set of triplets T (p) in parallel to obtain T ′(p), followed by a local triplet aggregation. Next, an all-reduce operation is performed on these locally aggregated triplet representations to obtain globally aggregated triplet representations. These globally aggregated representations are then used to update edges in parallel, and the process is repeated for nodes and ultimately for the global graph level representation. Figure 1b shows this.
In this framework, the highest order interaction attributes are never communicated across processors. Therefore, the communication time is bound by the number of lower order node interactions. In our example, the triplet representations are not communicated, while the edge and node representations are communicated once per EGN block, making the total communication cost equal to O(NvDv + NeDe), where Dv and De are the dimensions of node and edge representations. This makes it possible to work with a large number of triplets and large triplet representations. In section 3, we show concrete instantiations of this framework for two contemporary GNNs.

heading: 3 GRAPH PARALLELISM FOR ATOMIC SIMULATIONS
text: In this section, we present two concrete examples of using GNNs for the problem of predicting the energy and the forces for an atomic system, modeled as a graph whose nodes represent the atoms and whose edges represent the atoms’ neighbors. The GNN takes such a graph as input and predicts the energy of the entire system as well as a 3D force vector on each atom. Two paradigms have been proposed in the literature, that we call energy-centric and force-centric approaches. These approaches differ in how they estimate forces and whether they are energy conserving, which is an important physical property. We begin by describing the components shared by both approaches, followed by one recent model in each paradigm.

heading: 3.1 INPUTS AND OUTPUTS
text: The inputs to the network are 3D positions xi ∈ R3 and atomic number zi for each atom i ∈ {1 . . . n}. The outputs are the per-atom forces fi ∈ R3 and the energy E ∈ R of the entire structure. The distance between atoms i and j is denoted as dij = ||xi − xj ||. If edges (i, j) and (j, k) exist, then (i, j, k) defines a triplet in the graph and we denote the angle between the edges as αkj,ji.
The input graph is constructed with each atom t (target) as a node and the edges representing the atom’s neighbors s ∈ Nt where Nt contains all atoms s (source) within a distance δ, which is treated as a hyperparameter. Each edge has a corresponding message mst that passes information from source atom s to target atom t.

heading: 3.2 ESTIMATING FORCES AND ENERGY
text: As previously stated, there are two paradigms for estimating the energy and forces for an atomic system: energy-centric models and force-centric models. In energy-centric models, the model first computes the energy E as a function of the inputs by applying a forward pass of the GNN: E = GNN(x, z), where x, and z represent the atomic positions and atomic numbers respectively. The forces are then computed as the gradient of the energy with respect to atomic positions by using backpropagation: f = ∇xE. In force-centric models, the energy and forces are both computed directly during the forward propagation: E,F = GNN(x, z) where F represents the matrix of all atomic forces. Force-centric models tend to be more efficient in terms of computation time and memory usage compared to the energy-centric models. However, energy-centric models guarantee that the forces are energy conserving which is an important physical property satisfied by atomic systems. In this work, we demonstrate the benefits of scaling up GNNs in both paradigms.

heading: 3.3 ENERGY-CENTRIC MODEL: DIMENET++
text: Dimenet++ (Klicpera et al., 2020a) is a recently proposed energy-centric model for predicting properties from atomic systems. In this model, the edges are represented by a feature representation mji, which are iteratively updated using both directional information (via bond angles) as well as
the interatomic distances. The edges are initially represented using a radial basis function (RBF) representation e(ji)RBF of their lengths dji. The triplets are represented using a spherical basis function (SBF) expansion a(kj,ji)SBF of the distances dkj as well as bond angles α(kj,ji). In each block, the messages are updated as:
m′ji = fupdate(mji, ∑
k∈Nj\{i} fint(mkj , e
(ji) RBF ,a (kj,ji) SBF )) (1)
where the interaction function fint corresponds to the TripletUpdate function, the summation corresponds to the TripletAggr function, and the update function fupdate to the EdgeUpdate function respectively. The interaction function consists of a hadamard product of the embeddings, followed by a multi-layer perceptron. To speed up these computations, the embeddings are projected down to a smaller dimension before computing the interactions, and later projected back up to the original dimension.
The updated messages are then input to an output block which are summed up per atom i to obtain a per-atom representation: hi = ∑ j m ′ ji (EdgeAggr function). These are then transformed by another MLP to obtain the node representation v′i (NodeUpdate function).
Thus, Dimenet++ is a special case of an EGN, and we can closely follow the recipe from section 2.2 to parallelize it.

heading: 3.4 FORCE-CENTRIC MODEL: GEMNET
text: Gemnet (Klicpera et al., 2021) extends Dimenet++ in a number of ways, including the addition of quadruplets as well as a force-centric version. Here, we focus only on the force-centric Gemnet-T model since it was recently shown to obtain state-of-the-art results on the OC20 dataset1. Gemnet-T largely follows the structure of the Dimenet++ model, but includes some modifications.
First, Gemnet-T uses a bilinear layer instead of the Hadamard product for the interaction function. This is made efficient by optimally choosing the order of operations in the bilinear function to minimize computation. Second, Gemnet-T maintains an explicit embedding for each atom that is first updated by aggregating the directional embeddings involving that atom, similar to Dimenet++. Next, the updated atom embedding is used to update each of the edge embeddings. This creates a second edge update function, EdgeUpdate′, that is run after the node embeddings are updated. Third, Gemnet makes use of symmetric message passing, that is the messages mji and mij that are on the same edge, but in different directions, are coupled. In a parallel implementation, this step requires an additional all-reduce step since messages mji and mij could be on different processors.
Thus, Gemnet-T largely follows the EGN framework, with a few minor deviations from the standard formulation. The distributed EGN implementation described in section 2.2 can be used for the Gemnet-T as well, but with additional communication steps to account for the second edge update function and symmetric message passing.

heading: 4 EXPERIMENTS
text: In this section, we present the results of our scaling experiments on the OpenCatalyst 2020 (OC20) dataset (Chanussot et al., 2021). The OC20 dataset contains over 130 million atomic structures used to train models for predicting forces and energies during structure relaxations. We report results for three tasks: 1) Structure to Energy and Forces (S2EF) that involves predicting energy and forces for a given structure; 2) Initial Structure to Relaxed Energy (IS2RE) that involves predicting the relaxed energy for a given initial structure; and 3) Initial Structure to Relaxed Structure (IS2RS) which involves performing a structure relaxation using the predicted forces. Dimenet++ and Gemnet-T are the current state-of-the-art energy-centric and force-centric models respectively.
1https://opencatalystproject.org/leaderboard.html

heading: 4.1 EXPERIMENTAL SETUP
text: Dimenet++-XL. Our Dimenet++ model consists of B = 4 interaction blocks, with a hidden dimension of H = 2048, an output block dimension of D = 1536, and intermediate triplet dimension of T = 256. This model has about 240M parameters, which is over 20× larger than the Dimenet++large model used in (Chanussot et al., 2021). We call this model Dimenet++-XL. The model was trained with the AdamW optimizer (Kingma & Ba, 2015; Loshchilov & Hutter, 2019) starting with an initial learning rate of 10−4, that was multiplied by 0.8 whenever the validation error plateaus. The model was trained with an effective batch size of 128 on 256 Volta 32GB GPUs with a combination of data parallel and graph parallel training: each graph was split over 4 GPUs with data parallel training across groups of 4 GPUs.
Gemnet-XL. Our Gemnet model consists of B = 6 interaction blocks, with an edge embedding size of E = 1536, triplet embedding size of T = 384 and embedding dimension of the bilinear layer of B = 192. We found that it was beneficial to use a small atom embedding size of A = 128, much smaller than the previous SOTA Gemnet model. 2. This model has roughly 300M parameters, which is about 10× larger than the previous SOTA model. However, since we reduced the atom embedding dimension and increased the edge and triplet dimensions, the total amount of compute and memory usage is significantly larger. We call this model Gemnet-XL. We followed the same training procedure as with Dimenet++, except for a starting learning rate of 2× 10−4.

heading: 4.2 STRUCTURE TO ENERGY AND FORCES (S2EF)
text: The Structure to Energy and Forces task takes an atomic structure as input and predicts the energy of the entire structure and the per-atom forces. The S2EF task has four metrics: the energy and force Mean Absolute Error (MAE), the Force Cosine Similarity, and the Energy and Forces within a Threshold (EFwT). The EFwT metric indicates the percentage of energy and force predictions that are below a preset threshold.
Table 1 compares the top 4 models on the Open Catalyst Project leaderboard3 with the Gemnet-XL model. Gement-XL obtains a roughly 16% lower force MAE and an 8% lower energy MAE relative to the previous state-of-the-art. Further, Gemnet-XL improves the EFwT metric more than 50% relative to the previous best, although the value is still very small. These results indicate that model scaling is beneficial for the S2EF task.

heading: 4.3 INITIAL STRUCTURE TO RELAXED STRUCTURE (IS2RS)
text: The Initial Structure to Relaxed Structure (IS2RS) task involves taking an initial atomic structure and predicting the atomic structure that minimizes the energy. This is performed by iteratively predicting the forces on each atom and then using these forces to update the atomic positions. This process is repeated until convergence or 200 iterations. There are three metrics for this task: Average Distance within Threshold (ADwT), that measures the fraction of final atomic positions within a distance threshold of the ground truth; Forces below Threshold (FbT), which measures whether a true energy minimum was found (i.e., forces are smaller than a preset threshold); and, the Average Forces below Threshold (AFbt), which averages the FbT over several thresholds.
2https://discuss.opencatalystproject.org/t/new-gemnet-dt-code-results-model-weights/102 3https://opencatalystproject.org/leaderboard.html
Table 2 shows the results on the IS2RS task, comparing our models with the top few models on the OpenCatalyst leaderboard. Both the Dimenet++-XL and Gemnet-XL models outperform all existing models. The Dimenet++-XL model obtains a relative improvement of 53% on the AFbT metric, and more than triples the FbT metric compared to the Dimenet-large model. The Gemnet-XL model obtains similar improvements compared to the smaller Gemnet-T model. These results underscore the importance of model scaling for this task.

heading: 4.4 INITIAL STRUCTURE TO RELAXED ENERGY (IS2RE)
text: The Initial Structure to Relaxed Energy (IS2RE) task takes an initial atomic structure and attempts to predict the energy of the structure after it has been relaxed. Two approaches can be taken to address this problem, the direct and the relaxation approaches (Chanussot et al., 2021). In the direct approach, we treat this task as a regression problem, and train a model to directly estimate the relaxed energy for a given atomic structure. The relaxation approach first estimates the relaxed structure (IS2RS task) after which the energy is estimated using the relaxed structure as input. Relaxation approaches typically outperform the direct approaches, though they are generally two orders of magnitude slower during inference time due to the need to estimate the relaxed structure.
Table 3 compares our Gemnet-XL model to the top three models from the Open Catalyst Project leaderboard. There are two metrics for the IS2RE task: energy Mean Absolute Error (MAE) and the Energy within Threshold (EwT) which measures the percentage of time the predicted energy is within a threshold of the true energy. Table 3 shows that the Gemnet-XL model obtains a roughly 8% lower energy MAE and a 12% higher EwT compared to the previous best, which is the smaller Gemnet-T model, demonstrating the benefits of scaling up IS2RE models.
Since direct IS2RE models are very fast at inference time, there are applications where they are more useful than relaxation based approaches. It is possible to convert a trained S2EF model into a direct IS2RE model by fine-tuning it on the IS2RE training data. We finetune our Gemnet-XL model in this manner for 5 epochs, starting with an initial learning rate of 3 × 10−5 that is exponentially decayed by multiplying with 0.95 at the end of each epoch. We call the resuting model Gemnet-XLFT, which obtains a roughly 9.5% better energy MAE and 24% better EwT compared to the current state-of-the-art direct approach (Table 3).

heading: Edge Dim Atom Dim Triplet Dim Bilinear Dim Num Params Num GPUs

heading: Weak Scaling (model size)

heading: 4.5 SCALING ANALYSIS
text: Weak scaling studies the effect on the throughput when computation is scaled proportional to the number of processors. We study weak scaling efficiency in terms of scaling batch size as well as scaling model size, both with graph parallelism. For these experiments, we use 4 different Gemnet-T models, ranging from 125.4 million parameters to 1.12 billion parameters (see Table 4).
For weak scaling efficiency by batch size, we train the smallest Gemnet-T model from Table 4 and scale up batch size linearly with number of GPUs. Figure 2a shows that using graph parallelism to scale to 8x the batch size on 8 GPUs achieves ∼46% of linear scaling. For reference, if the throughput (time per batch) with 8x the batch size on 8 GPUs was the same as 1x the batch size on 1 GPU, efficiency would be 100%. In this case, since the model fits on a single GPU, one can alternatively use data parallelism and we expect it to have better scaling efficiency with batch size because of significantly fewer quantities being communicated across GPUs than graph parallelism.
For weak scaling efficiency by model size, we train increasingly larger models on multiple GPUs that do not fit on a single GPU. Figure 2b shows that we are able to obtain a scaling efficiency of roughly 50% with 8 GPUs for our largest model with 1.12B parameters. Here, data parallelism is not even applicable since this model does not fit on 1 GPU. Further, we see that going from 500M parameters (on 4 GPUs) to 1.12B parameters (on 8 GPUs), scaling efficiency only decreases from ∼54% to ∼50%, suggesting that graph parallelism is promising in this large model size regime. Finally, Figure 2c shows that combining graph and data parallelism follows a similar linear scaling behavior as graph parallel alone, suggesting that the two approaches are complementary.

heading: 5 RELATED WORK
text: GNNs for simluating atomic systems Many GNN based approaches have been proposed for the task of estimating atomic properties such as (Schütt et al., 2017b; Gilmer et al., 2017; Jørgensen et al., 2018; Schütt et al., 2017a; 2018; Xie & Grossman, 2018; Qiao et al., 2020; Klicpera et al., 2020b), where the atoms are represented by nodes and neighboring atoms are connected by edges. An early approach for force estimation was the SchNet model Schütt et al. (2017a), which computed forces using only the distance between atoms without the use of angular information. SchNet proposed the use of differentiable edge filters which enabled constructing energy-conserving models by estimating forces as the gradient of the energy. Subsequent work (Klicpera et al., 2020a;b; 2021;
Liu et al., 2021) has extended on the SchNet model by adding bond angles and dihedral angles, which has resulted in improved performance. These models make use of higher order interactions among nodes which make them highly compute and memory intensive. An alternate approach for estimating forces is to directly regress the forces as an output of the network. While this approach does not enforce energy conservation or rotational equivariance, recent work (Hu et al., 2021) has shown that such models can still accurately predict forces.
Distributed GNN Training. Research on distributed GNN training has focused on the regime of training small models on a single, large graph, for instance by sampling local neighborhoods around nodes to create mini-batches (Jangda et al., 2021; Zhang et al., 2020; Zhu et al., 2019). While these approaches can scale to very large graphs, they are not suitable for the task of modeling atomic systems where it is important to consider the entire graph for predicting the energy and forces.
An alternate line of work, that is more similar to ours, keeps the entire graph in memory by efficiently partitioning the graph among multiple nodes (Jia et al., 2020; Ma et al., 2019; Tripathy et al., 2020). These methods still operate in a single graph regime that is partitioned ahead of time. Thus the focus of that work is on finding efficient partitions of the graph, which is not applicable to our problem since we operate on millions of graphs. Further, these works do not train very large GNNs, or GNNs that operate on higher-order interactions (e.g. triplets), that are important for atomic systems.
Model Parallelism methods focus on training large models that do not fit entirely on one GPU (even with a batch size of 1). GPipe (Huang et al., 2019) splits different sequences of layers into different processors, and splits each training mini-batch into micro-batches to avoid idle time per GPU. Megatron-LM (Shoeybi et al., 2020) splits the model breadth-wise, where Transformer layer weights are partitioned across multiple GPUs to distribute computation. We see model and graph parallelism as complementary approaches, and can both be combined to train even larger models.

heading: 6 DISCUSSION
text: In this paper, we presented graph parallelism, a new method for training large GNNs for modeling atomic simulations for which modeling higher-order interactions is critical. We showed that training larger Dimenet++ and Gemnet-T models can yield significant improvements on each task of the OC20 dataset. And although we demonstrated graph parallelism for just two GNNs, it is more broadly applicable to a host of message passing GNNs, including equivariant networks, that can be cast in the GN / EGN framework (Sec. 2.1) by an appropriate choice of update and aggregate functions. To give a couple more examples, the model proposed in Jing et al. (2021) uses geometric representations for nodes and edges that are updated through message passing. The update functions involve the use of Geometric Vector Perceptrons (GVPs) that preserve equivariance. This model can be cast in the GN framework by using an update function that includes GVPs within it. Schütt et al. (2021) present another equivariant GNN model called PAINN, that is composed of a series of Message and Update blocks. By viewing the computation within the message block as the EdgeUpdate function, the message passing step as the EdgeAggregate function and the Update block as the NodeUpdate function, we can see that this model can also be implemented as a GN.
Further, it is possible to combine graph parallelism with model parallel methods such as GPipe (Huang et al., 2019) to train even larger models, which could yield further improvements. Although we have not yet explored this, for force-centric GNNs, it should be possible to use graph parallel for ‘breadth-wise’ scaling i.e. to split higher-order computations (e.g. triplets) across GPUs, and GPipe for ‘depth-wise’ scaling i.e. to scale to larger number of message-passing blocks, sequentially split across GPUs. For energy-centric GNNs e.g. DimeNet++, this combination is less obvious since these models require an additional backward pass through the network to compute forces as gradients of energy with respect to atomic positions. Energy-centric GNNs are common for atomic systems because they enforce the physical constraint of energy conservation. As we demonstrate with our DimeNet++-XL experiments, graph parallelism is applicable to energy-centric GNNs.
We see scaling to large model sizes as a necessary (but not sufficient) step to effectively model atomic simulations from large, diverse datasets of adsorbates and catalysts. Further progress may require marrying large scale with ways to better capture 3D geometry and physical priors.

heading: A APPENDIX

heading: A.1 ADDITIONAL RESULTS
text: Tables 5, 6 and 7 show results from each of the four test sets for the S2EF, IS2RE and IS2RS tasks respectively. Dimenet++-XL and Gemnet-XL achieve the best results for each test set along each metric.

heading: OOD Both

heading: OOD Cat

heading: OOD Ads

heading: OOD Both

heading: OOD Cat

heading: OOD Ads

heading: OOD Both

heading: OOD Cat

heading: OOD Ads

heading: Model Training IS2RS TestDataset AFbT↑ ADwT↑ FbT↑

Review summaries:
This paper proposes a method to train large scale graph neural networks containing up to billions of parameters, called Graph Parallelism. The method is used to train large scale versions of the DimeNet++ and GemNet models containing 10-20x more parameters that the vanilla versions. These large GNN models are evaluated on a set of tasks from the Open Catalyst 2020 (OC20) benchmark and show improved performance compared to the smaller baselines. 

This paper proposes a distributed training method for large graph neural networks (GNN) up to billion parameters. The method first distributes the triplet update operations to multiple GPUs and aggregate the updated vectors by global synchronization. It then distributes the edge update operations to GPUs and aggregate the edge vectors (another global synchronization). Finally, it applies node update in parallel and global node aggregation at the end. The method is applied to GemNet and DimeNet up to 1.12 billion parameters and achieved state-of-the-art results on OC20 benchmark.

The authors present an approach to train graph neural networks with many parameters across multiple GPUs. The approach is specifically demonstrated for graphs representing molecular structure and two graph neural networks of similar flavor that have previously been presented to learn from such structures. Higher-order interaction terms make these networks very compute intensive. Leveraging their parallel approach to scale up the existing GNNs, the authors demonstrate compelling results on the OpenCatalyst benchmark. 

The paper is based on the idea that most of the computation costs to extended graphs come from triplets. The authors proposed a way to parallelize the computation of triplets in a distributed way. The author also discussed two models that fit into this framework and showed their increase in performance due to the larger parameter count enabled by their parallel framework.