MegaGraph: Towards Efficient Training of Large-Scale Graph Transformers with Automated Hybrid Parallelism

Tong Qiao; Ao Zhou; Yingjie Qi; Chunming Hu; Jianlei Yang

cs.LGSep 28, 2026

MegaGraph: Towards Efficient Training of Large-Scale Graph Transformers with Automated Hybrid Parallelism

Authors: Tong Qiao, Ao Zhou, Yingjie Qi, Chunming Hu, Jianlei Yang

Organizations: School of Computer Science and Engineering, Beihang University, China · School of Software, Beihang University, China · State Key Laboratory of Complex and Critical Software Environment, Beihang University, China · Qingdao Research Institute, Beihang University, Qingdao, China

Abstract

Graph Transformers (GTs) offer superior representation capabilities by overcoming the depth limitations and over-smoothing issues of traditional Graph Neural Networks (GNNs). However, scaling GTs to large graphs poses critical bottlenecks. Specifically, the attention score matrix and its associated topology-aware bias matrix jointly incur significant per-layer memory overhead, and heavy graph embedding layers result in severe workload imbalances. These characteristics are unique to GT training and are not addressed by parallelism techniques designed for either conventional GNNs or Transformers, making a dedicated solution necessary. This paper introduces MegaGraph, the first automated hybrid parallelism framework designed for efficient GT training. MegaGraph designs three specialized strategies, namely graph-aware context parallelism, heterogeneous pipeline parallelism, and hybrid data parallelism, to support efficient training on large-scale graphs. However, coordinating these three parallelism strategies yields an exponentially large configuration space. To address this complexity, an automatic search engine leverages precise cost models via a Profile - Model - Search workflow to identify the optimal parallelism configuration. Evaluations demonstrate that MegaGraph enables training on large-scale graphs where state-of-the-art baselines fail due to out-of-memory (OOM) errors. The framework reduces per-device peak memory by up to 77.8% and achieves up to 4.51 $\times$ training speedup while maintaining model accuracy.

Figures & tables

Fig. 1: Scalability of graph representation learning.

Fig. 3: The overview of MegaGraph framework.

Fig. 4: Graph-aware context parallelism: KV blocks rotate through the ring while bias columns are dynamically sliced to match each incoming KV block.

Fig. 5: Comparison of pipeline parallelism strategies for Graph Transformers. Heterogeneous pipeline parallelism achieves better workload balance through non-uniform layer allocation.

Fig. 6: Peak memory footprint breakdown under different pipeline layer partition strategies.

Fig. 7: Validation of the piecewise latency cost model with $P_{dp}=2$ .

Fig. 8: End-to-end performance comparison using a 16-layer, 8-head GT with hidden and FFN dimension 64. OOM entries indicate training failure. TorchGT-1G ′ denotes a reduced batch size of 6800.

Fig. 9: Accuracy of the memory and latency cost models. Diagonal lines indicate perfect prediction ( $Y=X$ ).

Config	Time/ep	Mem 0	Mem 1	Strategy
(1,1,1) Baseline	8.13s	70.56 GB	N/A	None
(1,2,1) CP=2	4.92s	13.50 GB	13.50 GB	GA-CP
(2,1,1) PP=2 balanced	4.35s	36.83 GB	34.06 GB	Hetero. PP
(2,1,1) PP=2 unbalanced	6.00s	53.68 GB	17.21 GB	Naive PP

TABLE I: Ablation of parallelism strategies. Each configuration is denoted as a

(P_{pp},P_{cp},P_{dp})

triple. Mem 0 and Mem 1 report peak memory per GPU.

Explore similar work

Scalable and Adaptive Parallel Training of Graph Transformer on Large Graphs
Apr 17, 2026Jun-Liang Lin, Kamesh Madduri, Mahmut Taylan KandemirMultiplex Graph Transformers Cpu-Gpu Hybrid Designs
The Pennsylvania State University University Park, PA, USA
On Efficient Scaling of GNNs via IO-Aware Layers Implementations
May 29, 2026Daria Fomina, Daniil Krasylnikov, Alexey Boykov +3Graph Neural Networks Graphics Processing Unit Kernels
1Yandex · 2HSE University · 3ITMO University
FAST: A Holistic Framework for Optimizing Memory-I/O, Computation, and Sampling in Temporal GNN Training
Jul 6, 2026Yushu Cai, Qingrui Zhu, Lei Liu +3Temporal Graph Neural Networks Cpu-Gpu Hybrid Designs
Xidian University Guangzhou, Guangdong, China · Hunan University Changsha, Hunan, China

Apr 17, 2026, cs.DCJ/K move · Enter open · S save

Scalable and Adaptive Parallel Training of Graph Transformer on Large Graphs

Jun-Liang Lin, Kamesh Madduri, Mahmut Taylan Kandemir

The Pennsylvania State University University Park, PA, USA

Multiplex Graph Transformers Cpu-Gpu Hybrid Designs Scalability Pretraining

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Scalable and Adaptive Parallel Training of Graph Transformer on Large Graphs

On Efficient Scaling of GNNs via IO-Aware Layers Implementations

FAST: A Holistic Framework for Optimizing Memory-I/O, Computation, and Sampling in Temporal GNN Training