使用TC优化量化图神经网络QGTC: Accelerating Quantized Graph Neural Networks via GPU Tensor Core

PPoPP’22

Abstract

提出了第一个基于Tensor Core (TC) 的、支持任意位宽QGNN计算的框架，QGTC

引入了基于低位数据表示（low-bit data representation）和位分解计算（bit-decomposed computation）的量化低位计算设计；通过3D堆叠位压缩技术（3D-stacked bit-compression technique）、zero-tile jumping、non-zero tile reuse制作了一个TC定制的CUDA kernel；采用了一种带宽优化子图打包策略（bandwidth-optimized subgraph packing strategy），优化CPU和GPU之间的数据传输

推理速度：DGL 2.7X

1. Introduction

quantized GNN(QGNN)，以很小的精度损失，减少了计算和内存开销。

GNN适合量化的原因如下：

• 邻接矩阵可以量化。邻接矩阵只需要使用0/1来表示边是否存在，可以使用low bit来表示这些信息，以减少内存和计算开销。
• 权重和节点embedding可以量化。量化带来的微小精度损失，可以被多轮邻居聚合来弥补；量化带来的微小精度损失，可以作为微小的输入扰动，来应对某些对抗性攻击，提高模型鲁棒性。（对抗性攻击旨在对输入数据进行微小的扰动，从而使模型做出错误预测）

在实际的GPU上进行GNN量化面临的挑战如下：

• 现在流行的GNN框架（PyG、DGL、NeuGraph）是针对CUDA Core定制的，总体性能受到CUDA Core峰值吞吐量的限制，并且仅支持基于字节的数据类型（比如int32）

使用Tensor Core (TC)的原因如下：

• TC支持位级运算（比如XOR、AND）（这是量化计算的主要部分）
• 在常规NN操作上（如线性转换、卷积），TC的峰值吞吐量更大（比CUDA Core高10X以上）

直接将Tensor Core (TC)用于QGNN的计算面临的挑战如下：

• 当前的TC仅支持某些位宽（比如1位、4位），有些用户期望的位宽可能不被支持（比如2位）
• TC是为GEMM而设计的，但GNN中有大量的稀疏计算，而TC的输入矩阵尺寸是严格的（比如$8\times 128$），因此需要大量的0填充，浪费计算和内存资源
• 低位计算可能带来兼容性问题，因为现在流行的DL框架（TensorFlow、PyTorch）都不支持在低位数据类型上的操作

因此，使用TC进行QGNN的计算，需要考虑以下几方面：

• 硬件层面的支持：寻找支持QGNN计算的GPU特征，虽然没有支持任意位宽计算，但是图灵架构GPU支持基于TC的1位GEMM计算，因此可以间接实现任意位宽的GEMM计算
• 软件层面的优化：GNN计算具有高度的稀疏性和不规则性，这和为传统NN操作定制的密集的GPU计算流是不合适的。因此要采取一些优化，比如子图划分（从原始的稀疏图生成一系列密集的子图）
• 框架层面的融合：集成到主流NN框架中（PyTorch）

提出QGTC，第一个使用GPU TC支持任意位宽QGNN的框架。

• 在输入水平上，合并METIS图划分，从稀疏的输入图中生成一系列稠密的子图。（真实世界图中的节点很可能形成一个个集群，可以用来提高GNN计算效率）
• 在算法水平上，任意位宽的QGNN计算可以被分解为1位计算，因此使用基于原子1位类型的量化低位数据类型和位分解计算。
• 在GPU内核水平上，为基于批处理的稠密子图的QGNN计算定制了低位计算设计。使用1位表示子图邻接矩阵中边是否存在；使用3D堆叠位压缩技术（3D-stacked bit-compression technique）来维护量化低位节点embedding特征和权重；通过zero-tile jumping和non-zero tile reuse来充分利用子图内部的稀疏性，进一步避免不必要的计算、提高数据局部性。
• 在框架水平上，将QGTC集成到Pytorch中，引入了bit-Tensor数据类型和bit-Tensor计算的新概念，并将它们封装为一组新的Pytorch API。

主要贡献：

• 第3章：针对任意位宽，提出了1-bit composition technique，来支持不同精度的QGNN
• 第4章：基于GPU Tensor Core，引入了QGNN的高效实现。使用了一系列计算优化（比如：subgraph partitioning and batching、zero-tile jumping）、内存优化（比如：3D-stacked bit-compression、non-zero tile reuse）
• 第5章：将QGTC集成到P有torch，引入bit-Tensor数据类型和bit-Tensor计算
• 实验证明，超过DGL

2. Background and Related Work

