[论文解读] QGTC: Accelerating Quantized Graph Neural Networks via GPU Tensor Core

Yuke Wang, Boyuan Feng+1 more|arXiv (Cornell University)|Nov 18, 2021

Advanced Graph Neural Networks被引用 4

一句话总结

QGTC 是一种基于 GPU Tensor Core 的新型框架，可实现量化图神经网络（QGNNs）的高性能、任意位宽推理。通过在 TC 优化的 CUDA 内核中引入位分解算术、三维堆叠位压缩、零块跳过和非零块重用技术，QGTC 在多种 GNN 工作负载上相较最先进水平的 DGL 实现了平均 2.7× 的加速，充分展示了 GPU Tensor Core 在稀疏、低比特 GNN 计算中的有效利用。

ABSTRACT

Over the most recent years, quantized graph neural network (QGNN) attracts lots of research and industry attention due to its high robustness and low computation and memory overhead. Unfortunately, the performance gains of QGNN have never been realized on modern GPU platforms. To this end, we propose the first Tensor Core (TC) based computing framework, QGTC, to support any-bitwidth computation for QGNNs on GPUs. We introduce a novel quantized low-bit arithmetic design based on the low-bit data representation and bit-decomposed computation. We craft a novel TC-tailored CUDA kernel design by incorporating 3D-stacked bit compression, zero-tile jumping, and non-zero tile reuse technique to improve the performance systematically. We incorporate an effective bandwidth-optimized subgraph packing strategy to maximize the transferring efficiency between CPU host and GPU device. We integrate QGTC with Pytorch for better programmability and extensibility. Extensive experiments demonstrate that QGTC achieves on average 2.7x speedup compared with the state-of-the-art Deep Graph Library framework across diverse settings.

研究动机与目标

• 解决尽管量化图神经网络（QGNNs）具有低内存和计算开销，但在现代 GPU 上部署时仍存在性能差距的问题。
• 克服现有 GPU 框架依赖 CUDA 核心、无法高效利用低精度 GNN 的位级操作的局限性。
• 实现在 GPU Tensor Cores 上的任意位宽计算（如 1-、2-、4 位），这些硬件在当前 GNN 工作负载中尚未被充分利用。
• 设计一个软件栈，将低比特 GNN 计算与 PyTorch 等高层深度学习框架无缝衔接，以提升可用性和可扩展性。
• 通过带宽优化的子图打包和针对 TC 的内核优化，最大化 GPU 利用率，以应对稀疏 GNN 工作负载。

提出的方法

• 提出一种位分解算术设计，通过原生支持 1 位和 4 位操作的 Tensor Core 实现任意位宽 GNN 计算的模拟。
• 实现一种三维堆叠位压缩技术，将多个低比特值打包进一个 32 位字中，减少内存占用并提升数据重用效率。
• 设计一种针对 TC 优化的 CUDA 内核，引入零块跳过机制，跳过邻接矩阵中零值块的无谓计算，最小化冗余内存访问。
• 引入非零块重用机制，通过在多个 GEMM 操作间复用先前加载的非零邻接矩阵块，减轻全局内存带宽压力。
• 开发一种带宽优化的子图打包策略，通过合并大小和访问模式相似的子图，最大化 CPU 主机与 GPU 设备之间的数据传输效率。
• 将 QGTC 集成至 PyTorch，以确保可编程性、可扩展性，并与现有深度学习工作流保持兼容。

实验结果

研究问题

• RQ1 GPU Tensor Cores 是否能被有效利用，以超越传统基于 CUDA 核心的框架，加速量化 GNN 推理？
• RQ2 如何在硬件仅原生支持 1 位和 4 位操作的前提下，高效模拟任意位宽 GNN 计算（如 2 位、4 位）？
• RQ3 在 Tensor Cores 上，针对稀疏且不规则的 GNN 计算，哪些内核级优化最能有效缓解性能瓶颈？
• RQ4 子图打包与内存访问优化在多大程度上能提升 GNN 推理中 CPU 与 GPU 间的数据传输效率？
• RQ5 位级压缩、零块跳过与非零块重用的组合对 QGNN 推理的整体吞吐量和 GPU 利用率产生何种影响？

主要发现

• QGTC 在多种 GNN 模型和数据集上相较最先进水平的 DGL 框架实现了平均 2.7× 的推理加速。
• QGTC 的 1 位 GEMM 内核性能随子图规模增大而表现出良好可扩展性，尤其在 512 至 16,384 个节点的范围内，由于计算强度提升，GPU SM 利用率随之提高。
• 非零块重用显著提升了大矩阵（如 N=8,192）的吞吐量，通过减少全局内存访问开销，尤其在节点嵌入使用更高比特宽（4–16 位）时效果更明显。
• 三维堆叠位压缩技术有效减少了内存占用，并实现了高效的位级并行，从而提升了 Tensor Core 上的算术强度。
• 采用带宽优化的子图打包策略显著提升了 CPU 与 GPU 间的数据传输效率，减少了空闲时间，增强了整体系统吞吐量。
• 该框架在现代 GPU 上展现出强大的可扩展性和资源利用率，由于完全激活 SM，在大规模子图下性能已接近峰值极限。