T OWARDS T RAINING B ILLION P ARAMETER G RAPH N EURAL N ETWORKS FOR A TOMIC S IMULATIONS
A BSTRACT
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 for the S2EF task and 2) 21% on the
AFbT metric for the IS2RS task, establishing new state-of-the-art results.
I NTRODUCTION
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; Zitnick et al., 2020; Schütt et al., 2017a; Xie & Grossman, 2018).
These models operate on graph-structured inputs, where nodes of the graph represent atoms, and edges represent bonds or atomic neighbors.
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 signiﬁcant performance gains.
Most previous approaches to scaling GNNs have 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, ﬁxed graph, leading to the development of methods like neighborhood sampling (Jangda et al., 2021; Zheng et al., 2021)
(Jia et al., 2020; Ma et al., 2019; Tripathy et al., 2020).
These methods do not apply to atomic simulation datasets that contain millions of smaller graphs where it is necessary to consider the entire graph for prediction.
Our focus is on the complementary problem of scaling to very large models for a dataset of many moderately-sized graphs ( ⇠ 1 k nodes, 500 k 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 i.e. , 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., 2022).
These interactions are necessary to capture the geometry of the underlying system, critical in making accurate predictions.
Scaling 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 1 k nodes could contain several million triplets of atoms.
In this paper, we introduce Graph Parallelism , an approach to scale up such GNNs with higher-order interactions to billions of parameters, by splitting up the input 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 134 M training examples spanning a wide range of adsorbates and catalyst materials.
A GNN that can accurately predict per-atom forces and system energies on OC 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.
G RAPH P ARALLELISM 2. E XTENDED G RAPH N ETS
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 deﬁne 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 deﬁned as a 3 -tuple G = ( u , V, E ) , where u represents global attributes about the entire graph; V = { v i } i =1: N v is the set of all nodes, with v i representing the attributes of node i ; and E = { ( e k , r k , s k ) } k =1: N e is the set of all edges where e k represents the attributes for the edge from node s k to r k .
A GNN then 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 deﬁned as a 4 -tuple G = ( u , V, E, T ) , where u , V, and E are deﬁned as in the Graph Network, and T = { ( t m , e m 1 , e m 2 , . . . )
} is the set of higher-order interaction terms involving edges indexed by m 1 , m 2 , . . . .
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 v i and edge attributes e k 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 t m 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 ( t m , e m 1 , e m 2 ) .
In the EGN framework, GNNs then 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 0 , V 0 , E 0 , T 0 ) .
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 ﬁrst 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 ﬁgure 1a.
Many GNNs can be cast in the EGN framework using appropriate update and aggregation functions.
2. G RAPH P ARALLELISM FOR E XTENDED G RAPH N ETS
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 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
[Table]
(a) Forward computation of an EGN block
First, each triplet representation is up- dated ( TripletUpdate function), followed by aggregation of these updated representations ( TripletAggr ).
Next, these aggregated values are used to update edge representations ( EdgeUpdate ), followed by edge aggregation ( EdgeAggr ), and ﬁnally node and global updates.
All Reduce
[Figure]
All Reduce All Reduce
[Figure]
[Table]
[Figure]
[Table]
problem by using a very low-dimensional representation for the triplets.
However, we found this to signiﬁcantly reduce model capacity leading to underﬁtting 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 ﬁrst update each set of triplets T ( p ) in parallel to obtain
T 0 ( 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 ( N v D v +
N e D e ) , where D v and D e 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.
G RAPH P ARALLELISM FOR A TOMIC S IMULATIONS
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.
3. I NPUTS AND O UTPUTS
The inputs to the network are 3D positions x
i 2 R 3 and atomic number z i for each atom i{ 1 . . .
n } .
The outputs are the per-atom forces f i 2 R 3 and the energy E 2 R of the entire structure.
The distance between atoms i and j is denoted as d ij =
||
x i  x j || .
If edges ( i, j ) and ( j, k ) exist, then ( i, j, k ) deﬁnes 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 2 N t where N t contains all atoms s ( source ) within a distance  , which is treated as a hyperparameter.
Each edge has a corresponding message m st that passes information from source atom s to target atom t .
3. E STIMATING F ORCES AND E NERGY
As previously stated, there are two paradigms for estimating the energy and forces for an atomic system: energy-centric and force-centric .
In energy-centric models, the model ﬁrst computes the energy E 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 negative gradient of the energy with respect to atomic positions by using backpropagation: f = r x E .
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 efﬁcient 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 satisﬁed by atomic systems.
In this work, we demonstrate the beneﬁts of scaling up GNNs in both paradigms.
3. E NERGY -C ENTRIC MODEL : D IME N ET ++
DimeNet++ (Klicpera et al., 2020a) is a recently proposed energy-centric model for atomic systems.
In this model, the edges are represented by a feature representation m ji , 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 d ji .
The triplets are represented using a spherical basis function (SBF) expansion a ( kj,ji ) SBF of the distances d kj as well as bond angles ↵ ( kj,ji ) .
In each block, the messages are updated as:
[Equation]
[Equation] [Equation] where the interaction function f int corresponds to the TripletUpdate function, the summation corresponds to the TripletAggr function, and the update function f update 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 interactions, and later projected back up.
The updated messages are then fed as input to an output block that sums them up per atom i to obtain a per-atom representation: h i =
P j m 0 ji ( EdgeAggr function).
These are then transformed by another MLP to obtain the node representation v {+
+}
i ( NodeUpdate function).
Thus, DimeNet++ is a case of EGN, and we closely follow the recipe from Sec.
2.2 to parallelize it.
3. F ORCE -C ENTRIC MODEL : G EM N ET
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 dataset 1 .
GemNet-T largely follows the structure of the DimeNet++ model, but includes some modiﬁcations.
First, GemNet-T uses a bilinear layer instead of the Hadamard product for the interaction function.
This is made efﬁcient 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 ﬁrst 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 0 , that is run after the node embeddings are updated.
Third, GemNet makes use of symmetric message passing, that is the messages m ji and m ij 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 m ji and m ij 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.
E XPERIMENTS
In this section, we present the results of our scaling experiments on the Open Catalyst 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.
4. E XPERIMENTAL S ETUP
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 240 M 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. [Table] [Table]
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 beneﬁcial 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 signiﬁcantly 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 .
4. S TRUCTURE TO E NERGY AND F ORCES (S2EF)
The Structure to Energy and Forces task takes an atomic structure as input and predicts the energy of the entire structure and 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).
EFwT indicates the percentage of energy and force predictions below a preset threshold.
Table 1 compares the top models on the Open Catalyst Project leaderboard 1 with the GemNetXL model.
GemNet-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 beneﬁcial for the S2EF task.
4. I NITIAL S TRUCTURE TO R ELAXED S TRUCTURE (IS2RS)
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 ﬁnal 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.
Table 2 shows the results on the IS2RS task, comparing our models with the top few models on the Open Catalyst leaderboard.
Both the DimeNet++-XL and GemNet-XL models outperform all [Table] [Table] 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 GemNetXL model obtains similar improvements compared to the smaller GemNet-T model.
These results underscore the importance of model scaling for this task.
4. I NITIAL S TRUCTURE TO R ELAXED E NERGY (IS2RE)
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 ﬁrst 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 beneﬁts 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 ﬁne-tuning it on the IS2RE training data.
We ﬁnetune 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 .
at the end of each epoch.
The resuting model – GemNet-XL-FT – obtains a relative improvement of ⇠ 2% on energy MAE compared to 3D-Graphormer (Ying et al., 2021), the current state-of-the-art direct approach (Table 3).
4. S CALING A NALYSIS
Weak scaling studies the effect on the throughput when computation is scaled proportional to the number of processors. We study weak scaling efﬁciency in terms of scaling the model size propor-
[Figure]
Dashed lines show the same with 32-way data parallel training. We do not show the results of training the largest models with pipeline parallelism as those runs ran out of GPU memory. tional to the number of GPUs.
For these experiments, we use 4 different GemNet-T models, ranging from 120M parameters to nearly 1B parameters (see Table 4).
We train increasingly larger models on multiple GPUs that do not ﬁt on a single GPU.
Figure 2a shows that we are able to obtain a scaling efﬁciency of roughly 79% with 8 GPUs for our largest model with nearly a billion parameters.
This shows that our graph parallel training can be scaled up to billion parameter GemNet-T models while obtaining reasonably good scaling performance. Figure 2a further shows the scaling efﬁciency for each model combining graph parallelism with 32-way data parallelism. Pure data parallel training with 32 GPUs obtains a scaling efﬁciency of 75% for the smallest model, showing the effect of network communication and load imbalance between the GPUs.
Combining graph and data parallel training with 256 GPUs only reduces the scaling efﬁciency to 47% for the largest model compared to the 1 GPU case, suggesting the graph parallelism is promising for training extremely large models on hundreds of GPUs. Finally, ﬁgure 2b shows the raw performance of running these models on V100 GPUs in terms of
TeraFLOPs per second as a function of the number of GPUs. On a single GPU, the 120M parameter GemNet-T sustains 32 TeraFLOPs or roughly 25% of the theoretical peak FLOPS for a single GPU, and is thus a strong baseline. With 256 GPUs, the largest model sustains 3 .
5 PetaFLOPs. Figure 2b compares graph and pipeline parallelism (Huang et al., 2019) showing that graph parallelism outperforms pipeline parallelism for the models that we consider. Since each graph in the training data contains different numbers of nodes, edges and triplets, load balancing across GPUs is difﬁcult for pipeline parallelism. Graph parallelism is able to overcome this problem since the nodes, edges and triplets of a given batch are always distributed evenly across the GPUs, helping it outperform pipeline parallelism.
It is possible, however, that pipeline parallelism might outperform graph parallelism for very deep GNNs since inter-GPU communication overhead for pipeline parallelism is independent of the number of blocks.
Figure 2b also shows results with graph and pipeline parallelism combined, indicating that these methods are complementary to each other.
R ELATED W ORK
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 ﬁlters 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., 2022) 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 efﬁciently 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 ﬁnding efﬁcient 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 ﬁt 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 that can be combined to train even larger models.
D ISCUSSION
We presented Graph Parallelism, a new method for training large GNNs for modeling atomic simulations, where modeling higher-order interactions between atoms (triplets / quadruplets) is critical.
We showed that training larger DimeNet++ and GemNet-T models can yield signiﬁcant improvements on the OC20 dataset.
Although we demonstrated graph parallelism for just two GNNs, it is 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 appropriately picking update and aggregate functions.
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, as brieﬂy explored in Sec 4.5.
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 sufﬁcient) 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.
The carbon emissions from training large deep learning models are non-trivial and the work we have presented here is no exception (Strubell et al., 2019; Schwartz et al., 2019).
We estimate that training our GemNet-XL model with Tesla v100 32 GB GPUs on cloud resources in the US ranges from- 8052 kg of CO 2 eq.
(Lacoste et al., 2019).
In the worst case, the emissions are roughly equivalent to 16 round trip ﬂights from Los Angeles to New York.
Assuming that the training time is ﬁxed, emissions largely depend on the carbon intensity of the energy generation in a given region and the percentage of emissions offset with other investments by the resource provider. When choosing compute resources we recommend evaluating the stated carbon offset commitments and if possible, consider running experiments in regions that have more sustainable energy generation. The compute resources we utilized for this paper were committed to be 100% offset. A A PPENDIX A. A DDITIONAL RESULTS
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.
[Table]
[Table] [Table]