2.1 Graph Neural Networks

batched computation（读一读Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Networks）

batched computation已被用于PyG、DGL，用于处理由于太大而无法直接放入CPU/GPU的大规模图。

batched GNN computation的流程：

• 使用图划分工具集（比如METIS）划分输入图，最小化图结构信息损失，同时最大化子图内部的边数（提升子图的模块性）

• 然后，将这些子图输入到GNN中，生成节点embedding

• 最后，将这些embedding用于下游任务

2.2 Quantization of GNNs

我的理解：在尽量保证模型准确度的前提下，通过牺牲精度，来节省内存、提高速度。

2.3 Tensor Core on GPUs

CUDA Core：标量计算

Tensor Core：矩阵计算（直接计算矩阵乘矩阵）

矩阵被分为一个个tile，tile被进一步分为fragment，每个fragment是映射到线程寄存器的一组tile元素，每个warp对应一个fragment，warp中的每个线程，在自己的寄存器中存储fragment的一部分数据，这些寄存器中的数据是共享的。

3. QGTC Algorithm Design

3.1 1-bit Composition for Quantized Ops

GNN很适合做量化，因为精度损失可以被聚合操作弥补，使得非常低位的量化也不会带来明显的准确率损失。

给定一个quantization bit $q$ 和 32位浮点数 $\alpha \in R$ ，通过下面公式将 $\alpha$ 量化为一个 $q$ 位的值：
$${\alpha }^{(q)}=\lfloor \frac{\alpha -{\alpha }_{min}}{scale}\rfloor \text{\hspace{1em}(2)}$$
其中，${\alpha }_{min}$和${\alpha }_{max}$是经验下界和经验上界，$scale=|\frac{{\alpha }_{max}-{\alpha }_{min}}{2^q}|$

我的理解：${\alpha }^{(q)}=|\frac{\alpha -{\alpha }_{min}}{{\alpha }_{max}-{\alpha }_{min}}|\cdot 2^q$，也就是说 ${\alpha }^{(q)}\in [0,2^q-1]$，即可以用 $q$ 位来表示 $a^{(q)}$

对于任意位宽的量化值，提出一种基于原子1-bit计算的新的算数类型。

3.1.1 Any-bitwidth Scalar-Scalar Multiplication

假设一个3位的标量 $a$ 与一个2位的标量 $b$ 相乘：
$$\begin{gathered} a=a{t}_2\cdot 2^2+a{t}_1\cdot 2^1+a{t}_0\cdot 2^0 \\ b=b{t}_1\cdot 2^1+b{t}_0\cdot 2^0\text{\hspace{1em}(3)} \end{gathered}$$
比如：3位标量 $a$ 的值为6（三位二进制表示为110），那么$a{t}_2=1,a{t}_1=1,a{t}_0=0$
$$a\cdot b=a{t}_{2}b{t}_1\cdot 2^3+(a{t}_{2}b{t}_0+a{t}_{1}b{t}_1)\cdot 2^2+(a{t}_{1}b{t}_0+a{t}_{0}b{t}_1)\cdot 2^1+a{t}_{0}b{t}_0\cdot 2^0\text{\hspace{1em}(5)}$$

3.1.2 Any-bitwidth Vector-Vector Multiplication

从任意位宽标量乘法（Any-bitwidth Scalar-Scalar Multiplication）扩展得到任意位宽向量乘法（Any-bitwidth Vector-Vector Multiplication）

假设一个3位的向量 $\overrightarrow{v_i}$ 和一个2位的向量 $\overrightarrow{v_j}$ 相乘，它们都有 $k$ 个元素 （相当于把公式(5)累加k次）
$$\begin{gathered} \overrightarrow{v_i}\cdot \overrightarrow{v_j}=\sum _y^k{v}_i^{(y)}\cdot v_j^{(y)} \\ =\sum _y^{k}a{t}_2^{(y)}b{t}_1^{(y)}\cdot 2^3+\sum _y^k(a{t}_2^{(y)}b{t}_0^{(y)}+a{t}_1^{(y)}b{t}_1^{(y)})\cdot 2^2+\sum _y^k(a{t}_1^{(y)}b{t}_0^{(y)}+a{t}_0^{(y)}b{t}_1^{(y)})\cdot 2^1+\sum _y^{k}a{t}_0^{(y)}b{t}_0^{(y)}\cdot 2^0\text{\hspace{1em}(6)} \end{gathered}$$
大致流程： 将向量中的元素，按位分解，做bit-bit乘法，然后归约得到最终结果。

