Tango: rethinking quantization for graph neural network training on GPUs
Graph Neural Networks (GNNs) are becoming increasingly popular due to their
superior performance in critical graph-related tasks. While quantization is
widely used to accelerate GNN computation, quantized training faces
unprecedented challenges. Current quantized GNN training systems often have
longer training times than their full-precision counterparts for two reasons:
(i) addressing the accuracy challenge leads to excessive overhead, and (ii) the
optimization potential exposed by quantization is not adequately leveraged.
This paper introduces Tango which re-thinks quantization challenges and
opportunities for graph neural network training on GPUs with three
contributions: Firstly, we introduce efficient rules to maintain accuracy
during quantized GNN training. Secondly, we design and implement
quantization-aware primitives and inter-primitive optimizations that can speed
up GNN training. Finally, we integrate Tango with the popular Deep Graph
Library (DGL) system and demonstrate its superior performance over
state-of-the-art approaches on various GNN models and datasets.
Shiyang Chen, Da Zheng, Caiwen Ding, Chengying Huan, Yuede Ji, Hang Liu
# Tango: rethinking quantization for

###### Abstract.

Graph learning is becoming increasingly popular due to its superior performance in tackling many grand challenges. While quantization is widely used to accelerate Graph Neural Network (GNN) computation, quantized training faces remarkable roadblocks. Current quantized GNN training systems often experience longer training time than their full-precision counterparts for two reasons: (i) addressing the quantization accuracy challenge leads to excessive overhead, and (ii) the optimization potential exposed by quantization is not adequately leveraged. This paper introduces Tango which re-thinks quantization _challenges_ and _opportunities_ for graph neural network training on GPUs with three contributions: Firstly, we introduce efficient rules to maintain accuracy during quantized GNN training. Secondly, we design and implement quantization-aware primitives and inter-primitive optimizations to speed up GNN training. Finally, we integrate Tango with the popular Deep Graph Library (DGL) system and demonstrate its superior performance over the state-of-the-art approaches on various GNN models and datasets.

+
Footnote †: journal: Computer Vision and Pattern Recognition

+
Footnote †: journal: Computer Vision and Pattern Recognition
data format nor floating-point format with dynamically adjusted exponent and mantissa. The model trained using SWALP or clipped gradients needs more epochs to converge because the weights are updated less frequently. _(ii) The optimization potential offered by quantization is not well-utilized._ ActNN (Wang et al., 2017), TinyKG (Wang et al., 2017), and EXACT (Wang et al., 2017) quantize the tensors to save memory and dequantize them back to full-precision for computation, increasing the overall training time (Wang et al., 2017; Wang et al., 2017). For example, TinyKG with 8-bit quantization is 54.1% slower than using FP32. Of note, Degree-Quant (Zheng et al., 2019) performs Quantization-Aware Training (QAT) (Wang et al., 2017; Wang et al., 2017; Wang et al., 2017; Wang et al., 2018; Wang et al., 2018; Wang et al., 2019) that uses full-precision to "simulate quantization" in training to reduce the error for quantized inference, which, again, experiences longer training time.

This paper introduces Tango, the first GPU-based quantized GNN training system that both _maintains the model accuracy_ and _reduces turnaround time_ when compared to the full precision counterpart. Particularly, Tango encompasses three contributions:

* We introduce several lightweight rules to maintain accuracy for quantized GNN training. The rules include GPU-accelerated stochastic rounding, derivation of proper quantization bit count, novel quantization-aware GEMM, SPMM, and SDDMM, and full precision weight update and softmax.
* We design and implement a quantization-aware system to reduce the GNN training time on GPUs. Our techniques include GEMM with on-the-fly quantization, incidence-matrix-based adaptive SPMM, SDDMM with on-the-fly dequantization, and inter-primitive optimizations.
* For ease of use, we integrate Tango with DGL, which uses PyTorch as the backend. Therefore, all existing DGL-based models can enjoy the performance benefits from Tango without any changes. We demonstrate that Tango constantly outperforms state of the art for all evaluated GNN models while maintaining the training accuracy.

The remainder of this paper is organized as follows: Section 2 presents the background. Section 3 discusses the design and techniques in Tango. Specifically, Section 3.1 describes the challenges and opportunities of Tango, Section 3.2 illustrates the lightweight rules for maintaining training accuracy during quantization, and Section 3.3 presents the systematic effort on quantization accelerated training. We evaluate Tango in Section 4, describe related work in Section 5, and conclude in Section 6.

## 2. Background

### A running example for GAT training

This section uses a running example to illustrate how to express the forward and backward computations of GAT with three key primitives, i.e., GEMM, SPMM, and SDDMM.

**Primitives for forward computation.** Figure 1a presents the forward workflow of GAT on a toy graph, i.e., node projection (), attention computation (), and message aggregation ().

In step 1, GAT resorts to _GEMM primitive_ to perform a linear transformation for node features, i.e., \(\mathbf{H}^{\prime}=\mathbf{H}^{(l-1)}\cdot\mathbf{W}\). In \(\mathbf{H}^{(l-1)}\), where each row of \(\mathbf{H}^{(l-1)}\), in the same color, represents the features of one node. Each node feature contains two heads in \(\mathbf{H}^{(l-1)}\). Using the first row of \(\mathbf{H}^{(l-1)}\) as an example, [0.01, 0.40] is the first head, and the remaining two values belong to the second head. \(\mathbf{W}\) is the learnable weight matrix for linear transformation.

In step 2, GAT consolidates each head of the feature vector into one scalar by GEMM, i.e., \(\mathbf{S}=(\mathbf{H}^{\prime}\cdot\mathbf{a}_{src})^{T}\) and \(\mathbf{D}=(\mathbf{H}^{\prime}\cdot\mathbf{a}_{dst})^{T}\). Using node \(\mathbf{v}_{0}\) of 1 as an example, for the source feature, [0.59, 0.73]\(\times\)[0.91, 0.90]\({}^{T}\) = 1.20, and [0.51, -0.65]\(\times\)[0.42, 0.62] = -0.19. Similarly, we can derive the entire S and D.

In step 3, the source (**S**) and destination (**D**) node feature matrices are combined by an _SDDMM primitive_ to arrive at the edge feature \(\mathbf{E}\). Formally, the SDDMM is defined as \(\mathbf{E}=\mathbf{G}\odot(\mathbf{S}\oplus\mathbf{D}^{T})\), where every row of \(\mathbf{S}\) computes against every row of \(\mathbf{D}\) with the customized operation 1, and \(\odot\) is a Hadamard product operator. The resultant matrix \(\mathbf{E}\) is masked out by the sparse adjacency matrix \(\mathbf{G}\) of the graph so \(\mathbf{E}[i][j]=0\) when there is no edge between \(v_{i}\) and \(v_{j}\). In Figure 1a, \(\oplus\) denotes addition. Using edge \(e_{3}\) as an example, since it connects source \(v_{0}\) and destination \(v_{3}\), we arrive at [1.20, -0.19] + [0.20, 0.05] = [1.40, -0.14] for \(e_{3}\). Then an element-wise LeakyReLU is applied to the edge features. Particularly, each non-negative entry in \(\mathbf{E}\) is unchanged while negative ones become close to 0, which we use 0 to represent. Hence [1.40, -0.14] becomes [1.40, 0.00] in Figure 1a.

In step 1, edges of the same destination come together to compute the head-wise attention scores through softmax operation.

