Towards Training Billion Parameter Graph Neural Networks for Atomic Simulations

1 INTRODUCTION . 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 . 2 GRAPH PARALLELISM . 2.1 EXTENDED GRAPH NETS . 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 . 2.2 GRAPH PARALLELISM FOR EXTENDED GRAPH NETS . 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 .

Summary: 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.