PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2609.34420v1cs.LG

TL;DR

Graph Transformers’ global attention and topology-aware bias create substantial memory demands, while graph embedding and data construction complicate workload balance and existing parallelism approaches do not directly fit. MegaGraph combines three GT-specific parallelism strategies with automated cost-model-based search. It enables large-graph training where baselines fail with OOM, reporting lower peak memory and faster training while maintaining accuracy.

  • Problem

    Graph Transformers combine global attention and topology-aware bias with embedding and data-construction overheads, creating memory and workload-balance challenges that existing GNN and Transformer parallelism techniques do not directly address.

    Source: B. Motivations · pp. 2–3

  • Method

    MegaGraph combines graph-aware context, heterogeneous pipeline, and hybrid data parallelism with Profile →Model →Search automation using cost models and memory-aware pruning.

    Source: §I · p. 2

  • Results

    77.8% lower peak memory and 4.51× training speedup compared with TorchGT, while maintaining model accuracy and enabling training where baselines fail with OOM.

    Source: §V · p. 8

  • Takeaways & Limitations

    MegaGraph supports training large-scale Graph Transformers with deep architectures through specialized hybrid parallelism and automatic configuration selection.

    Source: §I · p. 2

Abstract

from arXiv · show

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.

I. INTRODUCTION

Graph Transformers can support deeper graph models, but global attention drives per-layer memory toward the single-GPU ceiling, and existing GNN or language-model parallelism does not fit their workloads. MegaGraph combines specialized parallelism with automated search and reports large-graph training with lower memory and faster training.

I. INTRODUCTION

  • Global attention gives Graph Transformers an O(N^2) per-layer memory footprint, limiting deep models under a single GPU’s memory ceiling.The introduction notes that accuracy improves with depth while GPU memory grows rapidly.Source: §I · p. 1
  • Existing Transformer and GNN parallelism techniques do not directly fit GT workloads, and coordinating hybrid strategies creates a combinatorially growing configuration space.Transformer methods assume dense, regular workloads, while GNN systems target shallow message passing; neither supplies the needed GT scheduling and attention partitioning.Source: §I · pp. 1–2
  • MegaGraph combines graph-aware context, heterogeneous pipeline, and hybrid data parallelism with an automatic Profile →Model →Search engine for selecting configurations.The search engine uses precise cost models and memory-aware pruning.Source: §I · p. 2
  • 77.8% lower peak memory and 4.51× training speedup versus the best-performing baseline, while maintaining model accuracy.MegaGraph also enables training on large-scale graphs where existing baselines fail with OOM.Source: §I · p. 2

A. Preliminaries: Graph Transformer

Hierarchical Graph Transformers coarsen graphs and construct batches whose topology-aware bias enters attention. The motivations identify distinct challenges in distributing this attention, balancing stage workloads, and searching the resulting configuration space.

A. Preliminaries: Graph Transformer

  • Graph coarsening compresses an input graph into super nodes, and batch construction derives a topology-aware bias for attention between local nodes and super nodes.This structural encoding injects topological relationships as additive bias terms in multi-head attention.Source: A. Preliminaries: Graph Transformer · p. 2

B. Motivations

  • Even after coarsening, each hierarchical GT layer stores an N × N attention-score matrix and an equally sized topology-bias matrix.Standard context parallelism rotates KV blocks, but the locally stored bias columns no longer align with the incoming blocks, producing incorrect scores.Source: B. Motivations · p. 2
  • GT attention is all-to-all within a batch, so partitioning the graph into independent subgraphs would break global attention.This makes both standard Transformer context parallelism and GNN graph partitioning or neighbor sampling unsuitable as-is.Source: B. Motivations · p. 2
  • The node embedding table and batch data construction concentrate memory overhead in the first pipeline stage, making uniform layer allocation imbalanced and creating pipeline bubbles.The embedding table scales as |V| × d, while data construction generates attention-bias tensors and feature buffers.Source: B. Motivations · p. 3
  • Context-parallel degrees, non-uniform layer partitions, and data-parallel degrees make the hybrid configuration space combinatorial and exhaustive search infeasible.For L = 24 and Ppp = 4, the passage gives 1771 layer-partition schemes before multiplying by data-parallel choices.Source: B. Motivations · p. 3

III. METHODOLOGY

MegaGraph separates distributed execution from configuration selection. Its backend supplies specialized parallelism strategies, while an automatic search engine uses cost models to choose a strategy for the graph and hardware.

III. METHODOLOGY

  • The automatic search engine predicts execution costs and selects a parallelism configuration through a Profile →Model →Search workflow.The backend provides profiling data; the selected strategy is then deployed back to the backend for distributed training.Source: §III · p. 3

A. MegaGraph Parallelism Backend

  • MegaGraph’s backend partitions attention memory, balances work across layers, and scales throughput across graph data using three specialized parallelism strategies.These are graph-aware context parallelism, heterogeneous pipeline parallelism, and hybrid data parallelism.Source: A. MegaGraph Parallelism Backend · p. 3

1) Graph-Aware Context Parallelism:

