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

Apr 17, 2026cs.DC

Scalable and Adaptive Parallel Training of Graph Transformer on Large Graphs

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.

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

Ask PDF

May 29, 2026cs.LG

On Efficient Scaling of GNNs via IO-Aware Layers Implementations

Graph Neural Networks (GNNs) are bottlenecked by sparse, irregular memory access. Popular frameworks such as DGL and PyTorch Geometric support general message passing, but complex layers often materialize edge-wise intermediates, increasing memory traffic and limiting scalability on large graphs. We take an I/O- and arithmetic-intensity--centric view and show that widely used layers fall into three kernel families: SpMM-based convolutions, reduction-based aggregations, and attention-based layers (GATv2/Graph Transformer). For each family, we develop GPU kernels that reduce data movement, improve locality, and remain robust across realistic graphs. We also study graph reordering and find that its impact depends on the kernel mapping: it benefits neighbor-parallel (gather-dominated) kernels more consistently than feature-parallel designs. Empirically, our fused attention kernels reach up to

\textbf{3.9}\times

speedup for Graph Transformer (median

\textbf{1.6}\times

), with Tensor Core (block-sparse) variants up to

\textbf{7.3}\times

on locally dense graphs; for GATv2 we reach up to

\textbf{8.5}\times

speedup (median

\textbf{2.0}\times

) while reducing peak memory by up to

\textbf{76}\times

(median

\textbf{6}\times

). Our degree-aware reduction kernels achieve up to

\textbf{10}\times

speedup (median

\textbf{2.6}\times

). For SpMM-based layers, properly cached cuSPARSE achieves up to

\textbf{8}\times

speedup over DGL and outperforms evaluated custom baselines in the majority of evaluations. We release our implementations as drop-in replacements to support reproducible, hardware-aware GNN acceleration.

Daria Fomina, Daniil Krasylnikov, Alexey Boykov +3

1Yandex · 2HSE University · 3ITMO University

Graph Neural Networks Graphics Processing Unit Kernels

Ask PDF

Jul 6, 2026cs.LG

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

Temporal Graph Neural Networks (TGNNs) are widely used for learning from dynamic graphs in applications such as recommendation, social network analysis, and traffic forecasting. However, scaling TGNN training to large dynamic graphs remains challenging due to three intertwined bottlenecks: memory I/O, irregular computation, and temporal neighbor sampling. Existing systems often optimize these stages in isolation, leaving substantial performance headroom on the table. We present FAST, a holistic framework that accelerates end-to-end TGNN training by jointly optimizing sampling, memory I/O, and computation. FAST introduces SlimCache, which exploits within-batch compression and cross-batch caching to reduce host-device data movement under limited GPU memory budgets. It further designs thread-efficient graph operators tailored to sparse temporal subgraphs, improving GPU cache locality and reducing the latency of aggregation and edge softmax. In addition, FAST employs a topology-aware sampling strategy that improves CPU cache locality and accelerates temporal neighbor sampling. Extensive experiments on real-world large dynamic graphs show that FAST achieves an average of 2.1x (up to 4.7x) speedup over state-of-the-art systems without sacrificing model accuracy.

Yushu Cai, Qingrui Zhu, Lei Liu +3

Xidian University Guangzhou, Guangdong, China · Hunan University Changsha, Hunan, China

Temporal Graph Neural Networks Cpu-Gpu Hybrid Designs

Ask PDF