TorchGT: A Holistic System for Large-scale Graph Transformer Training

TL;DR

TorchGT tackles the scalability barrier of graph transformers on real-world graphs by jointly optimizing algorithmic attention patterns and distributed system design. It introduces a three-pronged co-design: Dual-interleaved Attention to fuse local topology with global attention, Cluster-aware Graph Parallelism for topology-aligned distributed training, and Elastic Computation Reformation to map clustered patterns into memory-friendly sub-blocks, guided by an Auto Tuner. The approach yields up to 62.7x speedups and supports graph sequence lengths up to 1M while maintaining convergence and task-agnostic applicability to graph-level and node-level tasks. This work enables practical, scalable training of large-scale graph transformers on graphs with billions of edges, broadening the applicability of graph transformer models in real-world domains.

Abstract

Graph Transformer is a new architecture that surpasses GNNs in graph learning. While there emerge inspiring algorithm advancements, their practical adoption is still limited, particularly on real-world graphs involving up to millions of nodes. We observe existing graph transformers fail on large-scale graphs mainly due to heavy computation, limited scalability and inferior model quality. Motivated by these observations, we propose TorchGT, the first efficient, scalable, and accurate graph transformer training system. TorchGT optimizes training at different levels. At algorithm level, by harnessing the graph sparsity, TorchGT introduces a Dual-interleaved Attention which is computation-efficient and accuracy-maintained. At runtime level, TorchGT scales training across workers with a communication-light Cluster-aware Graph Parallelism. At kernel level, an Elastic Computation Reformation further optimizes the computation by reducing memory access latency in a dynamic way. Extensive experiments demonstrate that TorchGT boosts training by up to 62.7x and supports graph sequence lengths of up to 1M.

Figures (12)

• Figure 1: The test accuracy of graph transformers when trained with different sequence lengths $S$.
• Figure 2: Training iteration time breakdown when training Graphormer on ogbn-products in different sequence lengths on two types of GPUs: RTX 3090 and A100.
• Figure 3: Overview of TorchGT architecture and workflow.
• Figure 4: Detailed training process on one worker with TorchGT, which includes three key components.
• Figure 5: Three attention layouts after Dual-interleaved Attention, Cluster-aware Graph Parallelism and Elastic Computation Reformation respectively. (c) is obtained by compacting elements to adjacent neighbors inside clusters.