举个栗子： 将向量 $\overrightarrow{v_i}$ 和 $\overrightarrow{v_j}$ 按位分解后，可以得到向量 $\overrightarrow{v_i}$ 在bit位置2的值 $a{t}_2^{(y)}$ 和向量 $\overrightarrow{v_j}$ 在bit位置1的值 $b{t}_1^{(y)}$，其中$y\in [0,k)$，经过乘法和加法，可以得到 $\overrightarrow{v_i}\cdot \overrightarrow{v_j}$ 在bit位置3的值。这样的1-bit向量乘法可以被高效实现为：
$$an{s}_{i,j}=popcnt(\overrightarrow{v_i}\&\overrightarrow{v_j})\text{\hspace{1em}(7)}$$
其中，$popcnt()$返回结果中1的个数，比如 $popcnt(1101)=3$。

❓❓❓❓❓❓❓❓❓❓❓❓❓❓popcnt()是干嘛的？和向量乘法什么关系？ ❓❓❓❓❓❓❓❓❓❓❓❓❓❓

用类似的方法可以获得位置0 1 2的结果，然后移位、归约即可得到最终结果。

从vector-vector乘法可以衍生出matrix-matrix乘法，结果矩阵中的每个元素都由一个vector-vector乘法得到。

3.2 Quantized Computation in GNNs

两个特性：

• 邻接矩阵可以用1位表示边是否存在
• 节点embedding矩阵和权重矩阵可以选择任意位宽来满足精度要求

大致流程：

• 邻接矩阵A用1位，节点embedding矩阵X经位分解到s位，权重矩阵W经位分解到t位。A是1位，存储在A_bin；X存储在列表X_list，其中每一个元素是一个1位X矩阵；W存储在W_list，其中每一个元素是一个1位W矩阵。

  A_bin = bitDecompse(A, 1)[0];
  X_list = bitDecompse(X, s);
  W_list = bitDecompse(W, t);
  

• 计算AX：

  for xIdx in len(X_list) do
  	X_new_list.append(BMM(A_bin, X_list[xIdx]));
  end
  

• 计算XW：

  for xIdx in len(X_new_list) do
  	for wIdx in len(W_list) do
          bitIdx = xIdx + wIdx;
  		tmp_C = BMM(X_new_list[xIdx], W_list[xIdx]);
  		C_dict[bitIdx].append(tmp_C);
  	end
  end
  

• 归约：

  for bitIdx in len(C_dict) do
      for Idx in len(C_dict[bitIdx]) do
          X += C_dict[bitIdx][Idx] << bitIdx;
  	end
  end
  // 为防止溢出，X应是全精度数据类型（全位数据类型），比如int32
  

❓❓❓❓❓❓❓❓❓❓❓❓❓❓两个疑问❓❓❓❓❓❓❓❓❓❓❓❓❓❓

• 论文伪代码第9行，应该是W_list而不是W_new_list吧
• 论文伪代码第17行，应该是X而不是X[Idx]吧

4. Implementation

4.1 Subgraph Partitioning and Batching

使用METIS进行子图划分（子图的数量由用户决定），然后batch

这样可以从两个粒度级别上控制工作流：

• 工作负载的粒度由子图数量决定。通过设置不同的子图数量，可以改变子图的大小、子图内部的稠密性。一般来说，较小的子图，其内部的稠密性更大，局部性更强。
• 处理粒度由batch数量决定。batch数量决定了GPU上每一轮次执行的子图的数量。batch的适当选择可以最大化GPU的利用率。

4.2 3D-stacked Bit Compression

现在的NN框架都只支持全精度，有两个挑战：

• 不能直接使用现在NN框架的数据类型，因为它们只支持全精度数据类型，使用这样的数据类型，无法体现出量化对内存、计算的优势
• 低位量化不能很好地适应字节边界对齐（起始地址是8的整数倍），这导致一个字节内部可能有不同数据的一部分，这给解释实际值带来困难（如何组合才能得到实际值）

因此提出3D-stacked Bit Compression，将任意位宽的输入调整为32位对齐。如Fig. 4所示：

对于 $C=A\times B,其中A:\text{3-bit}\times M\times K,B:\text{2-bit}\times K\times N$

主要三点创新：

• Tensor Core的计算尺寸：$M=8,N=8,K=128$，因此需要对那些不能被8、128整除的矩阵维度进行填充，使用 $PAD8,PAD128$

• 在压缩时，使用32位格式（32位作为一个基本单元）。

• 对 $A$ 进行列压缩（因为对 $A$ 的访问是跨列的），对 $B$ 进行行压缩（因为对 $B$ 的访问是跨行的）。

