EXACT: S CALABLE G RAPH N EURAL N ETWORKS T RAINING VIA E XTREME A CTIVATION C OMPRESSION
A BSTRACT
Training Graph Neural Networks (GNNs) on large graphs is a fundamental challenge due to the high memory usage, which is mainly occupied by activations (e.g., node embeddings).
Previous works usually focus on reducing the number of nodes retained in memory.
In parallel, unlike what has been developed for other types of neural networks, training with compressed activation maps is less explored for
GNNs.
This extension is notoriously difﬁcult to implement due to the miss of necessary tools in common graph learning packages.
To unleash the potential of this direction, we provide an optimized GPU implementation which supports training
GNNs with compressed activations.
Based on the implementation, we propose a memory-efﬁcient framework called “EXACT”, which for the ﬁrst time demonstrate the potential and evaluate the feasibility of training GNNs with compressed activations.
We systematically analyze the trade-off among the memory saving, time overhead, and accuracy drop.
In practice, EXACT can reduce the memory footprint of activations by up to 32 × with 0 .
5% accuracy drop and 10 - 25% time overhead across different models and datasets.
We implement EXACT as an extension for Pytorch Geometric and Pytorch.
In practice, for Pytorch Geometric,
EXACT can trim down the hardware requirement of training a three-layer fullbatch GraphSAGE on ogbn-products from a 48GB GPU to a 12GB GPU.
The code is available at https://github.com/warai-0toko/Exact .
I NTRODUCTION
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 simpliﬁed 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:
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 difﬁcult 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, signiﬁcantly 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-efﬁcient GPU implementation for supporting common operations in GNNs with compressed activations.
Equipped with our implementation, this paper asks:
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 ﬁrst 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 cannot further push forward the memory saving by these two methods, e.g., we cannot 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 signiﬁcantly worse model performance and larger time overhead?
Following the questions, we make three major contributions as follows:
• We provide a space-efﬁcient GPU implementation for training GNNs with compressed activations as an extension for Pytorch.
Based on our implementation, we are the ﬁrst 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.
We implement EXACT as an extension for Pytorch Geometric and Pytorch.
For Pytorch
Geometric, 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.
T HE MEMORY CONSUMPTION OF GNN S
Background.
Let G =
( V , E ) be an undirected graph with V
= ( v 1 , · · · , v |V| ) and E = ( e 1 , · · · , 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 A i,j = 1 if ( v i , v
j ) ∈
E else
A i,j = 0 .
In this work, we are mostly interested in the task of node classiﬁcation , where the goal is to learn the representation h v for all v ∈ V such that the label y v can be easily predicted.
To obtain such a representation, GNNs follow the neighborhood aggregation scheme.
Speciﬁcally, GNNs recursively update the representation of a node by aggregating representations of its neighbors.
Formally, the l th layer of a GNN can be written as h ( l +1) v = U PDATE (cid:18) h ( l ) v , (cid:76) u
[Equation]
v ) M SG (cid:0) h ( l ) u , h ( l ) v (cid:1)(cid:19) , where h ( l ) v is the representation of node v at the l th 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 M SG ( · ) .
Then these messages are aggregated using a permutation-invariant aggregation function (cid:76) .
The aggregated features at v are updated by U PDATE ( · ) .
For example, the Graph
Convolutional Network (GCN) (Kipf & Welling, 2017) layer can be deﬁned as
[Equation]
where H ( l ) is the node embedding matrix consisting of all nodes’ embeddings at the l th layer and H (0)
=
X .
Θ ( l ) is the weight matrix of the l th GCN layer.
[Equation]
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, ˆ AH ( l ) Θ ( l ) =
SPMM ( ˆ A ,
H ( l ) Θ
( l ) ) .
Thus, the computation graph of Equation 1 can be written as
[Equation]
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. Speciﬁcally, 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 ) ∈
R D × 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.
M ETHODOLOGY
From the above analysis, the memory consumption of GCNs can be signiﬁcantly reduced by storing only the compressed activations.
As shown in Figure 1, due to the simplicity and small time overhead of quantization (Section 3.1) and random projection (Section 3.2), we ﬁrst explore these two methods for compressing the activations of GNNs.
We show that these two simple methods
[Figure]
can moderate ratio negligible loss in accuracy highest
- pression rates methods are either limited by hardwares or limited by the accuracy drop
Motivated by GNN’s endless demand on memory, we then explore the direction of combining quantization with random projection into one framework termed “EXACT”, to aggressively maximize the memory saving (Section 3.3).
We show that random projection and quantization 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
3. S TORING Q UANTIZED A CTIVATIONS
Inspired by ActNN (Chen et al., 2021a), as shown in Figure 1a, we ﬁrst explore the direction of compressing activation maps in lower numerical precision.
Speciﬁcally, during the forward pass, each layer’s accurate activations (e.g., J ( l ) and H ( l ) ) is used to compute those for the subsequent layer.
Then the accurate activations will be quantized into the compressed activations (e.g., H ( l ) INT and J ( l ) INT ) with lower numerical precision, which overwrites the the accurate activations and is retained in the memory.
During the backward pass, we dequantize the compressed activations back to full-precision (e.g., ˆ H ( l ) and ˆ J ( l ) ).
Then the gradients are computed based on the dequantized activations.
We note that all operations (e.g., MM and SPMM ) are done in full-precision since most of the GPUs do not support operands with bit-width other than full-precision and half-precision.
For convenience, we use the term “precision” to present “numerical precision” in the following of this paper.
Below we introduce the technical details of (de)quantization.
Speciﬁcally, each node embedding vector h ( l ) v will be quantized and stored using b -bit integers.
Let
B = 2 b − 1 be the number of quantization bins.
The integer quantization can be expressed as
[Equation]
where Z ( l ) v = min { h ( l ) v } is the zero-point, r ( l ) v = max { h ( l ) v } − min { h ( l ) v } is the range for h ( l ) v , and (cid:98) · (cid:101) is the stochastic rounding operation (Courbariaux et al., 2015).
H ( l ) INT in Figure 1a is the matrix containing all quantized h ( l ) v INT .
During the backward pass, each h ( l ) v INT will be dequantized as
[Equation]
ˆ H ( l ) in Figure 1a is the matrix containing all dequantized embeddings ˆ h ( l ) v .
The following proposition characterizes the effect of quantization, which is adopted from Chen et al.
(2021a).
[Equation]
Proposition 1 (Details in Appendix E)
The above quantization and dequantization are unbiased operations, i.e., E [ˆ h ( l ) v ] =
E [Dequant(Quant( h ( l ) v ))]
= h ( l ) v and Var(ˆ h ( l ) v ) =
D [ r ( l ) v ] 2 6 B 2 .
From the unbiased nature of quantization illustrated in Proposition 1, the calculated gradient is also unbiased.
The approach imposes extra noise (i.e., the variance term in Proposition 1) to the gradient during the backward pass.
From Proposition 1, the noise effect is inversely-correlated with the number of quantization bins B .
To evaluate how the noise affects the model performance, we train three popular GNN models, namely, GCN, GraphSAGE, and GAT, on the ogbn-arxiv dataset (Hu et al., 2020) with different precision (Table 1).
The detailed experiment setting is elaborated in
Appendix G.2.
Here we emphasize that all these three models are trained with full-batch data .
[Table]
Observations.
For all three models, the loss in accuracy is negligible ( ≈ 0 .
2% ) even using the vanilla 2 -bit quantization.
In contrast, for CNNs, adopting the vanilla INT1 or INT2 quantization, will cause a signiﬁcant accuracy drop (usually > 5% )
(Chen et al., 2021a; 2020a).
This observation can be explained by the following mechanism.
We show in Appendix E.4 that the approximation error will compound layer-by-layer during the backward pass.
The layer-depth of GNNs is much smaller than CNNs due to the over-smoothing problem.
And hence GNNs is much more noisetolerant than CNNs.
Our observation suggests that in practical scenarios, there is no need for an expensive sophisticated quantization approach (e.g., mixed precision quantization) as considered in previous works (Chen et al., 2021a; Fu et al., 2020).
3. S TORING R ANDOMLY P ROJECTED A CTIVATIONS
Here we explore the direction of random projection.
The key idea is to project the activation maps into a low-dimensional space that keeps the original information as much as possible.
In this way, similar to the framework in Section 3.1, we only need to store the dimension reduced activation maps in memory.
Figure 1b illustrates the workﬂow of GCN layers.
During the forward pass, each layer’s accurate activations (e.g., h ( l ) v ) are used to compute those for the subsequent layer.
Then the accurate activations will be projected into low-dimensional space, which can be expressed as:
[Equation]
(5) where R ∈ R D
× R is a random matrix ( R < D ) which satisﬁes E [ RR (cid:62) ]
= I .
H ( l ) proj ∈
R |V|× R is the matrix containing all projected node embeddings h ( l ) v proj .
In this paper, R is set as the normalized
Rademacher random matrix (Achlioptas, 2001) due to its low sampling cost (detailed introduction in Appendix D).
After projection, we store only H ( l ) proj in memory and hence the compression ratio is DR .
During the backward pass, the projected node embeddings are inversely transformed by
[Equation]
where IRP( · ) is the inverse projection operation.
ˆ H ( l )
= H ( l ) proj RR (cid:62) ∈
R |V|× D is the recovered activation map containing all ˆ h ( l ) v .
The following proposition shows the effect of random projection.
Proposition 2 (Proof in Appendix E)
The above RP and IRP are unbiased operations, i.e.,
E [ ˆ H ( l ) ] =
E [IRP(RP( H ( l ) ))]
= H ( l ) .
For each ˆ h ( l ) v in ˆ H ( l ) , we have Var(ˆ h ( l ) v ) =
D − 1 R
|| h ( l ) v
|| 22 .
From the unbiased nature of random projection illustrated in Proposition 2, the calculated gradient is also unbiased.
The approach here also only imposes extra variance to the calculated gradient, where the variance linearly scales with the D R ratio.
To quantitatively study the effect of the extra variance
[Table]
in scenarios of practical interest, following the same setting in Section 3.
[Table]
we show the performance of GCN, GraphSAGE
[Table]
and GAT trained on the ogbn-arxiv dataset with different DR ratio in Table 2.
Observations.
We make two main observations: (1) For all three models, the loss in accuracy tends to range from negligible ( ≈ 0 .
to moderate ( ≈ 0 .
5% ) when DR ≤ 8 .
(2) Under the same compression ratio, the quantization is better than random projection in terms of the loss in accuracy (e.g., compare the DR = 16 result in Table 2 with the INT2 result in Table 1).
3. EXACT: C OMBINING RANDOM PROJECTION AND QUANTIZATION
We experimentally show that GNNs are noise-tolerant to inaccurate activations compressed by quantization and random projection.
However, the highest compression ratio of these two methods is limited.
Namely, the precision cannot be less than 1-bit for quantization and DR cannot surpass eight for random projection due to a large accuracy drop.
However, GNN’s demand on memory is endless since real-world graphs can contain hundreds of millions of nodes and the graph size is ever growing.
Motivated by this fact, we try to further push forward the memory saving as long as the loss of accuracy and the time overhead are both acceptable.
We explore the direction of combining these two methods into one framework termed “EXACT” to maximize the memory saving.
Speciﬁcally, we compress the activations and store only the ˜ h
( l ) v = Quant(RP
( h ( l ) v )) in memory during the forward pass.
During the backward pass, the node embedding is recovered as ˆ h ( l ) v = IRP(Dequant(˜ h ( l ) v )) .
EXACT can achieve a superior compression ratio, e.g., the compression ratio of EXACT is roughly 128 × if DR = 8 and the precision is INT2.
However, one practical question is that, will the variance explode, leading to signiﬁcantly worse performance when applying random projection and quantization sequentially?
Here we try to answer the above question both theoretically and experimentally.
Below we ﬁrst theoretically show that after applying random projection, the quantization range of projected node embeddings is bounded.
Proposition 3 (Proof in Appendix E)
For each projected node embedding h ( l ) v proj = RP( h ( l ) v )
=
h ( l ) v R , for ∀ (cid:15) > 0 , by choosing s = ||
h ( l ) v || 2 (cid:113) 2ln(2 R/(cid:15) )
R , we have P ( || h ( l ) v proj || ∞ ≤ s ) ≥ 1 − (cid:15) .
From Proposition 3, after projection, the maximal absolute value among weights in h ( l ) v proj is bounded.
Following the fact that the quantization range r ( l ) v proj = max { h ( l ) v proj } − min { h ( l ) v proj } ≤ 2 || h ( l ) v proj || ∞ , r ( l ) v proj is also bounded.
Recall that the variance of quantization scales with [ r ( l ) v proj ] 2 B(Proposition 1), applying random projection and quantization sequentially will not lead to the variance explosion.
We also experimentally visualize the inﬁnity norm of projected embeddings in
Appendix E Figure 4.
The inﬁnity norm of projected embeddings may be even less than those of original embeddings when R is larger than a threshold ( R = 0 .
5 D in Figure 4).
Thus, the extra variance is expected to be dominated by the random projection, and hence the loss in accuracy is largely determined by the DR ratio.
We also experimentally investigate the effect of
EXACT by plotting the test accuracy against each of quantization precision and the DR ratio.
Due to the page limit, we only present one representative result in Figure 2, namely, training GCNs on ogbn-arxiv with EXACT.
More similar results can be found in Appendix I.
Here we do not consider INT1 precision since its model performance drop is already near 0.5%.
We summarize two main observations.
First, the performance of EXACT is largely determined by the DR ratio of random projection.
Second, when DR ≤ 8 , the loss in accuracy generally ranges from ≈ 0 .
2% to ≈ 0 .
5% , regardless of the quantization precision.
System Implementation.
In EXACT, each operations (e.g., SPMM , ReLU, and BatchNorm) is reimplemented with different conﬁgurations using CUDA kernels.
For example, BatchNorm only supports quantization, while the random projection is not applicable to it.
Since Pytorch only supports precision down to INT8, we convert quantized tensors into bit-streams by CUDA kernels to maximize the memory saving.
The conﬁguration and implementation details are given in Appendix F.
R ELATED W ORK AND D ISCUSSION
Due to the page limit, we brieﬂy review and discuss the relationship between existing works and
EXACT.
A more comprehensive discussion can be found in Appendix B. Existing works can be roughly divided into scalable/efﬁcient GNN inference and scalable GNN training according to the problem they try to solve.
Almost all previous GNN quantization works try to solve the ﬁrst problem, which is much simpler than the second problem that EXACT tries to address.
Speciﬁcally, most of them try to enable the usage of low precision integer arithmetic during inference (Tailor et al., 2021;
Zhao et al., 2020) by simulating the quantization effect, which may even increase the memory consumption during training.
Feng et al.
(2020) tries to address the second problem by proposing a heterogeneous quantization framework which assigns different bits to node embeddings in each layer.
However, this framework is impractical on off-the-shell hardwares.
For scalable GNN training methods, EXACT is orthogonal to most of them, including distributed training (Zheng et al., 2020b;
Jia et al., 2020; Zhu et al., 2019; Wan et al., 2021), subgraph sampling (Hamilton et al., 2017; Zeng et al., 2020; Chiang et al., 2019), and historical embedding-based methods (Fey et al., 2021).
We discuss the potential beneﬁt of integrating EXACT with them in Appendix B .
E XPERIMENTS
The experiments are designed to answer the following research questions.
RQ1 : How effective is
EXACT in terms of model performance at different compression rates (Section 5.1)?
RQ2 : Is the training process of deeper GNNs also robust to the noise introduced by EXACT (Section 5.1)?
RQ3 :
How sensitive is EXACT to its key hyperparameters (Appendix I.4)?
RQ4 : What is the running time overhead of EXACT (Section 5.2)?
RQ5 : Is the convergence speed of GNNs impacted by EXACT (Appendix I.3)?
RQ6 :
To what extent can EXACT reduce the hardware requirement for training
GNNs on large graphs (Section 5.3)?
Datasets and Models: To evaluate the scalability of EXACT, we adopt ﬁve common large scale graph benchmark datasets from different domains.
Namely, Reddit, Flickr, Yelp, ogbn-arxiv , and ogbn-products .
We evaluate EXACT under both the mini-batch training and full-batch training setting.
For the mini-batch training setting, we integrate EXACT with two state-of-the-art subgraph sampling methods, namely, Cluster-GCN (Chiang et al., 2019) and GraphSAINT (Zeng et al., 2020).
For the full-batch training setting, we integrate EXACT with three popular models: two commonly used shallow models, namely, GCN (Kipf & Welling, 2017) and GraphSAGE (Hamilton et al., 2017), and one deep model GCNII (Chen et al., 2020b).
To avoid creating confusions, GCN,
GraphSAGE, and GCNII are both trained with the whole graph at each step .
For a fair comparison, we use the mean aggregator for GraphSAGE, Cluster-GCN, and GraphSAINT throughout the paper.
Details about the hyperparameters of models and datasets can be found in Appendix G.
Hyperparameter Settings: From Section 3.3, we show that the performance of EXACT is largely determined by the D R ratio of random projection, and the model performance drop with INT2 quantization is negligible.
Thus, to balance the accuracy drop and memory saving ratio, we adopt INT precision with different DR ratio for EXACT.
Speciﬁcally, EXACT (INT2) indicates it only applies 2bit quantization to the activation maps of all applicable operations (see Table 8).
EXACT(RP+INT2) indicates it applies the random projection followed by a 2-bit quantization to the activations of all applicable operations.
We perform a grid search for DR ratio from { 2, 4, 8 } .
Detailed DR conﬁguration for EXACT(RP+INT2) can be found in Table 12.
[Table]
[Figure]
The time overhead of EXACT is roughly 12% - 25% .
We discuss about the swapping and gradient checkpointing in Appendix I.2.
5. T RADE -O FF BETWEEN S PACE AND M ODEL P ERFORMANCE
To answer RQ1 and RQ
[Figure]
we ﬁrst analyze the trade-off between the model performance and memory saving.
Table 3 summarizes the model performance and the memory footprint of the activation maps.
Besides the memory usage of activation maps, a detailed analysis about the overall memory saving ratio is provided in Appendix H. We make two main observations.
First, EXACT aggressively reduces the memory footprint with ≤ 0 .
5% loss in accuracy.
From
Table 3, EXACT (INT2) reduces the memory footprint of activation maps from 7 × to 16 × , with negligible loss in accuracy ( ≈ 0 .
in almost all experiments.
EXACT (RP+INT2) generally reduces the memory footprint of activation maps from 10 × to 32 × , and the loss in accuracy tends to range from lossless ( 0 . 0% ) to moderate ( ≈ 0 .
Although EXACT can aggressively compress the memory usage of activations associated with MM and SPMM by up to 128 × , the compression ratio for other operations usually cannot surpass 32 × (e.g., the activation maps of ReLU and Dropout take exact one bit to store, see Appendix F).
As a result, the overall compress ratio is pulled down.
We show in Appendix H that the measured compression ratio is consistent with the theoretical one.
Second, the training process of deeper models is also robust to extremely compressed activations.
As we analyzed in Appendix E.4, the vanilla quantization and random projection obtain unreasonable good performance might because the commonly used GNNs are shallow.
Thus, one open question is that, to what extent can deeper GNNs robust to extremely compressed activation maps ( RQ2 )?
The ablation study of GCNII in Table 3 tries to experimentally answer this question.
Namely, the training process of deeper GNNs (up to 16-layer, see Table 17) is also robust to the noise introduced by the extremely compressed activations, with up to 32 × less memory footprint of activation maps.
We note that building deeper GNNs is still an open direction due to the oversmoothing and the extensive memory consumption problem (Oono & Suzuki, 2019; Li et al., 2021;
Chen et al., 2020b).
From the practical usage perspective, our GCNII ablation study indicates that
EXACT may enable to train deeper GNNs on large graphs with minimal loss in performance.
To answer RQ3 , we present the sensitivity study of EXACT (RP+INT2) to DR ratio in Appendix I.4.
In general, the model performance drop of EXACT (RP+INT2) increases with the DR ratio.
In summary, when DR = 8 , the loss in accuracy of EXACT (RP+INT2) is ≤ 0 .
5% in roughly two-third of experiments and ≤ 1% in almost all experiments.
5. T RADE -O FFS BETWEEN S PACE AND S PEED
To answer RQ4 , we compare the training throughput of EXACT with the baseline using a single
RTX 3090 (24GB) GPU.
Figure 3 shows the results among EXACT with different conﬁgurations and the baseline (see Table 9 for detailed package information).
We make two main observations.
First, the overhead of EXACT (INT2) is roughly 12% ∼ 25% .
The overhead of EXACT (INT2) comes entirely from quantization.
We note that EXACT (INT2) adopts vanilla integer quantization with negligible loss in accuracy, which has the lowest overhead among all quantization methods.
Sophisticated quantization strategies (e.g., mixed precision quantization and non-uniform quantization) further increase the time overhead, because they require an extra search process to ﬁnd precision or quantization step for each layer/sample (Chakrabarti & Moseley, 2019;
Fu et al., 2020; Chen et al., 2021a).
From Figure 3, the overhead of EXACT is roughly 12% ∼ 25% .
Second, the running speed of EXACT (RP+INT2) is comparable or even faster than quantization only.
The above counter-intuitive observation can be explained by the following mechanism.
Although random projection introduces two extra matrix multiplication operations, the total number of elements to be quantized is DR times less than directly quantizing the activations.
Considering that quantization is the main bottleneck, compared to EXACT (INT2), EXACT (RP+INT2) achieves roughly twice overall compression ratios with a comparable or even smaller time overhead.
RQ4 focuses on the actual running speed measured by the hardware throughput.
Here we investigate the convergence speed from the optimization perspective ( RQ5 ).
To answer RQ5 , due to the page limit, we present the training curve of models trained on Reddit dataset with EXACT in
Appendix I.3.
The converge speed here is measured by the number of epochs it takes to reach a certain validation accuracy.
In summary, the convergence speed of EXACT is slightly slower than the baseline, with limited impact on ﬁnal accuracy.
5. A BLATION S TUDIES
Train full-batch GraphSAGE on ogbn-products using a GPU with 11GB memory.
[Table]
Table 3 shows that EXACT reduces the memory footprint of activation maps by up to 32 × .
However, as analyzed in Appendix H, besides the activation maps, input data (including feature matrices X , adjacency matrix
A , and labels), activation gradients also occupy memory in the full-batch training.
For ogbn-products , the input data occupies 4.7GB (see Appendix H).
From the
OGB leader-board, training a full-batch GraphSAGE requires GPUs with at least 48GB memory.
We note that EXACT can be integrated with AMP (see
Appendix F) to further squeeze out the memory.
To answer RQ6 , by integrating AMP and EXACT (RP+INT2) with DR = 4 , we successfully train a full-batch GraphSAGE on ogbn-products on a single GTX 1080 Ti with 11GB memory, with moderate loss ( ≈ 0 .
5% ) in accuracy (Table 4).
C ONCLUSION AND F UTURE W ORK
In this paper, we propose EXACT, a simple-yet-effective framework for training GNNs with compressed activations.
We demonstrate the potential of EXACT for the real-world usage by systematically evaluating the trade-off among the memory-saving, time overhead, and accuracy drop.
We show that EXACT is orthogonal to most of the existing existing solutions and discuss how EXACT can be integrated with them in Appendix B. Future work includes (1) evaluating EXACT under multi-GPU settings for distributed training; (2) combining EXACT with historical embedding-based method; (3) combining EXACT with memory-efﬁcient training system, such as swapping.
A CKNOWLEDGEMENT
We would like to thank all the anonymous reviewers for their valuable suggestions.
Thank all the members of Samsung Research America advertisement team for your feedbacks, and everyone who has provided their generous feedback for this work.
This work is, in part, supported by NSF IIS1750074.
The views and conclusions contained in this paper are those of the authors and should not be interpreted as representing any funding agencies
A N OTATIONS
[Table]
matrix
[Table]
of the vector h , where Var( h )
=
E [ h (cid:62) h ] − ( E [ h ]) (cid:62) ( E [ h ])
[Table]
The 2-norm of vector h , where | h || 2 = √ h (cid:62) h | h || ∞
The inﬁnity norm of vector h , where | h || ∞ = max | h |
B R ELATED W ORK AND D ISCUSSION
Due to the page limit, we discuss the relationship between EXACT and existing works in detail here.
GNN is notoriously inefﬁcient and non-scalable during both the training and inference phase.
Previous works can be divided into two categories according to the problem they try to solve.
Namely, scalable/efﬁcient GNN inference and scalable GNN training.
Below we introduce and discuss the relationship between them and EXACT.
B. S CALABLE /E FFICIENT GNN I NFERENCE
Scalable GNN inference is an orthogonal direction to the problem that EXACT tries to address.
Below we introduces two popular directions for scalable GNN inference.
Namely, GNN quantization and graph sparsiﬁcation.
GNN Quantization.
Most of the previous GNN quantization works focus on enabling the usage of low precision integer arithmetic during inference.
Speciﬁcally, Tailor et al.
(2021) proposes a Quantization Aware Training method tailored for GNNs, which emulates inference-time quantization during the training phase.
Zhao et al.
(2020) proposes to jointly search the quantization bit-width and GNN architectures.
Note that both of them do not actually convert the node embeddings into lower numerical precision during training.
Instead, they use full precision data type to simulate the effect of real quantization.
Feng et al.
(2020) proposes a heterogeneous quantization framework which assigns different bits to node embeddings in each layer while maintaining the weights at full precision.
However, due to the mismatch in operands’ bit-width, it is impractical to use in general purpose hardwares, such as CPUs and GPUs.
Graph Sparsiﬁcation.
The adjacency matrices are usually stored as sparse matrices.
When hardwares perform SPMM , the sparse matrix format leads to random memory accesses and limited data reuse due to its irregular structure.
Thus, GNNs have higher inference latency than other neural networks.
Graph sparsiﬁcation can alleviate the issue by removing redundant edges from original graphs.
Zheng et al.
(2020a) proposes a learning-based graph sparsiﬁcation method which removes potentially task-irrelevant edges from input graphs Li et al.
(2020) formulates the graph sparsiﬁcation problem as an optimization objective which can be solved by alternating direction method of multipliers (ADMM).
Chen et al.
(2021b) proposes to co-simplify the input graph and GNN model by extending the iterative magnitude pruning to graph areas.
B. S CALABLE GNN T RAINING
In this subsection, we introduce previous works that try to train GNNs on large graphs.
We conclude that EXACT are parallel to most of the previous works.
Subgraph Sampling/Mini-Batch Training.
As we introduced in the main body, most of the previous scalable GNN training methods fall into this category.
The key idea is 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.
Based on this idea, various sampling techniques have been proposed, including the node-wise sampling (Hamilton et al., 2017;
Chen et al., 2017; Markowitz et al., 2021), layer-wise sampling (Zou et al., 2019; Chen et al., 2018;
Huang et al., 2018), and subgraph sampling (Chiang et al., 2019; Zeng et al., 2020).
Generally, methods in this category are orthogonal to EXACT, and they can certainly be combined.
Similar to EXACT, subgraph sampling also introduce the gradient noise during the training process.
However, the gradient noise does not necessarily result in the model performance drop.
In contrast, many previous works experimentally show that subgraph sampling based methods may even yield better performance than the whole graph training (Zeng et al., 2020; Hu et al., 2020).
Our experiment results also exhibit a similar phenomenon that EXACT may achieve better performance than the original one.
This observation can be explained by the regularization effect introduced by the gradient noise (Hu et al., 2020).
In this paper, we show in experiments that EXACT can be combined with subgraph-sampling based methods with only limited accuracy drop.
From the practical usage perspective, EXACT can enlarge the maximal subgraph size for subgraph training-based method.
As a result, the error introduced by
EXACT can be compensated by using larger subgraphs, which previously cannot ﬁt into GPUs.
We leave it as one future direction.
GNNAutoScale (Fey et al., 2021).
GNNAutoScale follows the idea of utilizing historical embeddings from prior training iterations (Fey et al., 2021).
Speciﬁcally, it stores a historical embedding for each node as an ofﬂine storage in CPU memory.
At each training step, it ﬁrst samples a minibatch of nodes.
Then for out-of-mini-batch nodes, GNNAutoScale swaps their historical embeddings from CPU memory to GPU memory and utilizes their non-trainable historical embeddings for propagation.
In this way, only the node embeddings inside the current mini-batch and those of their direct 1-hop neighbors are retained in memory.
The time overhead of GNNAutoScale is mainly from the swapping step.
EXACT is orthogonal to GNNAutoScale.
It would be interesting to store node embeddings compressed by EXACT as historical embeddings for GNNAutoScale.
Considering that quantization is more time-efﬁcient compared to swapping, the time overhead of GNNAutoScale can be greatly reduced since the time cost of swapping scales with the total number of bits to be swapped.
Distributed Training.
Unlike the data in other domains, the graph data cannot be trivially divided into mini-batches due to its connectivity between samples.
The graph distributed training methods require to split the graph into mini-batches that minimizes the communication overhead.
Speciﬁcally, AliGraph (Zhu et al., 2019) is a distributed GNN framework on CPU platforms, which does not support GPUs acceleration.
ROC (Jia et al., 2020) learns to optimize the graph partitioning by predicting the execution time of performing a GNN operation on an input subgraph.
PipeGCN (Wan et al., 2022) proposes to train GNNs with stale features and stale feature gradients under the distributed training setting.
Both BNS-GCN (Wan et al., 2021) and DistDGL (Zheng et al., 2020b) leverage METIS (Karypis & Kumar, 1998) to performance the graph partitioning to minimize the communication cost.
However, one main drawback of applying METIS is that the ahead-of-training overhead of METIS is relatively large on huge graphs.
EXACT is also orthogonal to the distributed training methods and can be used to decrease the hardware requirements.
It would be interesting to investigate the trade-off among the accuracy drop and the memory saving of EXACT under multiGPU settings.
Memory efﬁcient training system.
Gradient checkpointing (Chen et al., 2016; Jain et al., 2019;
Kirisame et al., 2020; Shah et al., 2020) trades computation for memory by dropping some of the tensors activations in the forward pass and recomputing them in the backward pass.
Swapping (Meng et al., 2017) utilizes the huge amount of available CPU memory by swapping tensors between CPU and
GPU.
EXACT is also parallel to both the gradient checkpointing and swapping.
As a result, EXACT can be combined with them to further squeeze out the memory, at the cost of larger time overhead.
Other Activation Compression Techniques.
JPEG-ACT (Evans et al., 2020) extends the widelyused JPEG method to compress the activations.
Cusz is a prediction-based lossy compression method with GPU support (Tian et al., 2020; Tao et al., 2017).
Both of them can be applied to compress the activations of GNNs.
However, one practical problem is that they have hyperparameters which are hard to tune.
For example, in JPEG-ACT, there is a hyperparameter α controls the space/accuracy trade-off.
This problem can be alleviated by incorporating one recent proposed method AC-GC (Evans & Aamodt, 2021).
Given a preset allowable increase in loss, AC-GC can automatically ﬁnd the hyperparameter such that the compression rate can adapt to the preset training conditions.
We note that EXACT and AC-GC approaches the problem from different perspective.
For EXACT, we value the compression rate more than the accuracy drop.
Namely, given a preset memory budget, we try to ﬁt the model with acceptable accuracy drop.
For AC-GC, they value the accuracy change more than the compression rate in the sense that the compression rate is adapted to the given allowable accuracy drop.
C A NALYZING THE MEMORY CONSUMPTION OF GCN S
Generally, each training step contains a forward pass phase and a backward pass phase.
During the forward pass, activation maps of each operation are kept in memory for computing the gradients.
During the backward pass, the weight gradients and activation gradients are calculated and stored in memory.
For weight gradients, their memory usage is negligible since they share the same shape with model weights.
For activation gradients, their memory usage is dynamic and will eventually be cleared to zero.
In standard deep learning libraries, activation gradients will be deleted as soon as their reference counting becomes zero (Paszke et al., 2019; Abadi et al., 2016).
Thus, storing activation maps take up most of the memory.
The memory usage of activation maps in GCN layers are given in Algorithm 1.
D R ANDOM P ROJECTION
In this paper, we use the normalized Rademacher random matrix (Achlioptas, 2001) for projecting the activation maps.
Speciﬁcally, for a D × R normalized Rademacher random matrix R , each element in R is i.i.d.
sampled from the following distribution:
[Equation]
where (cid:113) 1 R is the normalize factor such that E [ RR (cid:62) ]
= I .
From above, we can see that the sampling cost of the normalized Rademacher random matrix is relatively low since sampling from
Bernoulli distribution is much cheaper than sampling from Gaussian distribution.
E T HEORY E. P ROOF OF P ROPOSITION
Proposition 1 (Details in Appendix E)
The above quantization and dequantization are unbiased operations, i.e., E [ˆ h ( l ) v ] =
E [Dequant(Quant( h ( l ) v ))]
= h ( l ) v and Var(ˆ h ( l ) v ) =
D [ r ( l ) v ] 2 6 B 2 .
Proposition 1 is adopted from the theoretical analysis in ActNN (Chen et al., 2021a).
For completeness, we prove it here with more details.
The conclusion of E [ˆ h ( l ) v ] = h ( l ) v follows from the fact that the stochastic rounding is an unbiased operation (Courbariaux et al., 2015).
Speciﬁcally, ∀ scalar h , the stochastic rounding can be expressed as
[Equation]
where (cid:100)·(cid:101) is the ceil operation and (cid:98)·(cid:99) is the ﬂoor operation.
First, we have E [ (cid:98) h (cid:101) ]
= h because
[Equation]
= h (following the fact that (cid:100) h (cid:101) − (cid:98) h (cid:99) = 1 )
Hence,
[Equation]
Regarding the variance, let ¯ h
= h
[Equation]
l
[Equation]
v − Z
[Equation]
B = ( h 1 , · · · , h D ) .
Suppose ∀ i, h
i − (cid:98) h
i (cid:99) = σ ∼
Uniform (0 , 1) , we have.
[Equation]
(substitute (cid:100) h
[Equation]
(cid:101) with (cid:98) h
[Equation]
(cid:99) + 1 )
[Equation]
(cid:98)
[Equation]
h 1 −
[Equation]
By taking expectation w.r.t.
σ on the both side, we have Var(ˆ h ( l ) v ) =
[Equation]
(cid:3)
E. P ROOF OF P ROPOSITION
Proposition 2 (Proof in Appendix E)
The above RP and IRP are unbiased operations, i.e.,
E [ ˆ H ( l ) ] =
E [IRP(RP( H ( l ) ))]
= H ( l ) .
For each ˆ h ( l ) v in ˆ H ( l ) , we have Var( ˆ h ( l ) v ) =
D − 1 R
|| h ( l ) v || 2 2 .
We can get E [ ˆ H ( l ) ] =
E [IRP(RP( H ( l ) ))]
=
H ( l ) directly following the fact that
E [IRP
[Equation]
Regarding the variance, let P = RR (cid:62)
∈
R D × D .
For the sake of notation convenience, we ignore the subscript and superscript in this subsection.
Namely, we use the notation ˆ h to represent ˆ h ( l ) v , and use the notation h to represent h ( l ) v .
Also, we assume the shapes of ˆ h and h are both 1 × D in this subsection.
We have
[Equation]
Next, we will characterize the distribution of P P (cid:62) .
Recall that P = RR (cid:62) , thus,
[Equation]
where each a i,j = ± (cid:113) 1 R with equal probability.
Thus,
[Equation]
From Equation 9, it is easy to show that E [ P P (cid:62) ] is a diagonal matrix.
Namely, ∀ i (cid:54) =
j , we have E
[ P P (cid:62) ] i,j =
following the fact that each a i,j are independent and E [ a i,j ] = 0 .
For elements on the diagonal, we have
[Equation]
= 1 + (cid:88) m (cid:54) =
i (cid:18) R (cid:88) k
=
E [( a i,k a m,k ) 2 ] + 2 (cid:88) p,q : p<q
E [ a i,p a m,p a i,q a m,q ] (cid:19) =
1 + (cid:88) m (cid:54) =
i (cid:18) R (cid:88) k
=
E [( a i,k a m,k ) 2 ] (cid:19) =
1 + D −
R .
Let h
= ( h 1 , · · · h D ) , we have
Var(ˆ h ) =
h E
[Equation]
P P (cid:62) ] h (cid:62)
[Equation]
(cid:62) =
h  
[Equation] [Equation]

[Equation]
(cid:62)
[Equation]
(cid:62) =
D (cid:88) i =
(1 + D
[Equation]
R )
h 2 i
[Equation]
D (cid:88) i =
h 2 i
=
D
[Equation]
R ||
h
|| 22 .
(11) (cid:3)
E. P ROOF OF P ROPOSITION
Proposition 3 (Proof in Appendix E)
For each projected node embedding h ( l ) v proj = RP( h ( l ) v )
=
h ( l ) v R , for ∀ (cid:15) > 0 , by choosing s = ||
h ( l ) v || 2 (cid:113) 2ln(2 R/(cid:15) )
R , we have P ( || h ( l ) v proj || ∞ ≤ s ) ≥ 1 − (cid:15) .
For the normalized Rademacher random matrix R ∈ R D
×
R , we have R i,j = ± 1 √ R with equal probability.
Let h ( l ) v =
( h 1 , · · · , h D ) and h ( l ) v proj =
( u 1 , · · · , u R ) .
Since h ( l ) v proj = h ( l ) v R , we have u 1 = (cid:80) Di
=1 a
i h i , where a i = ± 1 √ R with equal probability.
First, we have the following inequality:
[Equation]
(following the fact that cosh x ≤ e x 2 2 )
[Equation]
For notation convenience, we use the notation “ C ” to represent the node embedding norm || h ( l ) v || 2 .
For ∀ s, q > 0 , according to Chernoff bound, we have:
[Equation]
Note that Equation 13 holds for ∀ s > 0 .
Given a speciﬁc (cid:15) , by setting s = C (cid:113) 2ln(2 R/(cid:15) )
R , we have
P ( | u i | ≥ s ) ≤ (cid:15)R .
Also,
[Equation]
Here we also experimentally verify Proposition 3 by visualizing the inﬁnity norm of projected node embeddings in Figure 4.
To avoid creating confusions, DR = 1 in Figure 4 means we apply the random projection with R
= D on the activation maps.
We can observe that in general, the inﬁnity norm increases with the compression ratio DR .
However, we note that the inﬁnity norm of projected embeddings may less than those of original embeddings when R is larger than a threshold.
This suggests that when R is below a certain threshold, quantizing projected embeddings only have limited inﬂuence on model performance.
This claim is also veriﬁed in the main body of this paper.
E. T HE C OMPOUND E FFECT OF A PPROXIMATION E RRORS
In this subsection, we analyze why GNNs can tolerate the extremely compressed activations.
From the numerical analysis (Bauer, 1974), the jacobian matrix (gradient) between two tensors can be viewed as a sum over all ”paths” connecting these two tensors.
Moreover, the architecture of GNNs is considerably shallower than CNNs and transformers.
Informally, for GNNs, the key intuition is that these paths are signiﬁcantly ”shorter” compared to those of CNNs.
As a result, the approximation error (i.e., the variance from quantization) along the path is hard to get accumulated.
We adopt Theorem 3 in ActNN (Chen et al., 2021a) to make the above statement more rigorous.
Let ˆ
C ( m ) be the compressed context (either by quantization, random projection, or quantized random projection).
Let ∇ Θ ( l ) and ∇ H ( l ) be the gradient of Θ ( l ) and H ( l ) , respectively.
ˆ ∇ Θ ( l ) and ˆ ∇ H ( l ) are the calculated gradient using the compressed context, respectively.
Further, we use the notation G ( l ∼ m ) Θ ( ˆ ∇ H ( m ) , ˆ C ( m ) ) to represent the variance introduced by utilizing the compressed context ˆ C ( m ) .
Speciﬁcally, for a L layer GNN, we have
[Equation]
where Var( ·| ˆ ∇ H ( m ) ) is the conditional variance, and Var( ∇ Θ ( l ) ) is the variance is from the minibatch sampling variance.
The key insights from Equation 14 are two folds.
First, since the variances introduced by compressed contexts at different layers will accumulate, it suggests that the noise introduced by the compressed context is relatively small for shallow models.
Considering that most of GNNs are usually less than four layers, this may explain why the loss in accuracy is negligible when using the vanilla INT2 quantization and a relatively large DR ratio.
Second, under the subgraph training setting, the extra variance introduced by quantization can be compensated when using larger batch size (i.e., the size of subgraph in the graph learning).
In Table 1, we evaluate GNNs under the full-batch setting.
Therefore, in addition to the reason of mentioned shallower architecture, it is also possible that the activation compressed training beneﬁts from the large batch size, which leads to a much smoother gradient.
To quantitatively study the effect of the batch size, below we present the ablation studies of two representative sampling based methods with much smaller batch size.
Namely, GraphSAINT (Table 6) and Cluster-GCN (Table 7).
For the baseline, we use the same hyperparameters reported in the corresponding paper.
We make two main observations.
First, INT2 quantization works under a much smaller batch size.
Namely, for both Cluster-GCN and GraphSAINT, the accuracy drop is negligible even when the sampled subgraph contains only ≈ 500 nodes.
This observation implies that the activations of GNNs can be aggressively compressed regardless of the batch size.
Second, using a much smaller batch size in general will lead to an accuracy drop.
However, for Yelp, a much smaller batch size may even improve the performance.
[Table] [Table]
[Table]
F S YSTEM I MPLEMENTATION OF EXACT F. I NDIVIDUAL L AYERS C ONFIGURATIONS OF EXACT
The operation conﬁgurations are shown in Table 8.
For all graph convolution operations, EXACT can apply both the quantization and random projection to their saved activation maps.
In EXACT, the SPMM and its variants (i
[Table]
e., SPMM MAX , SPMM MEAN
[Table]
and SPMM MIN ) are implemented based on those provided in Pytorch Sparse 3 , with extra supporting for compressing activation maps by quantization and random projection.
For BatchNorm layers, we found that quantizating its saved activation maps only impact model performance a little, which is experimentally veriﬁed in the experiments.
However, we experimentally found that randomly projecting its saved activation maps will lead to divergence.
We note that this observation is consistent with previous ﬁnding that BatchNorm is very sensitive to noise (Micikevicius et al., 2017).
Hence, EXACT only apply the quantization to BatchNorm layers.
Regarding ReLU operations, we have y = ReLU ( x )
= x 1 x > 0 and ∇ y =
∇ y 1 x > 0 .
Hence,
ReLU operations only need to store 1 x > 0 in the context for the backward pass, which takes a single bit per element to store, without introducing any errors.
Since standard deep learning framework only support down to INT8 precision, here we convert the mask matrix into bit-stream and fuse this process into the CUDA kernel of ReLU.forward to minimize the time overhead.
Regarding Dropout operations, let p be the dropout probability.
During the training process, we have
[Equation]
) where g ∼ Bernoulli (1 − p ) is a binary vector sharing the same shape as x . 11 − p is the normalization factor such that E [ y ] = x .
Similarly, Dropout Operations also only need to store g in the context for the backward pass, which takes a single bit per element to store, without introducing any errors.
Similarly, we convert g into bit-stream and fuse this process into the CUDA kernel of
Dropout.forward to minimize the time overhead..
To be clear, “ﬁxed 1 bit” in Table 8 means the activation maps of ReLU and Dropout operations take only 1 bit per element to store.
And the bit-width of all other operations can be adjusted.
F. I MPLEMENTATION D ETAILS
We provide simple API for converting modules in Pytorch Geometric and Pytorch to its corresponding version in EXACT.
For example, replacing torch geometric.nn.
GCNConv with exact.
GCNConv and replacing torch.dropout with exact.dropout .
Currently, EXACT only support Pytorch Geometric and Pytorch.
In future we will try to integrate EXACT with other popular graph learning packages, such as DGL.
Following ActNN (Chen et al., 2021a), to obtain the highest compression ratio, the quantization range r ( l ) v and the zero point Z ( l ) v are stored in the bﬂoat16 data type 4 .
The quantization and dequantization modules are both implemented using CUDA kernels.
We note that Pytorch only support data types down to INT8.
To obtain highest compression ratio, the quantized data is compressed into bit streams such that it can be decoded later during the dequantization process.
All CUDA kernels in EXACT support both full-precision and half-precision (i.e., bﬂoat16 and ﬂoat16).
Thus, EXACT can also be integrated with the automated mixed precision training, i.e.,
AMP 5 , to further decrease the memory consumption.
G E XPERIMENT S ETTINGS G. D ATASETS , F RAMEWORKS , AND H ARDWARES
We give the detailed statistics and the download URLs for all datasets used in our experiments in
Table 10.
We follow the standard data splits and all datasets are directly downloaded from Pytorch
Geometric or the protocol of OGB (Hu et al., 2020).
We implement all models based on Pytorch and Pytorch Geometric.
Almost all experiments are done on a single NVIDIA GeForce RTXwith 24GB GPU memory.
During our experiments, we found that the version of Pytorch, Pytorch
Sparse, and Pytorch Scatter can signiﬁcantly impact the running speed of the baseline.
Here we list the details of our used packages in all experiments in Table 9.
3 https://github.com/rusty1s/pytorch sparse/blob/master/torch sparse/matmul.
[Table]
[Table]
G. M ODEL H YPERPARAMETER C ONFIGURATIONS OF T ABLE 1 AND T ABLE
[Table]
We adopt three popular GNNs.
Namely, GCN (Kipf & Welling, 2017), GraphSAGE (Hamilton et al., 2017)
[Table]
and GAT (Veli ˇ ckovi ´ c et al., 2017) in Table 1 and Table 2.
We follow the hyperparameter conﬁgurations and codebases provided on the OGB (Hu et al., 2020) leader-board.
Speciﬁcally, the hyperparameter conﬁguration is given in Table 11.
The optimizer is Adam (Kingma & Ba, 2014) for GCN and GraphSAGE, while the optimizer is RMSprop for GAT.
All methods terminate after a ﬁxed number of epochs.
We report the test accuracy associated with the highest validation score.
The “Gradient Clipping” in Table 11 indicate the maximum norm for gradients.
“Gradient
Clipping = 0 . 0 ” means we do not clip the gradients in that experiment.
G. H YPERPARAMETER C ONFIGURATIONS OF EXACT (RP+INT2)
The EXACT (RP+INT2) has only one hyperparameter, namely, D R .
The D R conﬁguration of EXACT (RP+INT2) can be found in Table 12.
G. M ODEL H YPERPARAMETER C ONFIGURATIONS IN S ECTION
The optimizer used in all experiments in Section 5 is Adam (Kingma & Ba, 2014).
We use the default hyperparameters for Adam optimizer, except for the learning rate.
All methods terminate after a ﬁxed number of epochs.
We report the test accuracy/F1-micro associated with the highest validation score.
Regarding Reddit, Flickr, and Yelp dataset, we follow the hyperparameter conﬁgurations reported in the respective papers as closely as possible.
We clips the gradient during training.
The “Gradient Clipping” in below tables indicate the maximum norm for gradients.
“Gradient Clipping = 0 . 0 ” means we do not clip the gradients in that experiment.
Reddit
Flickr
Yelp ogbn
[Table]
arxiv ogbn
[Table]
products
[Table]
GraphSAINT
GCNGraphSAGE
GCNIIRegarding ogbn-arxiv and ogbn-products dataset, we follow the hyperparameter conﬁgurations and codebases provided on the OGB (Hu et al., 2020) leader-board.
Please refer to the OGB website for more details.
Table 13 and Table 14 summarize the hyperparameter conﬁguration of Cluster-GCN and GraphSAINT, respectively.
Table 15, Table 16, and Table 17 summarize the hyperparameter conﬁguration of full-Batch GCN, full-Batch GraphSAGE, and full-batch GCNII, respectively.
[Table] [Table]
[Table] [Table]
[Table]
H A NALYZING THE M EMORY U SAGE H. A NALYZING THE M EMORY U SAGE
We use torch.cuda.memory allocated for the memory measurement.
As we mentioned in Table 3, “Act Mem.”
is the memory occupied by activation maps.
Besides activation maps, the model, optimizer, input data, weight gradients, and activation gradients also occupy GPU memory.
We will analyze each of them below.
First , the memory occupied by the model, optimizer, and weight gradients is negligible because the number of parameters in most of GNNs is very small.
Second , the memory occupied by the input data depends on the graph size and often cannot be compressed.
This part can take up a lot of memory when the graph is large.
Third , the memory occupied by activation gradients is dynamic and hard to estimate, since these tensors are temporarily stored in GPUs.
For standard deep learning library, they will be deleted as soon as their reference counting becomes zero (Abadi et al., 2016; Paszke et al., 2019).
EXACT cannot compress activation gradients.
However, the memory occupied by activation gradients can be compressed using AMP, because activation gradients are stored in ﬂoat16 data type under AMP.
As illustrated in Appendix
F.2, we note that our EXACT framework can be integrated with AMP.
The memory usage of the model weights plus optimizer is about 2MB.
Since activation gradients are dynamic and hard to analyze, here we report the peak memory usage during the backward pass, which encompass the activation gradients and other intermediate variables.
For the memory usage of the input data, activation maps, and the peak memory usage during the backward pass, we provided a detailed analysis in Table 18.
We can observe that the memory usage is mainly occupied by activation maps.
For GCNII, the memory usage is dominated by the activation maps.
This is because we need to store all layers’ activation maps during the forward pass.
In contrast, during the backward pass, there is usually only one layer’ activation gradients are kept in memory.
[Table] [Table]
H. O VERALL C OMPRESSION R ATIO
The overall memory compression ratio is shown in Table 19.
We observe that (1) for shallow GNN models, the overall memory compression ratio ranges from 1 .
5 × to 4 × .
(2) for GCNII, the overall memory compression ratio ranges from 4 × to 18 × .
We note that we store the graph structure data (e.g., node ID and edge ID) in torch.
Long data type, which can safely cast to torch.
Int data type to save the memory.
Our implementation supports for using torch.
Int as the data type for the graph structure data.
However, for a fair comparison, we do not utilize this feature.
[Table]
[Figure]
[Table]
H.
A NALYZING THE M EMORY C OMPRESSION R ATIO
Below we analyze the compress ratio in Table 3.
We take a three-layer, 128-dimensional GCNs trained on ogbn-arxiv for example.
The computational graph of the baseline in Table 3 is:
Total bits =
( MM ) + 32( SPMM ) + 32( BN ) + 0( ReLU ) + 8( Dropout )
+ (the ﬁrst layer) 32( MM ) + 32( SPMM ) + 32( BN ) + 0( ReLU ) + 8( Dropout ) + (the second layer) 32( MM ) + 32( SPMM ) (the third layer)
Hence the baseline costs totally 240 bit per element.
The ﬁrst MM in the ﬁrst layer does not cost extra bits because its activation map is exactly the input feature matrix X , which has been stored in GPU memory (recall that we need to ﬁrst move the input data to GPU memory before training).
Pytorch will save this part of memory via the “pass by reference” mechanism.
The ofﬁcial ReLU operation in Pytorch does not need the extra space for saving activation maps (they can reuse the activation maps saved by the previous layer via passing by reference).
Regarding Dropout operation, Pytorch stores the mask matrix using UINT8 data type, which costs 8 bit per element.
The computational graph of “EXACT (INT2)” in Table 3 is:
Total bits =2 .
25( MM ) + 2 .
25( SPMM ) + 2 .
25( BN ) + 1( ReLU ) + 1( Dropout ) + (the ﬁrst layer) 2 .
25( MM ) + 2 .
25( SPMM ) + 2 .
25( BN ) + 1( ReLU ) + 1( Dropout ) + (the second layer) 2 .
25( MM ) + 2 .
25( SPMM ) (the third layer)
The 2.25 bit of MM , SPMM , and BN is from 2 (quantized activation maps) + 0.125 (the zero point tensor)
+
0.125 (the range tensor), where 0 .
125 =
(the bit-width of bﬂoat16) 128 (the dimension is
) (See Appendix F.2 for details).
Also for EXACT, we cannot leverage the “pass by reference” mechanism to save memory for ReLU since the exact activation maps of previous operation is replaced by the compressed one.
Hence in EXACT, the 1-bit masks of ReLU and Dropout is converted into bit-streams using CUDA kernels, which cost 1 bit per element to store.
Hence the “EXACT (INT2)” costs totally 22 bit per element.
And the theoretical compression ratio is 24022 = 10 . 9 .
The empirical compression ratio in Table 3 may have a small gap with the theoretical one.
This is because the attributes of the sparse tensor (adjacency matrix) in pytorch sparse is lazily initialized, i.e., it may be generated during the forward pass and be account for the memory of activation maps.
The computational graph of “EXACT (RP+INT2)” with DR = 8 (see Table 12) in Table 3 is:
Total bits =0 .
28( MM ) + 0 .
28( SPMM ) + 2 .
25( BN ) + 1( ReLU ) + 1( Dropout ) + (the ﬁrst layer) 0 .
28( MM ) + 0 .
28( SPMM ) + 2 .
25( BN ) + 1( ReLU ) + 1( Dropout ) + (the second layer) 0 .
28( MM ) + 0 .
28( SPMM ) (the third layer)
The 0.28 bit of MM and SPMM is from 28 + 16 8 × 128 + 16 8 × 128 = 0 .
We note that EXACT does not apply random projection for BatchNorm layers (see Appendix F.2).
Hence the “EXACT
(RP+INT2)” costs totally 10.18 bit per element.
And the theoretical compression ratio is 24010 .
18 = 23 .
Again, there exists gap between the empirical compression ratio and the theoretical one because the lazy initialization mechanism, the existence of temporary tensors, and the neglect of the memory usage of these random projection matrices.
We emphasize that the compression ratio also depends on the number of input features of the dataset.
(because the mentioned “input feature matrix X is passed by reference” mechanism).
For the above example, the number of input features and the hidden dimension are both 128 and hence it is easy to be analyzed.
Hence, the theoretical compression ratio may vary for different datasets.
I A DDITIONAL E XPERIMENT R ESULTS I. M ORE RESULTS ON A CCURACY AGAINST THE P RECISION AND DR
To support the claim that “the performance of EXACT is mainly determined by the DR ratio of random projection”, we present more results on the test accuracy against the precision and DR ratio here.
In Figure 5, we show the results of two models trained with EXACT using full-batch data on the ogbn-arxiv dataset.
In Figure 6, we show the results of two models trained with EXACT using mini-batch data on Yelp dataset.
I.
C OMPARISON AGAINST SWAPPING AND GRADIENT CHECKPOINTING
We also compare EXACT against the naive swapping and the naive gradient checkpointing.
Both swapping and gradient checkpointing are lossless compression, so they do not have any accuracy drop.
Below we present and discuss about their trade-off among the space and speed.
For swapping, we simply ofﬂoad all activation maps to CPU memory.
Thus, it can achieve the highest compression ratio for activation maps.
However, its running time overhead is roughly 50% - 80% , which is often unacceptable.
For gradient checkpointing, we utilize torch.utils.checkpoint to insert checkpoints at each
GNN layer.
For the time overhead, as shown in Figure 3, it is comparable to EXACT.
We present the memory usage of activations in Table 20.
In summary, its time overhead is comparable to EXACT, however, the memory compression ratio of activations is 1 .
7 ∼ 2 .
3 × , which is not large enough.
Thus, we still cannot train the GCNII on Yelp dataset using a RTX 3090 (24GB) GPU.
[Table]
I. T HE T RAINING C URVES OF EXACT
From the optimization theory, common convergence speed bounds (measured by the number of iterations) in stochastic optimization improve with smaller gradient variance (Bottou et al., 2018).
EXACT essentially trades the gradient variance in return for reduced memory.
Here we experimentally examine how EXACT affects the convergence speed.
Figure 7 shows the training curves of
GNNs trained with EXACT using different conﬁgurations on Reddit dataset.
We make two main observations.
First, EXACT with smaller DR typically converges faster per iteration.
This is consistent with the mentioned optimization theory.
Second, when increasing the DR ratio, the difference in the convergence speed is very small, where the convergence speed is measured by the gap in validation accuracy between consecutive epochs.
In practical scenarios, EXACT hence only have limited impact on the model performance.
I. H YPERPARAMETER S ENSITIVITY E XPERIMENT
As we analyzed in the main body, INT2 is a suitable precision and we only vary the DR ratio for
EXACT.
In this section, we investigate the sensitivity of EXACT (RP+INT2) to the DR ratio.
Table 12 shows the conﬁguration of EXACT (RP+INT2) in the experiment in the main body.
Here we present a comprehensive hyperparameter sensitivity study for EXACT.
Speciﬁcally, the sensitivity studies of EXACT with GraphSAINT, ClusterGCN, GCN, GraphSAGE, and GCNII are shown in
Figure 8, Figure 9, Figure 10, Figure 11, and Figure 12, respectively.
Here we summarize some key observations.
First, the model performance drop generally increases with the DR ratio.
Second, when
DR = 8 , the loss in accuracy of EXACT (RP+INT2) is below 0 .
5% on two third of experiments.
To be concrete, in Table 3, we totally adopt 22 combinations of different datasets and models.
As shown in Table 3 and Table 8, for 15 of the 22 experiments, the loss of accuracy is below or near 0 .
5% when DR = 8 .
Third, when DR = 8 , the loss in accuracy is below 1% in almost all experiments, excepting for two experiments.
Namely, ClusterGCN with ogbn-products and GraphSAINT with ogbn-products .
I. C OMPARISON BETWEEN EXACT AND S AMPLING M ETHODS I.5. C OMPARE EXACT TO S AMPLING M ETHODS UNDER A F IXED M EMORY B UDGET
As examined in Table 3, EXACT and subgraph sampling methods are orthogonal and can be applied over each others.
Here we present an ablation study of comparing EXACT to subgraph sampling methods in a standalone way.
We note that EXACT only focus on saving the memory for activations.
In contrast, subgraph sampling methods can simultaneously reduce the memory of the input data, activations, and activation gradients since they directly reduce the number of node embeddings retained in the memory.
Hence compared to subgraph sampling methods, one limitation of EXACT is that the overall memory saving ratio of EXACT often cannot surpass that of sampling methods.
Under the full-batch setting, one way to further improve the overall memory saving ratio is to incorporate the graph compression methods with EXACT, such as the graph sparsiﬁcation (Spielman & Srivastava, 2011) and graph coarsening (Cai et al., 2021a).
However, it is beyond the scope of this paper.
For a fair comparison, we control the batch size of subgraph subgraph sampling methods such that their activation memory usage equals to that of EXACT (INT2).
For both EXACT(INT2) and EXACT(RP+INT2), we quote the “GraphSAGE” results from Table 3.
For GraphSAINT and
Cluster-GCN, we tune their batch size such that they have the same activation memory usage as
EXACT (INT2).
The results are shown in Table 21.
We make two main observations.
First, for
Reddit and Flickr, full batch training with EXACT outperforms the two sampling methods.
For
Yelp, subgraph sampling methods are better than full batch training with EXACT.
However, we note that this gap is not from EXACT.
As shown in Table 3, without EXACT, the F1-micro of full batch
GraphSAGE is 62 .
05% , which is much lower than that of subgraph sampling methods.
Also, from
Table 6, one interesting observation is that for Yelp, a much smaller batch size ( ≈ 500 ) may even further improve the F1-micro of GraphSAINT from 63 .
20% to 64 .
In summary, subgraph sampling methods may outperform the full batch training on some datasets (e.g., Yelp).
In this case, we think full batch training with EXACT cannot outperform subgraph sampling methods.
If this is not the case, as shown in Table 21, full batch training with EXACT can outperform subgraph sampling methods under a ﬁxed activation memory usage.
[Table]
I.5. E NLARGING T HE B ATCH S IZE OF S AMPLING M ETHODS WITH EXACT
Here we present a case study of scaling up the batch size of GraphSAINT with EXACT on the ogbn-products dataset.
As shown in Table 14, the batch size of GraphSAINT (number of nodes in the sampled subgraphs) is controlled by the “Walk length” and “Roots”
[Table]
and roughly equals “Walk length” × “Roots” (Zeng et al., 2020).
For ogbn-products , we tripled the batch size by changing “Roots” from 20 , 000 to 30 , 000 .
All other hyperparameters are left unchanged.
The results are shown in Table 22.
We observe that EXACT may improve the accuracy when using larger batch size.
We note that EXACT can scale up the batch size to more than 3 × larger.
However, we found that if we further scale up the batch size to 4 × , there is an accuracy drop compared to the baseline
[Table]
with the original batch size, regardless applying EXACT or not.
This is consistent with previous ﬁnding that a larger batch size does not always lead to better performance for ogbn-products (Zeng et al., 2020; Hu et al., 2020).
We quote the sentences from the OGB paper (Hu et al., 2020) to explain this counter-intuitive observation.
“The recent mini-batch-based GNNs give promising results, even slightly outperforming the full-batch version of GraphSAGE that does not ﬁt into ordinary GPU memory.
The improved performance can be attributed to the regularization effects of mini-batch noise and edge dropout.”
(Hu et al., 2020).
In summary, in the ogbn-products ablation study, we utilize EXACT to triple the batch size of
GraphSAINT, which may even improve the accuracy over the original one.
Moreover, from our experiments, the 3 × larger batch size performs the best on ogbn-products .
Thus, for some datasets, there exist an optimal batch size for subgraph sampling methods.
Since this optimal batch size may be beyond the capacity of the hardware, the meaning of EXACT is to enlarge the “search space” for ﬁnding this optimal batch size.