Figure 1. GAT training on a toy graph, i.e., middle left of (a), with two heads. This example is used throughout this paper.
Using node \(v_{3}\) as an example, it has \(e_{3}\) and \(e_{4}\) as the incoming edges. Therefore, the attention scores of \(e_{3}\) and \(e_{4}\) are computed as \(0.63=\frac{e^{\bot 0}}{e^{\bot 0.64\bot 0.85}}\), and \(0.46=\frac{e^{0.46}}{e^{\bot 0.46\bot 16}}\) for \(e_{3}\), and \(0.37=\frac{e^{0.86}}{e^{\bot 0.86}}\) and \(0.54=\frac{e^{0.46}}{e^{\bot 0.46\bot 16}}\) for \(e_{4}\). Putting this example into a general formula, we use SPMM and SDDMM operations together to compute the denominator as follows: First, we use _SPMM_ to aggregate the in-edges for every node such as \(\mathbf{M}^{\prime}=(\mathbf{G}\odot exp(\mathbf{E}))\cdot\mathbf{1}\). Of note, \(\mathbf{1}\) is an all '1' dense matrix. Second, since the first step computed the denominator for each destination vertex, we use _SDDMM_ to assign this denominator back to each incoming edge via \(\mathbf{E}^{\prime}=\mathbf{G}\odot(\mathbf{1}\cdot\mathbf{M}^{\prime T})\). Subsequently, \(\boldsymbol{\alpha}=\frac{exp(\mathbf{E})}{\mathbf{E}}\).

In step 1, GAT performs an _SPMM_ to derive the new node embedding via \(\mathbf{H}^{(I)}=(\mathbf{G}\odot\boldsymbol{\alpha})\cdot\mathbf{H}^{\prime}\). Intuitively, this step derives the new embedding by computing the weighted sum of all the incoming neighbors to a destination vertex. Using \(v_{3}\) as an example, its incoming neighbors are \([v_{0},v_{2}]\). We arrive at \(\mathbf{H}^{(I)}[v_{3}]=\boldsymbol{\alpha}[e_{5}]\cdot\mathbf{H}^{\prime}[v _{0}]+\boldsymbol{\alpha}[e_{4}]\cdot\mathbf{H}^{\prime}[v_{2}]\), resulting in [0.49, 0.61, 0.77, -0.58].

**Primitives for backward computation.** Figure 0(b) is the backward pass of Figure 0(a). Steps 1 (_SPMM_) and 1 (_SDDMM_) of Figure 0(b) are the corresponding backward operations for step 1 in Figure 0(a). In the forward SPMM operation (1), \(\mathbf{H}^{(I)}=(\mathbf{G}\odot\boldsymbol{\alpha})\cdot\mathbf{H}^{\prime}\), we hence disperse the gradients of \(\mathbf{H}^{(I)}\) to both \(\mathbf{H}^{\prime}\) and \(\sigma\). First, we arrive at \(\partial\mathbf{H}^{\prime}=(\mathbf{G}^{T}\odot\boldsymbol{\alpha})\cdot \partial\mathbf{H}^{(I)}\), which is an _SPMM_ (1) on the reversed graph since the updated node feature \(\mathbf{H}^{(I)}\) is aggregated from the source node feature. Using \(\mathbf{H}[v_{1}]\) as an example, it receives gradients from both \(e_{0}\) and \(e_{2}\). Therefore, we arrive at \(\partial\mathbf{H}^{(I-1)}[v_{1}]=\boldsymbol{\alpha}[e_{0}]\cdot\partial \mathbf{H}^{(I)}[v_{0}]+\boldsymbol{\alpha}[e_{2}]\cdot\partial\mathbf{H}^{( I)}[v_{2}]\). That is, [1.56,1.57,-0.19,0.49]\(=1.0\times\)[0.54, 0.51] \(+1.0\times\)[1.02,1.06] \(\|_{concat}\)\(1.0\times\)[-0.26, -0.07] \(+1.0\times\)[0.07,0.56]. Second, we have \(\partial\boldsymbol{\alpha}=\mathbf{G}\odot(\partial\mathbf{H}^{(I)}\cdot \mathbf{H}^{T})\), which is an SDDMM operator on the original graph, where it performs row-wise dot-product (1). Using \(\partial\boldsymbol{\alpha}[e_{0}]\) as an example, it connects nodes \(v_{1}\) and \(v_{0}\), we arrive at \(\partial\boldsymbol{\alpha}[e_{0}]=\partial\mathbf{H}^{(I)}[v_{0}]\cdot \mathbf{H}^{\prime}[v_{1}]\). In the example, we get [0.78,-0.13] \(=\)[0.54,0.51] \(\times\)[0.76,0.73] \({}^{T}\|_{concat}\)[-0.26,-0.07]\(\times\)[0.79,-1.07]\({}^{T}\).

