Scalable and Adaptive Parallel Training of Graph Transformer on Large Graphs

Abstract: 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.

• The paper introduces distributed parallelization strategies (GP-AG and GP-A2A) to enable full-graph transformer training while overcoming GPU memory constraints.
• It develops an analytical throughput model that automatically selects the optimal strategy based on graph and hardware characteristics to maximize runtime efficiency.
• Empirical evaluation shows near-linear speedups, significant memory savings, and maintained accuracy across benchmark large-scale graphs.

Scalable and Adaptive Parallel Training of Graph Transformer on Large Graphs

Introduction

The paper "Scalable and Adaptive Parallel Training of Graph Transformer on Large Graphs" (2604.16715) addresses a fundamental scalability bottleneck for graph transformers (GTs) in the context of graph foundation models. While large-scale pretraining on full graphs is necessary to gain universal representations, prevailing GT implementations are restricted to single-GPU systems, making training on large graphs impractical due to excessive runtime and out-of-memory (OOM) constraints. The paper introduces novel distributed parallelization strategies and an automatic optimization framework, together enabling efficient, scalable, and adaptive multi-GPU training for graph transformers without compromising on global context or exhaustively partitioning the graph.

Distributed Graph Transformer Training: Parallelization Strategies

The core technical advancement is the formalization of two parallelization strategies—Graph Parallelism with All-Gather (GP-AG) and Graph Parallelism with All-to-All (GP-A2A)—designed to distribute and scale the execution of GTs across multiple GPUs while keeping communication and memory overhead manageable.

GP-AG partitions nodes across devices and leverages collective all-gather operations for the $K$ and $V$ matrices required in the attention computation. Each process gathers the entirety of the necessary feature matrices, computes sparse attention (SDDMM, SpMM) over its subgraph, and synchronizes gradients via reduce-scatter.

GP-A2A, leveraging the independence of attention heads, partitions the heads among devices and uses all-to-all communication to distribute Q, K, and V matrices along the head dimension. Each process computes attention for all nodes but for only a subset of heads, followed by reassembling results with another all-to-all.

Figure 1: The two parallelization strategies for attention—GP-AG (node partitioning, all-gather) and GP-A2A (head partitioning, all-to-all).

Each method trades off between communication volume (scaling with node count in GP-AG, attention heads in GP-A2A), active memory footprint, and suitability depending on graph and system characteristics.

Analytical Model and Automatic Strategy Selection

A significant contribution is the analytical throughput model correlating compute (dominated by edge count, sparse kernels) and communication (driven by node count, collective bandwidth), forming the basis for an automatic graph parallelism (AGP) algorithm. This AGP tool empirically profiles kernel times and collective operation costs, then predicts per-iteration runtime for candidate strategies as the system scale and graph structure vary.

This modeling enables automatic selection of the throughput-optimal parallelization strategy, ensuring portability and scalability as both dataset and hardware configurations change.

Empirical Evaluation: Scalability and Efficiency

The empirical evaluation utilizes multi-GPU nodes with both A100 and H100 GPUs, testing on benchmark large-scale graphs (ogbn-proteins, ogbn-products, Reddit). The results highlight:

• Near-linear scalability: Up to 6.1$\times$ speedup (8$\times$A100, proteins), 5.7$\times$ (8$\times$H100, proteins), and robust scaling for product and Reddit graphs.
• Memory efficiency: 78% reduction in graph attention memory usage compared to prior frameworks (TorchGT) at $N=512$K nodes.
• Execution time: 3.8$\times$ speedup over TorchGT for attention computation at large scales.
• Prediction accuracy: The distributed framework consistently matches the accuracy of the state-of-the-art field while supporting full-graph attention, unlike clustering/partitioning approaches.

  Figure 2: Speedup as the number of GPUs increases on an 8$\times$A100 server, different graphs.

  

  Figure 3: Speedup as the number of GPUs increases on an 8$\times$H100 server, different graphs.

  

Figure 4: Execution time for sparse graph attention with increasing nodes; the framework scales more efficiently than TorchGT.

Figure 5: Node classification accuracy parity between the framework and TorchGT, confirming no loss in model performance with the new distributed approach.

The AGP analytical estimates closely match measured iteration times, confirming the model's reliability in selecting optimal strategies under varied scenarios.

Practical and Theoretical Implications

Practically, the presented framework uniquely supports full-graph transformer training at scales not previously feasible without aggressive graph partitioning, offering new capabilities for foundation model pretraining in graph domains. It directly enables large-scale transfer learning, prompt tuning, and graph representation generalization needed for knowledge graphs, biological networks, and beyond.

Theoretically, the analytical decomposition of runtime into edge- and node-dependent components clarifies fundamental scaling bottlenecks for full-graph models and informs the design of future distributed GT systems. The method’s extensibility to other sparse-attention architectures and potential for integration with method-level optimizations (e.g., fused sparse kernels [Fused3S]) further broadens its impact.

Future Directions

Two main directions are evident:

• Scaling to larger and multi-node clusters: While the current results focus on 8-GPU nodes, the underlying primitives (all-gather, all-to-all) and AGP model maintain their properties as system size increases, contingent on inter-node bandwidth.
• Integration with hardware-specific optimizations: Techniques such as fused sparse attention and in-network aggregation can further reduce the compute-communication ratio, facilitating billion-scale pretraining and inference.

Conclusion

This work systematically advances distributed graph transformer training for large-scale graphs by introducing adaptive parallelism strategies and a robust automatic selection algorithm. The framework realizes substantial speedups, memory savings, and consistent accuracy across both hardware generations and graph domains. By enabling efficient, practical full-graph pretraining, it moves graph transformers closer to being deployable as universal graph foundation models in real-world, large-scale scenarios.