Scalable and Adaptive Parallel Training of Graph Transformer on Large Graphs
For researchers and practitioners training large graph transformers, this framework improves scalability and efficiency, addressing a key bottleneck in deploying graph foundation models.
The paper introduces a distributed training framework for graph transformers that automatically selects parallelization strategies based on graph structure and hardware, achieving up to 3.8x speedup in sparse attention and 78% memory reduction, with up to 6x speedup on 8 GPUs.
Graph foundation models have demonstrated remarkable adaptability across diverse downstream tasks through large-scale pretraining on graphs. However, existing implementations of the backbone model, graph transformers, are typically limited to single-GPU systems, leading to long training times or out-of-memory issues on large graphs. Moreover, parallelizing graph transformer training over the full graph is challenging, as efficiency depends heavily on both the graph structure and system characteristics, such as bandwidth and memory capacity. In this work, we introduce a distributed training framework for graph transformers, which automatically selects and optimizes parallelization strategies based on the graph structure and hardware configuration. With our implementation of distributed sparse operations, we accelerate sparse graph attention by up to 3.8x and reduce memory consumption by 78% compared to state-of-the-art frameworks. On large graph benchmarks, our proposed framework achieves up to 6x speedup with system scaling up to 8 GPUs. These results demonstrate that the proposed framework improves the scalability of graph transformers, bringing them closer to serving as practical graph foundation models.