最终（填充、压缩后）：
$$\begin{gathered} A:\text{3-bit}\times (PAD8(M)\times PAD128(K)/32) \\ B:\text{2-bit}\times (PAD128(K)/32\times PAD8(N)) \end{gathered}$$
总结： 矩阵A，其中的数据经量化后变为3位，在z轴上将其展开为3个bit matrix，每个bit matrix保持原有的行数和列数，然后做相应的填充、位压缩（以32位）

❓❓❓❓❓❓❓❓❓❓❓❓❓❓疑问： ❓❓❓❓❓❓❓❓❓❓❓❓❓❓

我是这样理解的：需要做对齐，是因为量化后的数据可能是3位、5位这样不能被8整除的位数，导致一个字节内可能有多个数据的一部分，导致了难以解释实际值，在取值时难以得到实际值。但以3位矩阵A为例，

作者的做法是：先把A分成3个bit matrix，然后在每个bit matrix内部以32位对齐做行/列压缩。这和难以解释实际值有什么关系？

4.3 Zero-tile Jumping

子图划分之后，仍然可能有一些子图包含全0的tile（Tensor Core将矩阵分为一个个tile）。

tile的尺寸是 $8\times 128,即8\times 4int32$，使用一个warp中的8个线程，每个线程对应一个uint4_v向量（uint4_v由连续存放的4个int32组成），分别

检查这8行是不是都是0，如果都是0则跳过。

❓❓❓❓❓❓❓❓❓❓❓❓❓❓**疑问：**❓❓❓❓❓❓❓❓❓❓❓❓❓❓

Fig. 5的代码中，为什么上面的代码设置0x000000FF，即低8个线程有效，而下面的代码设置0xFFFFFFFF，即32个线程有效

4.4 Non-Zero Tile Reuse

当1位邻接矩阵与3位embedding矩阵相乘时：

如下Fig. 6(a)的计算方式是：邻接矩阵分别与1位embedding矩阵相乘（$adj\times \text{1st-bit},adj\times \text{2nd-bit},adj\times \text{3rd-bit}$），共计3次，都计算完后再归约，称为Cross-bit Reduction，有个问题：需要重复加载邻接矩阵中的non-zero tile。

如下Fig. 6(b)的计算方式是：邻接矩阵中的non-zero tile分别与每个1位embedding矩阵中对应位置的tile相乘，计算完当前non-zero tile的乘法之后，先部分归约，然后再做邻接矩阵中下一个non-zero tile的乘法，然后再部分规约，当邻接矩阵中所有的non-zero tile都计算完并部分规约完之后，再归约得到最终结果。

如此一来，加载non-zero tile的时间复杂度从 $O(embedding位数)$ 降为 $O(1)$

4.5 Inter-layer Kernel Fusion

主要两点：

• 在计算kernel的最后进行数据量和和位分解，避免数据传入global memory
• 将激活函数kernel融合进计算kernel，直接操作共享内存

注意： 对于隐藏层，需要在层的最后进行量化和位分解；而对于最后一层，直接把全精度的结果输出给softmax层即可。

4.6 Bandwidth-Optimized Subgraph Packing

过去的做法：通过PCIe传输子图邻接矩阵和embedding矩阵

作者的做法：在CPU端对子图邻接矩阵和embedding矩阵进行量化和位分解，然后将处理后的矩阵传输到GPU端

5. Integration with PyTorch

将QGTC集成到PyTorch，面临两个挑战：

• PyTorch是构建在字节数据类型之上的、基于张量的框架，如何表示低位数据？
• 如何做低位数据和基于字节的数据的计算？（比如 $\text{3-bit}\times int32$）

因此提出两个技术：

• Bit-Tensor Data Type：使用3D-stacked Bit Compression打包量化数据，使用PyTorch中的32位数据类型IntTensor装载任意位宽的量化数据。
• Bit-Tensor Computation：我们处理两种不同类型的计算:1)只涉及位张量的操作和2)同时涉及位张量和浮点数int32张量的操作。对于第一种类型的操作，我们构建了两个api，基于我们是想要获得int32输出还是仍然作为位张量获得量化的低位输出。对于具有低位输出的任意位宽MM, API是bitMM2Bit(C, A, B, bit_A, bit_B, Bit_C)，其中A和B是位张量，bit_A/B/C是位宽参数。对于输出为int32的任意位宽MM, API是bitMM2Int(C, A, B, bit_A, bit_B)。对于第二种类型的操作，我们将首先使用ingtensor .to_val(nbits)将位张量解码为float/int32张量。然后我们调用PyTorch中的官方api进行常规的全精度计算。

6. Evaluation

Benchmarks: Cluster GCN、Batched GIN

Baselines:

