EXACT: Scalable Graph Neural Networks Training via Extreme Activation Compression

1 INTRODUCTION . Despite Graph Neural Networks ( GNNs ) have achieved great success across different graph-related tasks , training GNNs on large graphs is a long-standing challenge due to its extensive memory requirement ( Kipf & Welling , 2017 ; Zhang & Chen , 2018 ; Cai et al. , 2021b ) . The extensive memory consumption of GNN stems from its recursive neighborhood aggregation scheme , where each node aggregates embeddings of its neighbors to update its new embedding at each layer . Thus , training an L-layer GNN requires storing all L layers ’ intermediate node embeddings in GPU memory for computing the gradients , and this typically adds several times more memory than holding the node feature matrix ( see Algorithm 1 for a detailed analysis ) . Hence , storing these node embeddings is the major memory bottleneck for training GNNs on large graphs . Most of the existing works towards this problem can be roughly divided into two categories . First , some works propose to train GNNs with sampled subgraphs instead of the whole graph at each step . In this way , only node embeddings that are present in the current subgraph will be retained in memory ( Chiang et al. , 2019 ; Hamilton et al. , 2017 ; Zeng et al. , 2020 ; Zou et al. , 2019 ; Chen et al. , 2018 ; Huang et al. , 2018 ) . Second , another line of work tries to decouple the neighborhood aggregation from prediction , either as a preprocessing step ( Wu et al. , 2019 ; Klicpera et al. , 2018 ; Yu et al. , 2020 ) or post-processing step ( Huang et al. , 2020 ) , where the model is simplified as the Multi-Layer Perceptron ( MLP ) that can be trained with mini-batch data . In parallel , another orthogonal direction is to store only the compressed node embeddings ( i.e. , activations ) in memory for computing the gradients . Recent works propose to quantize the activations in lower numerical precision ( e.g. , using 8-bit integer ) during the forward pass ( Chakrabarti & Moseley , 2019 ; Fu et al. , 2020 ; Chen et al. , 2021a ; Evans & Aamodt , 2021 ) . This framework successfully trims down the memory requirement for training Convolutional Neural Networks ( CNNs ) by a large margin , at the cost of additional time overhead and loss of accuracy . Ideally , the real-world usage requires the training method should achieve a balanced trade-off among the following three aspects : 1 . Space . It should enable to train GNNs on large graphs using off-the-shelf hardwares , such as GPUs and CPUs ; 2 . Speed . The time overhead should be acceptable , ideally as small as possible ; 3 . Model Performance . The loss of accuracy should be acceptable , ideally as small as possible . Although storing the compressed activations successfully saves the memory for CNNs , up to our knowledge , there is no existing work extends this direction to GNNs and evaluates the mentioned trade-off for analyzing its feasibility . Despite the extension is conceptually straightforward , this direction is less-explored since it can be notoriously difficult to implement to fully leverage hardware potentials . This dilemma stems from the fact that the necessary tools for supporting this direction is usually missed in common graph learning packages . For example , operations in popular graph learning packages only support casting tensors down to 8-bit integer on GPUs , significantly limiting the memory saving ( Paszke et al. , 2019 ; Fey & Lenssen , 2019 ; Wang et al. , 2019 ) . As a result , previous GNN quantization works either emulates inference-time quantization via “ simulated quantization ” ( Tailor et al. , 2021 ; Zhao et al. , 2020 ) , or is impractical to use GPUs for accelerating ( Feng et al. , 2020 ) . To unleash the potential of this direction , we provide a space-efficient GPU implementation for supporting common operations in GNNs with compressed activations . Equipped with our implementation , this paper asks the following open question : To what extent can we compress the activations with both acceptable loss in accuracy and time overhead , for scalable GNN training ? To answer the open question , we first explore two different types of compression methods . Namely , one is “ quantization ” that compresses the activations into lower numerical precision . The other one is called “ random projection ” ( Achlioptas , 2001 ) that projects the activations into low-dimensional space . Both these two simple strategies can achieve near-lossless accuracy at a non-trivial compression ratio . For example , the loss in accuracy is negligible ( 0.2 % ) even under the vanilla 2-bit quantization . However , we can not further push forward the memory saving by these two methods , e.g. , we can not use a numerical precision below 1-bit . Considering that the real-world graphs often contain hundreds of millions of nodes , our main goal is to trim down the memory consumption to the maximum extent among the three aspects , as long as the other two are acceptable . We then naturally explore the direction of combining random projection and quantization , dubbed “ EXACT ” , to aggressively maximize the memory saving . EXACT is essentially applies random projection and quantization sequentially for compressing activations . Despite the superior memory saving , another following question is that , whether the combination bring significantly worse model performance and larger time overhead ? Following the questions , we make three major contributions as follows : • We provide a space-efficient GPU implementation for training GNNs with compressed activations as an extension for Pytorch . Based on our implementation , we are the first one training GNNs with compressed activations and demonstrating its potential for real-world usage . EXACT can complement the existing studies , as it can be integrated with most of existing solutions . • We propose EXACT , a simple-yet-effective framework which applies random projection and quantization sequentially on activation for scalable GNN training . We theoretically and experimentally show that applying random projection and quantization sequentially have an “ interaction ” effect . Namely , from the model performance aspect , after random projection , applying quantization only has a limited impact on the model performance . From the time aspect , EXACT runs comparable or even faster than quantization only . • Despite the simplicity , EXACT achieves non-trivial memory saving with both acceptable time overhead and loss in accuracy : EXACT can reduce the memory footprint of activations by up to 32× with roughly 0.5 % loss in accuracy and 10− 25 % time overhead across models and datasets . Noteworthily , EXACT trims down the hardware requirement of training a full-batch GraphSAGE ( Hamilton et al. , 2017 ) on ogbn-products ( Hu et al. , 2020 ) from a 48GB GPU to a 12GB GPU , which is affordable to most research labs . 2 THE MEMORY CONSUMPTION OF GNNS . Background . Let G = ( V , E ) be an undirected graph with V = ( v1 , · · · , v|V| ) and E = ( e1 , · · · , e|E| ) being the set of nodes and edges , respectively . Let X ∈ R|V|×d be the node feature matrix of the whole graph . The graph structure can be represented by an adjacency matrix A ∈ R|V|×|V| , where Ai , j = 1 if ( vi , vj ) ∈ E else Ai , j = 0 . In this work , we are mostly interested in the task of node classification , where the goal is to learn the representation hv for all v ∈ V such that the label yv can be easily predicted . To obtain such a representation , GNNs follow the neighborhood aggregation scheme . Specifically , GNNs recursively update the representation of a node by aggregating representations of its neighbors . Formally , the lth layer of a GNN can be written as h ( l+1 ) v = UPDATE ( h ( l ) v , ⊕ u∈N ( v ) MSG ( h ( l ) u , h ( l ) v ) ) , where h ( l ) v is the representation of node v at the lth layer . N ( v ) denotes the neighboring nodes of node v , not including v itself . The full table of notations can be found in Appendix A Table 5 . For node v , messages from its neighbors are calculated by MSG ( · ) . Then these messages are aggregated using a permutation-invariant aggregation function ⊕ . The aggregated features at v are updated by UPDATE ( · ) . For example , the Graph Convolutional Network ( GCN ) ( Kipf & Welling , 2017 ) layer can be defined as H ( l+1 ) = ReLU ( ÂH ( l ) Θ ( l ) ) , ( 1 ) where H ( l ) is the node embedding matrix consisting of all nodes ’ embeddings at the lth layer and H ( 0 ) = X . Θ ( l ) is the weight matrix of the lth GCN layer . Â = D̃− 1 2AD̃− 1 2 is the normalized adjacency matrix , where D̃ is the degree matrix of A + I . We note that Â is usually stored using the sparse matrix format . For GCN layers , the message and the aggregation function are fused into a Sparse-Dense Matrix Multiplication ( SPMM 1 ) operation . Namely , ÂH ( l ) Θ ( l ) = SPMM ( Â , H ( l ) Θ ( l ) ) . Thus , the computation graph of Equation 1 can be written as H ( l+1 ) = ReLU ( SPMM ( Â , MM ( H ( l ) , Θ ( l ) ) ) ) , ( 2 ) where MM ( · , · ) is the normal Dense-Dense Matrix Multiplication . Equation 2 resembles how GCNs are implemented in popular packages ( Fey & Lenssen , 2019 ; Wang et al. , 2019 ) . Here we analyze the memory consumption for the forward pass of GCN since most of the memory is occupied during the forward pass . For the memory usage of the backward pass , a detailed analysis is given in Appendix C. Specifically , for an L layer GCN , suppose the hidden dimensions of layer 0 , · · · , L − 1 are the same , which are denoted as D. The forward pass of GCN layers is shown in Appendix C Algorithm 1 . For layer l , it saves the following four variables in memory . ( 1 ) the weight matrix Θ ( l ) ∈ RD×D whose shape is independent to the graph size and is generally negligible . ( 2 ) the normalized adjacency matrix Â ( in CSR format ) whose space complexity is O ( |V|+ |E| ) . Note that there only needs to store one Â in memory which can be accessed by different layers . Thus , the memory consumption of Â is independent of the number of layers and is not the main memory bottleneck . ( 3 ) the intermediate result J ( l ) = MM ( H ( l ) , Θ ( l ) ) ∈ R|V|×D . For an L layer GCN , storing { J ( 0 ) , · · ·J ( L−1 ) } has a O ( L|V|D ) space complexity , which is the main memory bottleneck . ( 4 ) the node embedding matrix H ( l ) ∈ R|V|×D . For a L layer GCN , storing { H ( 0 ) , · · ·H ( L−1 ) } has a O ( L|V|D ) space complexity , which is the also main memory bottleneck . In this paper , we use the term “ activation maps ” to encompass H , J , and activation maps of other commonly used layers/operations such as BatchNorm , ReLU and Dropout .

Summary: This paper explores the training of GNN with compressed activation maps. It provides an optimized GPU implementation and comprehensively studies the trade-off among the memory saving, time overhead, and accuracy drop.  Experimental results show that the proposed framework can reduce the memory footprint of activations by up to 32x with only 0.2-0.5% accuracy drop and 10-25% time overhead.