Step 1 computes the gradient of edge features using attention scores. Using \(e_{3}\) as an example, we compute \(\partial\mathbf{E}[e_{3}]=\boldsymbol{\alpha}[e_{5}](\partial\boldsymbol{ \alpha}[e_{5}]-(\partial\boldsymbol{\alpha}[e_{5}]\boldsymbol{\alpha}[e_{5}] \boldsymbol{\alpha}[e_{5}]+\partial\boldsymbol{\alpha}[e_{4}]\boldsymbol{ \alpha}[e_{4}])\) based on the derivative of softmax operation. We first use _SPMM_ to aggregate the incoming edge features for every node \(\mathbf{P}=(\mathbf{G}\odot\boldsymbol{\alpha}\odot\boldsymbol{\alpha})\cdot \mathbf{1}\). For example, \(\partial\boldsymbol{\alpha}[e_{5}]\boldsymbol{\alpha}[e_{5}]+\partial \boldsymbol{\alpha}[e_{4}]\boldsymbol{\alpha}[e_{4}]\) is the aggregation on \(v_{3}\) as \(0.80\times 0.63+0.45\times 0.37=0.67\) for the first head. Then we compute the final gradient for every edge with _SDDMM_, \(\partial\mathbf{E}=\boldsymbol{\alpha}\odot(\partial\boldsymbol{\alpha}-( \mathbf{G}^{T}\odot(\mathbf{P}\cdot\mathbf{1}^{T})))\). That is, every node feature \(\mathbf{P}\) is assigned to their out-edges in the reversed graph, and then computed with \(\partial\boldsymbol{\alpha}\) and \(\boldsymbol{\alpha}\). The gradient of the first head of \(e_{3}\) is \(0.63\times(0.80\cdot 0.67)=0.08\).

In step 1 and 1, the gradient of edge attention score is used to compute the source feature and destination feature with two _SPMM_ operations, \(\partial\mathbf{S}=(\mathbf{G}^{T}\odot\mathbf{\partial E})\cdot 1\) and \(\partial\mathbf{D}=(\mathbf{G}\odot\mathbf{\partial E})\cdot\mathbf{1}\), where nodes aggregate their out-edge and in-edge attention scores, respectively. Still use \(v_{3}\) as an example, its gradient \(\partial\mathbf{S}[v_{3}]=\partial\mathbf{E}[e_{1}]=[0,0]\) and \(\partial\mathbf{D}[v_{3}]=\partial\mathbf{E}[e_{3}]+\partial\mathbf{E}[e_{4}]= [0,0.15]\). The gradients from multiple out-edges are accumulated.

Note steps 1 and 1, which do not depend on the graph structure, will follow the traditional DNN back propagation method for gradient computation. We skip the details.

### GNN models

There exist a variety of GNN models. Graph Convolutional Network (GCN) (Gai et al., 2017) derives a graph convolutional operator through spectral graph theory. It can be expressed by GEMM and SPMM operations. Later, GraphSAGE (Gai et al., 2017) uses sampling to encode the graph topology for inductive learning. The model is applicable for unseen nodes because it learns the feature from sampled sub-graph. _GraphSAGE can be implemented with GEMM and SPMM_. GAT (Gai et al., 2018) further introduces graph attention mechanisms that can attend to various neighbors with weights. This model contains GEMM, SPMM, and SDDMM primitives. Later, Relational GCN (RGCN) extends GCN via assigning different parameters to edges with different types (Shi et al., 2018; Wang et al., 2018). _Here, RGCN consists of GEMM and SPMM primitives._ HGT proposes a transformer-based GNN model for heterogeneous graphs (Wang et al., 2018). It contains different parameters for distinct edge and node types. _This model includes GEMM, SDDMM, and SPMM primitives_.

We study two GNN models, i.e., GCN and GAT, for two reasons: (i) These two models are the most popular and cover all the required primitives for most GNN models. (ii) These two models contain relatively large and complete training and testing datasets.

### Quantization

For a collection of values \(\mathbf{X}=\{\mathbf{X}_{i}\mid\mathbf{X}_{i}\in[\mathbf{X}_{min},\mathbf{X}_{ max}]\}\) which are represented in full precision, quantization uses fewer number of bits (i.e., \(B\)) to represent each \(\mathbf{X}_{i}\). Quantization scatters \(\mathbf{X}\) into \(2^{B}-1\) buckets. Subsequently, all the \(\mathbf{X}_{i}\)'s in the same bucket are represented as the same value, i.e., the bucket value.

For _uniform_ quantization, we assign each bucket the same value range, that is, \(s=\frac{\alpha-\beta}{2\beta-1}\), where \([\alpha,\beta]\) is the clipping range of \(\mathbf{X}\). There are also _nonuniform quantization_ whose quantized values are not necessarily uniformly spaced. If one wants to include the entire value range of \(\mathbf{X}\), one needs \(\alpha=\mathbf{X}_{min}\), and \(\beta=\mathbf{X}_{max}\). Formally,

\[\mathbf{X}_{i,Quant}=round(\frac{\mathbf{X}_{i}}{s})-Z, \tag{1}\]

where \(Z=\frac{\alpha+\beta}{2}\) is the zero point after quantization. One can recover the original value \(\mathbf{X}_{i}\) by dequantizing \(\mathbf{X}_{i,Quant}\):

\[\mathbf{X}_{i}\approx s\cdot(\mathbf{X}_{i,Quant}+Z). \tag{2}\]

Quantization further includes the following three configurations: (i) _Asymmetric_ vs. _symmetric_ quantization. Particularly, for \(s\), one can let \(-\alpha\)=\(\beta\)=max\((|\mathbf{X}_{max}|,\)\(|\mathbf{X}_{min}|)\). While asymmetric quantization will likely enjoy a more precise clipping range when compared to symmetric quantization, the latter design, however, simplifies the quantization function in Equation 1 as \(Z=\frac{\alpha+\beta}{2}=0\). (ii) _Quantization granularity_ concerns about the size of \(\mathbf{X}\). Using a matrix as an example, we can extract the same \(s\) for the entire matrix or one \(s\) per row/column of a matrix. The latter has a finer granularity than the former. (iii) _Static vs dynamic_ quantization determines whether we change \(s\) for the same tensor \(\mathbf{X}\) from iteration to iteration. The dynamic version does so, while the static one does not. In Tango, we adopt _symmetric, tensor-level granularity_, _dynamic_ quantization to maintain training accuracy and enhance training speed.
## 3. Tango: An Accuracy and Speed

Co-designed Quantization System

### Tango overview

Our key observation is that quantization presents both _challenges_ and _opportunities_ for GNN training. Tango aims to tackle the challenges efficiently while extracting quantization benefits as follows:

**Challenge: Maintaining training accuracy** poses three issues: (i) quantization could introduce additional computation tasks in addition to the steps in Figure 1. We need to reduce the overhead brought by those additional computations. (ii) For various operations in GNN training (in Figure 1), we need to decide what operations should be quantized and how we should quantize the tensors in each operator to meet the training accuracy requirements. In addition, (iii) those rules should expose optimization opportunities for Tango to accelerate the most time-consuming operations in GNN training with quantization.

In this paper, (i) we introduce GPU-friendly stochastic rounding and a lightweight operation to determine the required # of quantization bits, reducing the cost of meeting accuracy requirements. (ii) We determine that weight update and softmax operations should be performed in full precision, while GEMM, SPMM, and SDDMM can be performed in our novel quantization-aware manner. This minimizes the impact on training accuracy while providing critical optimization opportunities for reducing turnaround time. Notably, (iii) GEMM, SPMM, and SDDMM are the most time-consuming phases in GNN, and our quantization-aware design offers optimization opportunities (see below) to reduce computation costs in GEMM and memory costs for SPMM and SDDMM.

**Opportunity: Accelerating training by quantization.** Quantization offers two avenues to improve training speed, that is, higher computation throughput and less memory traffic with values in lower precision. We use quantized computing to accelerate the most computation-intensive primitives and operations, i.e., GEMM, SPMM and SDDMM. _However, the problem is that these primitives are highly optimized and fine-tuned by commercial libraries._ And CUBLAS GEMM and cuSPARSE SPMM, and SDDMM are closed-source. Integrating our proposed optimizations into these kernels and achieving the desired speedup is extremely challenging.

In this paper, (i) we utilize our novel quantization-aware GEMM to reduce computation time. Moreover, we identify an optimal tiling strategy to overlap the on-the-fly quantization of the matrix with the subsequent quantized computations. (ii) To address the memory-intensive nature of SPMM and SDDMM, we sequentially quantize the input tensor and write the quantized value in memory. The computation then randomly accesses the smaller quantized tensor, which provides better cache behavior than direct random access to full-precision tensors.

### Lightweight rules for maintaining training accuracy during quantized training

**GPU-accelerated stochastic rounding.** We adopt stochastic rounding to reduce the quantization error (Wang et al., 2017), with which the expectation of the quantization error should be 0 statistically. In particular, given a scaled floating-point number \(x\) between \([-2^{B-1}-1,+2^{B-1}-1]\) as the range of \(B\)-bit integers, we round it to integer based on:

\[x_{Quant}=\begin{cases}floor(x)+1,&w/probability\quad x-floor(x);\\ floor(x),&w/probability\quad 1-(x-floor(x)).\end{cases} \tag{3}\]

We design and implement a GPU-accelerated pseudo-random number generator to facilitate fast stochastic rounding, which is \(\sim\)20\(\times\) faster than the native cuRAND random number generator on GPU (Wang et al., 2017). Our key optimization is storing random generator states in GPU registers as opposed to in global memory, which is the case in the existing cuRAND library (Wang et al., 2017). Since the random number generator is a memory-bound operation, this optimization helps significantly improves the throughput. Of note, because cuRAND is closed-source, we cannot directly integrate this optimization into cuRAND. We thus implement our generator based upon xoshiro256++(Wang et al., 2017) with our memory optimizations.

**Lightweight rule for deriving # of desired quantization bits.** We develop a metric to measure the quantization error, which subsequently helps derive the # of desired quantization bits. During quantization, a value \(\mathbf{X}_{i}\) will be rounded to one of the quantization grid points \(\mathbf{X}_{i,Quant}\). We introduce the following metric to estimate the quantization error of a tensor \(\mathbf{X}\):

\[Error\mathbf{x}=\frac{1}{N}\sum_{i=1}^{N}\left|\frac{\mathbf{X}_{i}-\mathbf{ X}_{i,Quant}}{\mathbf{X}_{i}+\mathbf{X}_{i,Quant}+\epsilon}\right|, \tag{4}\]

where \(N\) is the number of elements in the tensor.

Intuitively, \(Error\mathbf{x}\) derives the relative quantization error of a tensor \(\mathbf{X}\), where the numerator, i.e., \(\left|\mathbf{X}_{i}-\mathbf{X}_{i,Quant}\right|\) is the absolute quantization error while the denominator is the sum of \(\mathbf{X}_{i}\), \(\mathbf{X}_{i,Quant}\), and \(\epsilon\). The denominator needs the sum of the three values for two reasons: (i) a small \(\epsilon\) to avoid dividing by zero, i.e., when \(\mathbf{X}_{i}=\mathbf{X}_{i}\),\(Quant=0\). Tango chooses \(\epsilon=0.0005\). Of note, Tango does not experience \(\mathbf{X}_{i}+\mathbf{X}_{i,Quant}=0\) when \(\mathbf{X}_{i}\neq 0\) and \(\mathbf{X}_{i,Quant}\neq 0\) because our quantization is symmetric. (ii) If we use \(\epsilon\) with only either \(\mathbf{X}_{i}\) or \(\mathbf{X}_{i,Quant}\) as the denominator, we could suffer from quantization error divided by \(\epsilon\). This would lead to an extremely large relative error for a particular \(\mathbf{X}_{i}\), overshadowing the relative quantization error of other \(\mathbf{X}_{i}\)'s.

Our proposed quantization error metric in Equation 4 is a relative error thus inductive. That is, this parameter could be used to compare the quantization error across tensors. Therefore, we can tune a desired \(Error\mathbf{x}\) that is generally applicable for various tensors. The value range of \(Error\mathbf{x}\) is [0, 1]. Particularly, if \(\mathbf{X}_{i}\) has no rounding error, the corresponding error is 0. When the rounding error of \(\mathbf{X}_{i}\) is significant, the term approaches 1.

We leverage Equation 4 to select the desired number of quantization bits as follows: we compute \(Error\mathbf{x}\) of the output tensor of the first GNN layer with quantization. Note that we do not apply this metric to the input tensor of the first layer because its quantization error can be recovered by learning from the graph structure (Wang et al., 2017). We also want to mention that the training process could potentially amend the quantization error when the bit count is even lower. Our bit count derivation metric derives a lower bound bit count that could maintain the training accuracy.

As shown in Figure 1(a), our heuristic demonstrates that when \(Error\mathbf{x}\)\(\sim\) 0.3, Tango can maintain the accuracy requirement across various datasets. Therefore, we let \(Error\mathbf{x}\) = 0.3 across all datasets.
Figure 1(b) shows that the desired number of bits for ogbn-arxiv, Pubmed, and ogbn-products are 8, 6, and 8, respectively.

The benefit of our metric is as follows: because our lightweight rule calculates the \(Error_{\mathbf{X}}\) solely for the first layer during the initial epoch. In contrast, determining the accuracy loss typically necessitates training the model until convergence (i.e., all epochs). The effectiveness of our approach is demonstrated by our empirical findings presented in Fig 1(a), which shows that \(Error_{\mathbf{X}}<=0.3\) is a general threshold to maintain the accuracy across datasets.

**Novel quantization-aware matrix multiplication with scaling factor computation.** Since the majority of the tensor primitives in GNN are either dense matrix multiplication or a variant of it, the accuracy analysis would be similar across these primitives. We hence restrict our accuracy analysis to dense matrix multiplication (i.e., GEMM) for brevity.

For two reasons, the resultant matrix of a quantized matrix multiplication has to be of higher precision. First, the result of a multiplication operation between two 8-bit integers could go beyond the value range of an 8-bit integer. Second, the subsequent accumulation of the multiplied values can again push the value beyond the range of an 8-bit integer.

Figure 2(b) shows that the desired number of bits for ogbn-arxiv, Pubmed, and ogbn-products are 8, 6, and 8, respectively.

The benefit of our metric is as follows: because our lightweight rule calculates the \(Error_{\mathbf{X}}\) solely for the first layer during the initial epoch. In contrast, determining the accuracy loss typically necessitates training the model until convergence (i.e., all epochs). The effectiveness of our approach is demonstrated by our empirical findings presented in Fig 1(a), which shows that \(Error_{\mathbf{X}}<=0.3\) is a general threshold to maintain the accuracy across datasets.

**Novel quantization-aware matrix multiplication with scaling factor computation.** Since the majority of the tensor primitives in GNN are either dense matrix multiplication or a variant of it, the accuracy analysis would be similar across these primitives. We hence restrict our accuracy analysis to dense matrix multiplication (i.e., GEMM) for brevity.

For two reasons, the resultant matrix of a quantized matrix multiplication has to be of higher precision. First, the result of a multiplication operation between two 8-bit integers could go beyond the value range of an 8-bit integer. Second, the subsequent accumulation of the multiplied values can again push the value beyond the range of an 8-bit integer.

Figure 3 presents this problem when we perform \(\mathbf{H}^{\{l-1\}}\cdot\mathbf{W}\) in quantized mode. After quantization, the first row of \(\mathbf{H}^{\{l-1\}}_{Quant}\), i.e., [1 58 101 28] multiplies with the first column of \(\mathbf{W}_{Quant}\), i.e., [-104 12 85 93]\({}^{T}\) experiences both issues mentioned above. In fact, all the entries in the resultant matrix (\(\mathbf{H}^{\{l-1\}}_{Quant}\cdot\mathbf{W}_{Quant}\))\({}_{int32}\) exceed the 8-bit range of [127, 127]. Therefore, we opt for a 32-bit data format to store the result to avoid this overflow problem. The good news is that storing the results in 32-bit integers introduces negligible overheads on commodity GPUs. Also, note that recent tensor core units on NVIDIA GPUs force the resultant matrix to be a 32-bit integer matrix for the input of two 8-bit integer matrices.

To reduce the quantization overheads, Tango directly dequantizes the GEMM results, i.e., \(\mathbf{H}^{\prime}\) into FP32 after computing the resultant matrix with our optimizations. In the meantime, Tango also derives the scaling factor \(\mathbf{spr}=166.26\) during the quantized GEMM operation, as shown in Figure 3. This design avoids a dedicated dequantization kernel, a scaling factor computation kernel, and the associated expensive global memory accesses.

**Full precision weight update.** To combat the round-off error, we update the model weights with dequantized FP32 gradients. The reason is that the magnitude of the model weight is often significantly larger than the gradients. In addition, the small learning rate further amplifies the difference. Previously, existing projects use shared exponent [27], Flexpoint [25], or delayed updates [50] to tackle the round-off error. Unfortunately, these designs could suffer from shortcomings of delayed convergence, unavailability on commodity GPUs, being slow to implement on GPUs, or multiple of these shortcomings [51, 52]. Of note, although the updated FP32 weights are quantized into 8-bit integers in the next iteration, quantizing the updated weights into 8-bit integers is often better than directly updating the quantized weights with quantized gradients as elaborated below.

Assuming \(W_{full}=W_{quant}+W_{round\_off}\) and \(\Delta W_{full}=\Delta W_{quant}+\Delta W_{round\_off}\), where \(W_{full}\) and \(\Delta W_{full}\) are weights and update values (e.g. gradients) of full precision, respectively. \(W_{quant}\) and \(\Delta W_{quant}\) are the output from the quantization function \(Q\). \(W_{round\_off}\) and \(\Delta W_{round\_off}\) are the correspondingly round-off errors. Below is our analysis:

\[Q(W_{full})+Q(\Delta W_{full})=W_{quant}+\Delta W_{quant}. \tag{5}\]

If we add the full precision before quantization, we arrive at:

\[\begin{split} Q(W_{full}+\Delta W_{full})&\sim W_{ quant}+\Delta W_{quant}\\ &+Q(W_{round\_off}+\Delta W_{round\_off}).\end{split} \tag{6}\]

One can observe that Equation 6 offers higher accuracy than Equation 5 as \(Q(W_{round\_off}+\Delta W_{round\_off})\) curbs the round off error.

**Full precision for the layer before Softmax.** The Softmax layer amplifies the quantization error due to its exponential operations. For simplicity, we consider a layer before Softmax with two outputs \(z_{0}\) and \(z_{1}\). The difference in Softmax score (\(D\)) between the two outputs is:

\[D=\frac{\frac{exp(z_{0})}{exp(z_{1})+exp(z_{1})}}{\frac{exp(z_{1})}{exp(z_{1})} }=\frac{exp(z_{0})}{exp(z_{1})}. \tag{7}\]

Once the quantization error \(e_{i}\) is introduced to \(z_{i}\), the perturbation of output difference follows:

\[D^{\prime}=\frac{exp(z_{0}+e_{0})}{exp(z_{1}+e_{1})}=D\cdot\frac{exp(e_{0})}{exp( e_{1})}=D\cdot\underbrace{exp(e_{0}-e_{1})}_{\text{Amplified error}}. \tag{8}\]

This analysis suggests that the exponential function applied to \((e_{0}-e_{1})\) will rapidly make the difference \(D\) either bigger or smaller, departing from the desired faithful \(D\). Therefore, we propose to use full precision to compute the layer before the Softmax.

Figure 3. Quantization for GEMM of step 1 in Figure 0(a).

Figure 2. The accuracy for different \(Error_{\mathbf{X}}\) and the required number of bits to retain the desired \(Error_{\mathbf{X}}\) for ogbn-arxiv, Pubmed, and ogbn-products datasets.
### Quantization accelerated training

**GEMM with on-the-fly quantization.** Figure 4 illustrates our GEMM with on-the-fly quantization and scaling factor (i.e., s) computation. Our quantized GEMM includes four steps: First, we quantize while loading the tiles of input matrices from global memory to shared memory (i.e., Tiles **A** and **B**). Note that the input matrices are usually needed for backward computation. Therefore, we store the quantized tiles back in global memory while computing, eliminating the round-trip memory latency in the naive design. Second, we store the resultant block in registers to minimize the write latency. Third, during computation, we pack four 8-bit integers into a 32-bit register and use one **DP4A** instruction for four multiply-accumulate operations between two packed registers. Fourth, we dequantize the resultant 32-bit integers in registers to floating-point. We also fuse the computation of parameter \(s\) in the kernel for the following primitives.

We develop a data tiling strategy to hide the data access latency behind the computations. First, when loading from global memory, we choose an appropriate tile size, i.e., \(128\times 32\) which is the ideal tile size to balance the computation capability and memory throughput on V100 GPUs. Below is our analysis: assuming the sizes of Tiles **A** and **B** are \(M\cdot k\) and \(N\cdot k\), to pipeline the loading of Tiles **A** and **B** with the computation, we hide the loading latency by arithmetic operations, \(loadLatency=arithmeticLatency\). We denote that the latency of loading one FP32 value from global memory is \(Latency_{global}\) and the latency of performing one multiply-accumulate operation is \(Latency_{compute}\). We arrive at \(loadLatency=(M\cdot k+N\cdot k)\cdot Latency_{global}\), and \(arithmeticLatency=(M\cdot N\cdot k)\cdot Latency_{compute}\). This leads to \(\frac{M+N}{MN}=\frac{Latency_{compute}}{Latency_{global}}\). On V100 GPU, we find \(Latency_{global}\approx 400\), \(Latency_{compute}\approx 4\). Without loss of generality, we let \(M=N\). We hence arrive at \(M=N\approx 200\). Since the size of \(M\) and \(N\) should be the power of two, we find \(M=N=128\) offer the best performance. Second, at the warp level, we let each warp load two blocks from Tiles **A** and **B** to compute four adjacent blocks in Tile **C** as shown in Figure 4. We derive the optimal block size that can hide the latency of accessing shared memory. Particularly, in each iteration, a thread loads 32 packed INT8 as the eight blocks colored on the right side of Figure 4. Then the 16 **DP4A** instructions cover the 18 cycles latency of loading for the next iteration.

For computation, we carefully schedule the threads to improve the computation intensity. First, when loading and quantizing Tile **A** from global memory to shared memory, we transpose the Tile because the access is in column while Tile **A** is row-major in global memory. Second, to avoid bank conflict, a warp stores Block **A** in shared memory with a 16-byte offset in each column. Each warp works on \(2\times 2\) C blocks to reuse Block **A** and **B**. Third, we schedule the threads as shown in the left side of Figure 4, mapping 32 threads within a warp to block **C** to increase shared memory throughput. For example, threads 0,1,2,3 access the same address in block **A** so the loaded data can be broadcast to 4 threads.

Of note, there exist frameworks that can generate GEMM kernels with efficient tiling and scheduling, but none of them can be used to implement GEMM with _on-the-fly quantization_. On the one hand, template-based frameworks, such as AutoTVM (Xu et al., 2018) and Ansor (Xu et al., 2018), optimize kernels by enumerating the combinatorial choices of optimizations (e.g., tile layout, tile size, and parallelization). Searching the design space is time-consuming, and the generated kernels are not guaranteed to be optimal. On the other hand, on-the-fly quantization is not supported by the existing templates. Moreover, the kernel generated by these frameworks, e.g., triton compiler (Xu et al., 2018), uses too many registers, resulting in unsatisfied performance.

**Incidence matrix-based adaptive SPMM.** Tango performs quantization in a separate kernel for SPMM. We use quantization to reduce the memory traffic for SPMM since quantization leads the node and edge feature matrices to a smaller size. Unlike GEMM, which performs sequential memory access for the input matrices, SPMM experiences random memory accesses. In this case, performing on-the-fly quantization would lead to random memory access for input matrices of 32-bit floating-point data type. Further, because of unpredictable access patterns, on-the-fly quantization could potentially lead to repeated quantization of the same data. Instead, a dedicated quantization kernel would read 32-bit input floating-point matrices sequentially once and write the 8-bit quantized matrices out, again, sequentially and once. Therefore, during SPMM, we perform random memory access to input matrices of 8-bit as opposed to 32-bit.

Tango further introduces two SPMM variant optimizations, that is, incidence matrix-based SPMM and adaptive SPMM to improve the quantized training performance.

**Incidence matrix-based SPMM** accelerates the SPMM variant, i.e., step 1 of Figure 0(b). Particularly, this SPMM computes the gradients of node features by aggregating the incoming edge features for each node. Using node \(v_{3}\) as an example, as shown in Figure 4(a), because \(v_{3}\) contributes to the edge features for \(e_{3}\) and \(e_{4}\), its gradient is the partial derivative of \(e_{3}\) and \(e_{4}\), that is, \(\partial v_{3}=\partial e_{3}+\partial e_{4}\). However, because DGL uses the adjacency matrix format for the graph, shown in Figure 4(a), we need three matrices for the SPMM, that is, this graph, \(\partial e\), and the node features with all "1"s.

Figure 4. Tango GEMM with on-the-fly quantization and scaling factor s computation.
The drawback of this design is two-fold: First, one needs to allocate and access the all "1" node feature matrix which is redundant and expensive. Second, although the operation in Figure 4(a) can be formulated as an SPMM operation, it includes three inputs, which is not supported by the state-of-the-art cuSPARSE library.

In Figure 4(b), Tango formulates this, computation as an incidence matrix-based SPMM. Particularly, the incidence matrix is a \(V\times E\) matrix where \(V\) and \(E\) are, respectively, the number of nodes and edges in the graph. Each row of the incidence matrix contains the incoming edges of each node by marking the associated entry as 1. This design allows us to compute step, by multiplying two matrices, i.e., the incidence matrix and edge feature. Because we only need two input matrices, Tango can now adopt high-performance cuSPARSE SPMM kernels for step, which is significantly faster than DGL's three matrices-based SPMM.

**Kernel count-based adaptation.** Chances are certain SPMM computations still involve three matrices, such as step in Figure 0(a). In Figure 6, we demonstrate how one could transform a three-matrix-based SPMM kernel into a collection of two-matrix-based SPMM kernels or, to an extreme, Sparse Matrix-Vector multiplication (SpMV) kernels. Once that transformation is completed, one could directly rely on cuSPARSE to perform each two-matrix-based SPMM or SpMV. Note, we prefer cuSPARSE over DGL primitives because our evaluation shows that a single cuSPARSE SPMM kernel is significantly faster than DGL's native two-matrix-based SPMM across various configurations.

However, the benefits brought by cuSPARSE kernels diminish along with the increment of the number of kernels, as kernel invocation cost soars when too many kernels are launched. In summary, neither DGL nor transformed cuSPARSE beats the other across all configurations. We hence adaptively leverage these two solutions to achieve the best performance of both worlds.

Figure 5(a) depicts the SPMM with both edge and node features as matrices. In this design, the first head of the node feature is scaled by the first element of the edge feature and similarly for the second head. We can use multiple optimized cuSPARSE SPMM kernels to replace the native kernel used in DGL. In this case, the two heads can be finished by two SPMM kernels. Figure 5(b) assumes we have four heads in step, of Figure 6. In this case, we arrive at four SpMV kernels for the original three-matrix-based SPMM.

**SDDMM with on-the-fly dequantization.** Similar to our SPMM design, we first perform sequential memory access to quantize the input matrices. During SDDMM, we perform random memory access on those quantized matrices of smaller sizes. Briefly, the existing SDDMM performs one round of random access on full precision matrices. In contrast, Tango performs one round of sequential access on full precision matrices and one round of random access to the low precision matrices. This will lead Tango to have a shorter turnaround time. Since SDDMM might perform addition or subtraction operations, one cannot directly compute the quantized values. This leads to our SDDMM with on-the-fly dequantization.

We use step  of Figure 0(a) to explain the reason. This SDDMM computes the edge feature by, assuming the scaling factor \(s\) for \(\mathbf{S}\) and \(\mathbf{D}\) are, respectively, \(\mathbf{s}\) and \(s_{\mathbf{D}}\). The addition in quantized format should be \(\mathbf{S}[q_{i}]+\mathbf{D}[q_{j}]\approx s_{\mathbf{S}}\cdot\mathbf{S}_{ Quant}[v_{i}]+s_{\mathbf{D}}\cdot\mathbf{D}_{Quant}[v_{j}]\). Because \(\mathbf{s}\) and \(s_{\mathbf{D}}\) are often not equal, one cannot directly add the quantized values \(\mathbf{S}_{Quant}[v_{i}]\) and \(\mathbf{D}_{Quant}[v_{j}]\). Therefore, Tango loads the quantized data to enjoy reduced memory traffic, subsequently on-the-fly dequantize the loaded values for addition/subtraction computation.

If SDDMM performs multiplication and division, we can conduct SDDMM directly on the quantized value. Using step, of Figure 0(b) as an example, one needs to compute \(\partial\mathbf{\alpha}[\mathbf{e}_{0}]=\partial\mathbf{\Pi}^{(I)}[v_{0}]\cdot\mathbf{H} ^{\prime}[v_{1}]\). In this case, assuming the scaling factor of \(\partial\mathbf{\Pi}^{(I)}[v_{0}]\) and \(\mathbf{\Pi}^{\prime}[v_{1}]\) are \(s_{0}\) and \(s_{1}\). The computation can be approximated as \(\partial\mathbf{\alpha}[\mathbf{e}_{0}]\approx(s_{0}\cdot\partial\mathbf{\Pi}^{(I)}[v_{0}]) \cdot(s_{1}\cdot\mathbf{\Pi}^{\prime}[v_{1}])\). We can further arrive at \(\partial\mathbf{\alpha}[\mathbf{e}_{0}]\approx(s_{0}\cdot s_{1})\cdot\partial\mathbf{\Pi} ^{(I)}[v_{0}]\cdot\mathbf{H}^{\prime}[v_{1}]\). This allows Tango to perform quantized multiplication directly. Division can also directly work on the quantized values.

**Inter-primitive optimization**. Noticing that the follow-up operators can reuse some quantized tensors, Tango caches these quantized tensors to reduce the quantization overhead. In general, there exist two caching opportunities: (1) caching forward pass for backward, (2) caching prior operators for subsequent operators. Tango develops a detection algorithm that runs on the computation graphs to automatically derive these reuse cases.

First, the backward computation can reuse the quantized tensors from the forward pass. For example, the GEMM of step, in Figure 1

Figure 5. Incidence matrix-based SPMM.

Figure 6. Transforming three-matrix-based SPMM, e.g., step, in Figure 0(a) into a collection of: (a) two-matrix-based SPMM and (b) two-matrix-based SpMV.
has the forward computation \(\mathbf{H}^{\prime}=\mathbf{H}^{(I-1)}\cdot\mathbf{W}\). The corresponding backward step contains the gradient computation for weight, that is, \(\partial\mathbf{W}=\mathbf{H}^{(I-1)}\cdot\partial\mathbf{H}^{T}\), and the gradient computation of node features, i.e., \(\partial\mathbf{H}^{(I-1)}=\partial\mathbf{H}^{\prime}\cdot\mathbf{W}^{T}\), as shown in step 1 in Figure 1b. Clearly, the quantized matrices \(\mathbf{H}^{(I-1)}\) and \(\mathbf{W}\) are used in both forward and backward computations. We thus save the quantized input \(\mathbf{H}^{(I-1)}\) and \(\mathbf{W}\) during the forward pass for the backward pass to avoid repeated quantization. Second, when two operators share the same tensor as input, we can cache the quantized tensor from the former operator and use it for the latter. For example, in Figure 1b, the SPMM in step 1 and SDMM in 1 both need the quantized \(\partial\mathbf{H}^{(I)}\). This way, we cache the quantized input tensor. We also intentionally schedule the computation orders such that the cached tensors can be reused.

We derive the caching opportunity on the computation graph, i.e., Figure 1a as follows. The computation graph consists of tensors as nodes and operators as edges. For nodes with more than one out edge, we can quantize once for multiple operators. For example, the tensor \(\partial\mathbf{H}^{(I)}_{\text{Quant}}\) in Figure 1b has two operators as out-edges, so we cache this tensor. Then we reverse the edges in the computation graph for the backward pass. In this backpropagation graph, we will check if the to-be-quantized tensors are already quantized in the forward graph in order to facilitate quantization sharing.

**Quantization overhead vs. benefit analysis.** While quantization helps reduce computation and data movement, it also brings overheads. Mainly, quantization introduces two types of overheads: parameter computing and data type casting. First, the parameter s in Equation 1 is derived by reducing the elements with the maximum absolute value. For a \(N\times N\) matrix, the reduction needs \(N^{2}\) operations to derive the absolute maximum values. Second, quantizing or dequantizing an element requires two operations: multiply and data type casting. Therefore, the quantization before the primitive performs \(4N^{2}\) floating-point operations.

For GEMM with the input matrices at sizes of \(M\times K\) and \(K\times N\), we perform \(4K(M+N)\) and \(2MN\) operations for quantization and dequantization, respectively. For GEMM, we reduce the number of multiply-accumulate instructions from \(MNK\) to \(\frac{MNK}{4}\), which is often significantly higher than the overheads. Sparse primitives quantize the node and edge feature matrices. We assume \(D\) as the size of features. Given a graph with \(N\) nodes and \(E\) edges, in SPMM, the node and edge features require \(4D(N+E)\) operations for quantization. Later, only the resultant node features are dequantized with \(2ND\) operations. Quantization in SDDMM performs \(4ND\) operations for node features. Dequantizing resultant edge features needs \(2ED\) operations. Regarding the benefits of sparse primitives, sparse primitives enjoy a better cache access pattern brought by quantization which reduces the sizes of the input matrices.

## 4. Experiments

### Experimental setup

**Datasets.** We conduct experiments on five graph datasets as shown in Table 1. The _obgn-arxiv_(Wang et al., 2017) and _Pubmed_(Wang et al., 2018) are citation graphs whose nodes represent papers and edges are the citations. The task is to predict the categories of papers. The _ogbn-products_(Wang et al., 2018) dataset depicts an product co-purchasing network, where nodes represent products sold in Amazon, and edges between two products indicate that they are purchased together. The task is to predict the category of a product. The _DBLP_ dataset is a co-authorship network where nodes are authors and edges are co-authorship (Wang et al., 2018). The _Amazon_ dataset contains products as the nodes and edges represent co-purchase (Wang et al., 2018). The tasks of _DBLP_ and _Amazon_ are to predict if a link exists between two nodes. We add the reverse edges for the directed graphs and self-connects edges to ensure the SPMM operation works for every node.

**Models.** We evaluate GCN (He et al., 2017) and GAT (He et al., 2017) from the example implementations of DGL and the models are trained with the same number of epochs and hyperparameter settings. Both models use the hidden size of 128 and two GNN layers; GAT has four attention heads. For node classification, the model generates the node embedding as the set probabilities of each category. For link prediction, we perform dot-product between two node embeddings as the score of edge existence. The training epochs for Pubmed, ogbn-arxiv, and ogbn-products are 30, 500, and 150.

**Implementation details.**_For ease of use, we integrate Tango in DGL. Therefore, all the models on DGL can enjoy the performance benefits brought from Tango without any changes._ We provide our optimized quantized CUDA kernels to replace the corresponding primitives in DGL. DGL employs the GEMM function from the cuBLAS library and sources its sparse primitives either directly from cuSPARSE (Wang et al., 2018) or through its own implementations. DGL's interface is designed in Python and its primitives are integrated as PyTorch functions. To ensure a fair comparison, the kernels of Tango are also invoked via PyTorch's auto-differential engine during training (Wang et al., 2018). Furthermore, DGL supports multiple graph data structure formats, and Tango leverages DGL's heuristics to determine the most efficient format for its primitives.

**Evaluation platforms.** We use Python 3.6.10 and CUDA 11.7 on six V100S GPUs and Intel(R) Xeon(R) Gold 6244 @ 3.60GHz CPU. The model is trained with PyTorch 1.13.0 and DGL 0.8. We also have access to a single A100 GPU for a limited time. We use that to compare GEMM on INT8 tensor core vs FP16 tensor core.

### Tango vs. state-of-the-art

We compare Tango with DGL (Wang et al., 2018) and EXACT (Wang et al., 2018). DGL trains the model in full precision, and EXACT trains the model with quantized tensors. EXACT aims to save memory by quantizing the saved tensors, and it has no optimization in computation. We set EXACT to use 8-bit quantization. Of note, the GNN models are implemented through PyTorch, which experiences significant high-level language overheads when calling Tango primitives. As a result, we observe significantly smaller model-level speedups than the primitive-level comparisons (detailed in Section 4.3).

**Training speed.** We evaluate the training speedup of Tango on GCN and GAT models. We train each model 5 times and report the average elapsed time achieving the same accuracy as the baseline,
including the forward and backward computations. As shown in Figure 8, Tango has \(1.2\times\) and \(1.5\times\) speedup on average on GCN and GAT models compared with DGL, respectively. The GAT model enjoys more benefits because it contains more quantized primitives than GCN. Further, larger graphs have more speedup for the GCN model, except the DBLP dataset, which has the smallest average degree among the five graphs. Overall, Tango achieves an average speedup of \(2.9\times\) on GCN and \(4.1\times\) on GAT compared with EXACT. The key takeaway is that applying quantization without appropriate optimizations will lead to a significant slowdown (e.g., EXACT).

**Accuracy study.** Figure 7 studies the accuracy impact of the techniques in Section 3.2 for GCN and GAT. In particular, we evaluate Tango, Tango with quantized layer before Softmax (**Test1**), and Tango without stochastic rounding (**Test2**). For clarification, the training crashes without quantization-aware matrix multiplication and full precision weight update. The baseline models are trained in FP32 with the same number of epochs as the quantized training.

Overall, Tango could achieve \(>\)99% accuracy of the full precision training with the same number of epochs. When quantizing the layer before Softmax (**Test1**), the models show noticeable accuracy loss except for the DBLP dataset, GCN model on Pubmed, and GAT model on ogbn-arxiv. The average relative accuracy drop is 9.7%. Despite the similar final accuracy, GAT on Pubmed and GCN on ogbn-arxiv converge slower than the baseline by 18 and 35 epochs, respectively. For quantization without stochastic rounding (**Test2**), we observe for GCN on Pubmed and both models on DBLP and Amazon, the quantization error changes the optimization direction in some epochs. Thus the models take more epochs to recover. Moreover, the models on ogbn-arxiv and ogbn-products suffer from significant accuracy drops. As shown in Figure (a)a, although the model can achieve convergence, the training process shows more instability than that with stochastic rounding.

**Multi-GPU training.** Figure 9 studies Tango's impact on multi-GPU training. We directly adopt DGL's mini-batch multi-GPU training. That is, each GPU trains the model on a batch of sampled subgraphs per epoch. Then, the gradients of all GPUs are updated by an all-reduce operation. _We compare the training speed between full precision baseline and Tango using the same number of GPUs_.

Tango achieves speedup over the full precision baseline via transferring the quantized node features and gradients. Since we perform stochastic rounding-based quantization, this process will introduce nontrivial turnaround time. In Tango, we overlap the feature quantization with the subgraph sampling. The overall trend is that more GPUs would enjoy higher speedup as the Peripheral Component Interconnect Express (PCI-E) congestion is better alleviated by our quantization. Particularly, the speedup increases from \(1.1\times\) to \(1.5\times\), and \(1.2\times\) to \(1.7\times\) from two to six GPUs on GCN and GAT, respectively.

### Turnaround time analysis

**Caching the quantized tensors.** Figure 10 shows the performance of caching the quantized tensor in forward for backward reuse. We test the GEMM primitive on different datasets. We test with two hidden sizes, \(D=128\) and \(D=256\). The result shows \(1.7\times\) and \(1.6\times\)

Figure 8. The speedup of training the GNN models with Tango and EXACT compared with DGL.

Figure 7. The convergence analysis of GCN and GAT with Tango. Test1 denotes Tango with quantized layer before Softmax. Test2 denotes Tango using nearest rounding instead of stochastic rounding.

Figure 9. Tango’s impact on multi-GPU training. The X-axis is the number of GPUs.
on average when \(D=128\) and \(D=256\), respectively. The saving is related to the data size; smaller graphs, such as Pubmed, enjoy more time savings.

**GEMM.** Figure (a)a shows the speedup of our quantized GEMM over cuBLAS GEMM with the hidden size \(D=256\) and \(D=512\). Of note, we include the quantization cost in Tango GEMM time. Our quantized GEMM primitive has \(2.2\times\) and \(2.5\times\) speedup on average when \(D=256\) and \(D=512\), respectively. The trend also suggests that quantization offers more speedup on the GEMM operator when the hidden size increases. In addition, we also compare our quantized INT8 GEMM with FP16 GEMM on A100 Tensor Core GPUs. Figure (b)b shows that our quantized GEMM primitives have \(1.9\times\) for \(D=256\) and \(1.8\times\) for \(D=512\). Since both baseline and Tango use Tensor Core, the speedup over baseline is smaller than using CUDA core because the performance difference of computing in INT8 and FP16 is \(2\times\) on A100 tensor cores. Further, we observe a speedup drop for a bigger \(D\) because our quantization needs to scan through a bigger tensor to extract the scaling factor \(s\).

Figure (a)a shows the profiling results of quantized GEMM. We profile the ratio of achieved computation throughput (operation/s), memory throughput (GB/s), Instruction Per Cycle (IPC), and the number of instructions compared with cuBLAS FP32 GEMM [64]. The average computation and memory throughput ratios are \(2.1\times\) and \(2.2\times\), respectively. Memory throughput is higher because our quantized GEMM writes the quantized matrix out. However, since GEMM is computation-intensive, our increased computation throughput dominates the performance impacts. Our further investigation into IPC and \(\bullet\) of instructions in Figure (b)b explains how Tango doubles the computation throughput. Our average IPC is \(\sim\)70% of the baseline, with the instruction number reduced to \(\sim\)31%. Together, we can roughly double the throughput of the baseline.

**SPMM.** Figure (a)a shows the performance of using incidence matrix SPMM for edge aggregation compared with DGL SPMM kernels. We set the edge features size ranging from 4 to 20. All dataset has an average \(2.1\times\) speedup. The ogbn-arxiv dataset has the best speedup of \(5.5\times\) on average because of the poor performance of its baseline kernel. That is, the randomness of the incidence matrix is much lower than that of the adjacency matrix from the baseline. Table 2 shows the achieved memory throughput using our incidence-based SPMM and baseline when the feature size is 16. The irregular access for the baseline on the ogbn-arxiv dataset leads to low memory throughput. Using the incidence matrix alleviates the irregularity because the edges incidents to a node are stored adjacent in memory.

Figure (b)b shows the performance of using multiple SPMM's with a small edge feature dimension in a multi-head graph attention operation compared with DGL SPMM kernels. We set the node feature dimension as (\(H\times D\)), and the edge features dimension as (\(H\times 1\)), where \(H\) represents the number of heads and \(D\) represents the hidden size of each head. Ours has \(2.1\times\), \(1.9\times\), \(2.0\times\), and \(1.8\times\) speedup over DGL's primitive on average respectively for (\(2\times 128\)), (\(4\times 128\)), (\(2\times 256\)), and (\(4\times 256\)). Increasing the number of heads leads to a smaller speedup when the hidden size is fixed because of the increased overhead of more kernel launches. The hidden size has little impact on speedup with the same head number.

Figure 14 shows the performance of using multiple cuSPARSE SpMV with a large edge feature dimension on ogbn-arxiv graph.

\begin{table}
\begin{tabular}{c|c|c|c|c|c} \hline Dataset & ogbn-arxiv & ogbn-products & Pubmed & DBLP & Amazon \\ \hline Ours (GB/s) & 344.06 & 491.72 & 353.38 & 331.57 & 342.93 \\ \hline Baseline (GB/s) & 41.26 & 244.22 & 131.89 & 297.67 & 105.88 \\ \hline \end{tabular}
\end{table}
Table 2. The achieved memory throughput using incidence-based SPMM and DGL baseline.

Figure 11. Tango GEMM vs cuBLAS GEMM.

Figure 12. The hardware profiling of quantized GEMM.

Figure 10. The speedup of caching the quantized tensors.
We test the feature size ranging from 2 to 12. When the size is smaller than 6, ours has a 1.6\(\times\) speedup on average over DGL's implementation. The result shows the increased turnaround time as the number of kernels increases.

**SDDMM.** Figure 15 shows the performance of quantized SDDMM compared with DGL SDDMM kernels. We evaluate two SDDMM variants, including the row-wise dot-product (step in Figure 1b) and element-wise addition (step in Figure 1a), denoted as _SDDMM dot_ and _SDDMM add_. The node features are matrices with the size of (4, 64). Our SDDMM add and SDDMM dot achieves 1.9\(\times\) and 1.6\(\times\) speedups over DGL, respectively.

### Speed impact for # of quantization bits

Because neither cuSPARSE nor DGL provides SPMM kernels of INT4, this section only studies INT4 GEMM and SDDMM which are implemented by Tango. Figure 16a shows the SDDMM performance using INT4 compared with full precision DGL primitives. The addition and dot-product kernels achieve, on average, 3.3\(\times\) and 1.8\(\times\), respectively. Dense graphs like _ogbn-arxiv_ and _ogbn-products_ enjoy more benefits from reduced memory traffic because the node embeddings are more likely to be reused by cache hit.

Figure 16b shows the GEMM performance using INT8, and INT4 compared with cuBLAS. Note that we run the tests on an A100 GPU with INT4 hardware support. Using INT8 and INT4 leads to 5.4\(\times\) and 6.2\(\times\) average speedup when hidden size \(D=256\). For \(D=512\), the average speedup is 8.1\(\times\) and 10.1\(\times\), respectively. Using fewer bits shows marginal improvement because the sub-byte access under-utilizes the shared memory bandwidth.

## 5. Related Work

Recent years have seen a surge of efforts on GNN (Zhu et al., 2017; Wang et al., 2018; Wang et al., 2019; Wang et al., 2019). For a comprehensive study about the history and advancements in quantization for DNNs and GNNs, we refer the readers to two surveys (Wang et al., 2018; Wang et al., 2019). In addition to the related work in Section 1, this section further discusses GNN primitives and inference.

**GNN operator optimization** projects often focus on improving the SPMM and SDDMM kernels in GNN. GE-SpMM (Wang et al., 2019) and DA-SpMM (Wang et al., 2019) propose optimizations for implementing SPMM on GPU for GNN workloads. QGTC (Wang et al., 2019) accelerates quantized GNN operations using Tensor Cores on GPU by representing the adjacency matrix as a 1-bit sparse matrix and quantizing node features in any bits, which can be computed using the 1-bit computation function on Ampere Tensor Cores. GE-SpMM, QGTC, and DA-SpMM do not support models with multi-edge features like GAT, while Tango revamps SPMM to support such models. In addition, Tango also supports quantized GEMM and SDDMM. FeatGraph (Wang et al., 2019) uses tensor compilers, providing a flexible programming interface to generate SPMM and SDDMM for various GNN operations. FeatGraph aims to exploit parallelism for customized operations between features. However, FeatGraph does not support quantization.

**GNN inference quantization** has received significant attentions recently (Wang et al., 2019; Wang et al., 2019; Wang et al., 2019; Wang et al., 2019; Wang et al., 2019). Unfortunately, none of these can achieve a shorter turnaround time for _training_ than not quantized GNN models. Particularly, (Wang et al., 2019) quantizes the GNN into binary with knowledge distillation to reduce accuracy loss. Since knowledge distillation needs to train a teacher model and a student model, the training time increases. SGQuant (Wang et al., 2019) is a quantization scheme aiming to reduce memory consumption. It assigns different bits to embeddings and attention tensors on different levels, but the mismatch of datatype incurs extra conversion overhead. In contrast, Tango introduces a variety of framework and primitive-level system optimizations, leading to a shorter turnaround time during quantized GNN training.

## 6. Conclusion

Tango identifies both the challenges and opportunities brought by quantization to GNN training. Particularly, Tango makes the following three major contributions. First, Tango introduces various lightweight rules to maintain the accuracy for quantized GNN training. Second, we design and implement quantization-aware primitives and inter-primitive optimizations to reduce the turnaround time for quantized GNN training. Third, we integrate Tango into DGL and evaluate it across a variety of GNN models and datasets to demonstrate the superior performance of Tango.

## Acknowledgement

We would like to thank the anonymous reviewers for their helpful suggestions. This work was in part supported by the NSF CRI Award No. 2331536, CAREER Award No. 2326141, and NSF 2212370, 2319880, 2328948, 2319975, 2331301 and Semiconductor Research Corporation (SRC) Artificial Intelligence Hardware program. Any opinions, findings conclusions, or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the funding agencies.

Figure 16. The turnaround time impacts by varying # of bits.

Figure 14. The performance of using multiple cuSPARSE SPMV with high edge feature dimension.

Figure 15. The performance of SDDMM operators.