Graph-aware context parallelism distributes attention and its topology bias together to preserve correctness across GPUs. During KV rotation, it dynamically slices the matching bias columns, reducing bias-matrix memory while retaining topological information.

1) Graph-Aware Context Parallelism:

  • GA-CP partitions attention and battn along the node dimension, then dynamically slices bias columns to match each incoming KV block during ring rotation.Each rank retains structural information for its local nodes relative to all nodes in the batch.Source: 1) Graph-Aware Context Parallelism · pp. 3–4, Figure 4 · p. 4
  • The bias-matrix memory footprint is reduced by approximately a factor of Pcp while preserving the topology encoded in battn.Each rank keeps a partial bias shard, and matching columns are selected as KV blocks rotate.Source: 1) Graph-Aware Context Parallelism · p. 4
  • The bias slice must enter each score before the online softmax update so its running normalization statistics reflect the full bias-augmented score.The described formulation allows attention computation to overlap with KV-block communication.Source: Figure 5 · p. 4

2) Heterogeneous Pipeline Parallelism:

MegaGraph addresses pipeline imbalance from the embedding and data-construction load by assigning layers non-uniformly. It also profiles computation and communication to support automated configuration search, while hybrid data parallelism distributes graph data across model replicas.

2) Heterogeneous Pipeline Parallelism:

  • For a 24-layer GT across 4 GPUs, uniform partitioning makes the first stage use approximately 65% more peak memory; partitioning layers as brings stage memory to comparable levels.The embedding-aware allocation gives fewer Transformer layers to the first stage and more to later stages to balance execution times and reduce pipeline bubbles; GPipe processes each batch as K micro-batches.Source: 2) Heterogeneous Pipeline Parallelism · p. 4

3) Hybrid Data Parallelism:

  • Each data-parallel device mesh holds a complete model replica, processes its own mini-batch partition, and synchronizes gradients with All-Reduce after each training step.The data construction module partitions mini-batches across Pdp device meshes.Source: 3) Hybrid Data Parallelism · pp. 4–5

B. Automatic Configuration Search

  • The automatic search module uses a Profile →Model →Search workflow to identify a parallelism configuration, with fine-grained layer assignments requiring more than degree-of-parallelism choices.MegaGraph profiles computation and communication in separate phases on hardware; measurements include model-component runtimes, point-to-point transfers, and collective gradient synchronization.Source: B. Automatic Configuration Search · p. 5, 1) Profiling · p. 5

1) Profiling:

  • For each candidate context-parallelism degree, computation profiling records forward and backward times for a Transformer layer and runtimes for the embedding and output projection layers.Communication profiling measures transfer time by message size, and uses warm-up iterations followed by three profiling runs averaged to estimate runtime overhead.Source: 1) Profiling · p. 5

2) Cost Model Construction:

MegaGraph models per-mesh memory and training-step latency from profiled costs. Its estimates account for stage-specific parameters and overheads, layer computation, pipeline bottlenecks, data-loading stalls, and gradient synchronization.

Memory cost model

  • Per-mesh memory comprises weights, activations, attention features and bias, and fixed framework and communication overheads.Weights include parameters, gradients, and optimizer states; the embedding and projection parameters add stage-dependent memory to the first and last pipeline stages.Source: 2) Cost Model Construction · p. 5
  • The model uses fitted coefficients for profiled memory costs, including attention bias proportional to layer count and batch nodes, while the indicator function selects stage-specific extra costs.The supplied memory-model expressions include cmlp, cattn, cbias, and cfeat coefficients.Source: 2) Cost Model Construction · p. 5

Latency cost model

  • The latency model increases layer-computation time by a profiled factor γ when layer count exceeds lbreak = cbreak/(b × d), reflecting the latency penalty after the working set surpasses L2 cache capacity.Embedding and projection computation are added at the first and last stages, respectively; Fig. 7 is described as validating the cache-related transition.Source: 2) Cost Model Construction · pp. 5–6
  • Pipeline cycle time is set by the slowest stage, with boundary activation size Sact = (b/K) × (N′/Pcp + Ns) × d transferred between stages.The step-latency estimate also includes CPU-side subgraph-sampling and attention-bias-construction stalls and All-Reduce over total parameter size.Source: 2) Cost Model Construction · p. 6

3) Optimal Strategy Search:

MegaGraph searches valid combinations of pipeline, context, and data parallelism together with heterogeneous layer partitions. Its dynamic program prunes partitions predicted to exceed the device memory limit and selects the configuration with minimum predicted step latency.

3) Optimal Strategy Search:

  • The search explores valid GPU-factor combinations and heterogeneous layer assignments, using dynamic programming to find minimum-latency partitions while pruning any stage predicted to exceed the device memory limit.The algorithm evaluates final step latency for feasible configurations and returns the best configuration and layer partition.Source: 3) Optimal Strategy Search · p. 6

A. Experimental Setup

The experiments cover graph datasets at several scales, compare MegaGraph with two complementary baselines, and measure accuracy, training time, and peak memory.