• end-to-end runtime performance comparison: DGL
• GNN aggregation kernel performance comparison: cuBLAS with int8 precision、CUTLASS with int4 precision

Datasets: Type I是许多GNN算法论文评估的流行GNN数据集；Type II是许多GNN算法研究框架中用于图kernel的流行基准数据集；Type III是在节点和边的数量上具有挑战性的GNN数据集，这些图在结构上表现出高度的不规则性。

Platforms & Metrics:

• Ubuntu server (16.04)、an 8-core 16-thread Intel Xeon Silver 4110 CPU@2.8GHz with 64GB host memory、an NVIDIA Ampere RTX3090 GPU with 24GB device memory、CUDA11.0、GCC 7.5.0 with the compilation option of “-std=c++14 -O3”
• the averaged latency of 200 rounds of end-to-end results

6.1 Compared with DGL

在本节中，我们对不同比特宽度的选择进行了与DGL框架的详细端到端比较。如图7(a)和图7(b)所示，相对于DGL，在三种类型的数据集上，QGTC在簇GCN和批量GIN方面的端到端推理速度平均提高了2.6倍和2.8倍。我们还注意到性能的增益与我们选择的比特宽度密切相关，从16位到32位，性能与2位到8位的设置相比表现出很大差异。接下来，我们提供对每种类型数据集的详细分析和见解。对于权重和节点嵌入特征，QGTC使用较少的比特更有可能达到更高的性能。这是因为较小的比特宽度会导致在位级别上的内存访问减少和计算减少。DGL由于以下原因达到了较差的性能：1) FP32计算与我们的QGTC低比特设计相比具有高计算复杂性；2) DGL只能依赖CUDA核心进行计算，其计算性能受到峰值计算性能的自然限制，而我们的QGTC在具有更高吞吐性能的TC上运行。与簇GCN相比，对批处理GIN的实验结果显示QGTC相对DGL的收益更高。这是因为批处理GIN在邻居聚合之前首先应用节点更新，这导致更高的计算与通信比。QGTC在Type III数据集上实现了相对较高的性能改进。主要原因是在相同数量的分区下，每个分区（子图）的大小将增加，因为节点/边的数量更多。这也提高了计算强度，突显了QGTC在GPU Tensor Cores上进行量化低比特计算的性能优势。

关于量化比特的准确性：为了构建QGNN模型，我们应用了量化感知训练，并在两个大型Type III数据集上评估了模型测试准确性与量化比特的关系，以演示在GCN模型上的情况。如表2所示，GNN模型对低比特量化具有鲁棒性，并且在很大程度上可以保持模型的准确性。将这些结果与我们上述在不同量化比特下的性能评估结果结合起来，我们可以得出结论：在运行时性能和模型准确性之间做出正确的权衡是有意义的，并且可以在不同的应用设置中带来好处。

6.2 Compared with other baselines

与cuBLAS-int8 on TC的比较：我们进一步比较了我们的低比特计算（从2位到7位）与基于Tensor Core的量化（int8）GEMM解决方案的cuBLASgemmEX在吞吐性能方面的最新情况。请注意，int8是cuBLAS目前支持的在Tensor Core上进行量化计算的最小比特数。在这项研究中，我们主要关注AX的计算（即，𝑁 × 𝑁 × 𝐷，其中𝑁是节点数，𝐷是节点嵌入维度）用于邻居聚合阶段。如图7©所示，相对于Tensor Core cuBLAS（int8）在低比特设置下，QGTC实现了显著的吞吐性能改进。主要原因是我们的QGTC设计有效地减少了位级别上的计算和数据移动，从而在GPU上实现了低比特量化的实际性能提升。当量化比特数接近8位时，由于位级别计算的增加，性能提升会减少。

与CUTLASS-int4 on TC的比较：我们还与最新版本的CUTLASS [25]（v2.7）进行了比较，采用了int4 Tensor Core GEMM，通过AX的吞吐量（TFLOPs）进行比较。结果总结在表3中，我们可以清楚地看到在吞吐量方面相对于CUTLASS实现的性能优势。请注意，所有报告的小数数字都是以TFLOPS为单位；N是邻接矩阵的大小，Dim是节点特征嵌入的维度。图的邻接矩阵存储在1位中。QGTC（2位）表示对嵌入矩阵采用2位表示。这种性能改进背后的主要原因是，我们的QGTC设计可以使用1位二进制表示图的邻接矩阵，并使用n位（n=1,2,3,4）表示节点嵌入矩阵，而CUTLASS int4仅支持4位× 4位。因此，在计算过程中，我们不得不对邻接矩阵和嵌入矩阵都使用4位表示。