A. Experimental Setup

  • Products is the large-scale end-to-end benchmark, Reddit and Reddit2 are medium-scale datasets, and Citeseer is used for small-scale evaluation.Computers supports ablation and configuration search; Arxiv provides additional cross-architecture accuracy validation.Source: A. Experimental Setup · p. 6
  • ANS-GT isolates system-level gains under the same hierarchical model, while TorchGT provides a cross-architecture accuracy reference.ANS-GT is single-GPU only; TorchGT supports multi-GPU sequence-parallel full attention.Source: A. Experimental Setup · p. 6
  • The end-to-end comparison uses 16 layers, 8 heads, and 64-dimensional hidden and FFN layers, with batch sizes of 1024 for small and medium graphs and 8192 for large graphs.TorchGT on one GPU uses a reduced batch size of 6800 to avoid OOM; all systems train for 100 epochs.Source: A. Experimental Setup · p. 6, Figure 8 · p. 7
  • End-to-end and ablation experiments run on dual A100 GPUs, cost-model validation on dual Tesla M40 GPUs, and Pareto search on four RTX 4090 GPUs.The measured metrics are test accuracy, steady-state training time per epoch, and peak GPU memory.Source: A. Experimental Setup · p. 7

B. End-to-End Performance

MegaGraph enables training on large graphs that exhaust ANS-GT, while reducing memory use and speeding training against TorchGT; the cost models are also validated across configurations.

B. End-to-End Performance

  • MegaGraph’s single- and two-GPU results differ by less than 0.2 percentage points across datasets, and MegaGraph exceeds TorchGT by over 10 percentage points on Products.On small and medium graphs, MegaGraph and ANS-GT achieve comparable accuracy; the parallel backend introduces no accuracy degradation.Source: B. End-to-End Performance · p. 7
  • 4.51× and 4.36× speedups over TorchGT on two GPUs are reported for Products and Reddit, respectively, while ANS-GT runs OOM on both datasets.MegaGraph’s one-to-two-GPU speedups are 1.88× on Products and 2.02× on Reddit.Source: B. End-to-End Performance · p. 7
  • 77.8% lower peak memory on Products is reported for MegaGraph-1G than TorchGT-1G′, despite TorchGT-1G′ using a reduced batch size to avoid OOM.MegaGraph-1G uses 16.53 GB, compared with 74.4 GB for TorchGT-1G′; MegaGraph-2G reduces memory by 70.5% over TorchGT-2G.Source: B. End-to-End Performance · p. 7

C. Accuracy of Cost Model

  • Across 216 configurations on dual Tesla M40, memory prediction achieves MAPE = 3.37% and R2 = 0.994; the piecewise linear model captures the ≈1.68× latency penalty when layers exceed lbreak ≈8.The penalty occurs when the working set surpasses L2 cache capacity.Source: C. Accuracy of Cost Model · p. 7
  • Cost-model parameters are fitted from profiling measurements to support cross-GPU-architecture transfer, while communication costs are measured directly on target hardware.The authors report that end-to-end Pareto-search recommendations consistently match observed performance.Source: C. Accuracy of Cost Model · p. 7

D. Automatic Configuration Search

The ablation and configuration search show that parallelism strategies trade memory against throughput, so the preferred configuration depends on hardware and constraints. The conclusion reports that MegaGraph trains large graphs where baselines fail with OOM, with lower memory use and faster training while maintaining accuracy.

D. Automatic Configuration Search

  • The ablation isolates graph-aware context parallelism for attention-memory pressure and heterogeneous pipeline parallelism for embedding and data-construction imbalance.The study uses the Computers dataset with a 32-layer model; uneven layer assignment tests the need for automatic partitioning.Source: D. Automatic Configuration Search · p. 7
  • 13.50 GB peak memory with graph-aware context parallelism is a 5.2× reduction from 70.56 GB, with a 13% slowdown versus the best pipeline configuration.The table compares peak memory per GPU across parallelism configurations.Source: D. Automatic Configuration Search · p. 7, Table I · p. 8
  • The automatic search selects a PP+DP hybrid under the 24 GB RTX 4090 constraint, and all recommended configurations fit device memory.Across 166 configurations in five strategy regions, the search finds the global optimum in approximately three seconds.Source: Figure 10 · p. 8, D. Automatic Configuration Search · p. 8
  • Heterogeneous pipeline parallelism achieves the best throughput at 4.35 s/epoch, 13% faster than context parallelism, but uses 36.83 GB peak memory, 2.7× context parallelism.The two strategies do not dominate across both throughput and memory, so the preferred choice depends on hardware and memory constraints.Source: D. Automatic Configuration Search · p. 8
  • Pure DP configurations are absent from the Pareto frontier because CP+DP matches their throughput with lower memory, while CP has the lowest memory at communication cost.DP dominates on throughput without reducing per-GPU memory; PP and hybrid configurations cover other trade-offs.Source: D. Automatic Configuration Search · p. 8

V. CONCLUSION

  • 77.8% lower peak memory and 4.51× training speedup versus TorchGT are reported while maintaining model accuracy, and MegaGraph trains large graphs where existing baselines OOM.The conclusion identifies dynamic graph architectures and elastic scheduling on larger heterogeneous systems as future work.Source: §V · p. 8

Related papers

Loading 2609.34420v1…