GSplit: Scaling Graph Neural Network Training on Large Graphs via Split-Parallelism

GSplit: Scaling Graph Neural Network Training on Large Graphs via Split-Parallelism

Speakers: Sandeep Polisetty, Juelin Liu, Yi Fung, Hui Guan, Marco Serafini

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

Graph Neural Networks (GNNs) have emerged as a pivotal technology for extracting insights from graph-structured data across diverse fields, from social network analysis and personalized recommendations to molecular biology and materials science. By learning powerful representations for nodes within a graph, GNNs enable critical downstream tasks such as node classification, link prediction, and graph classification. However, the burgeoning scale of real-world graphs presents significant challenges to GNN training. As graphs grow in size and complexity, the computational and memory demands of GNNs often exceed the capabilities of single-GPU systems, necessitating distributed training approaches.

This talk introduces GSplit, a novel system designed to address the scalability bottlenecks of GNN training on large graphs using a paradigm called split-parallelism. Presented by Juelin Liu, GSplit offers an optimized strategy for multi-GPU environments, aiming to significantly enhance training efficiency without compromising model accuracy. The core innovation lies in rethinking how mini-batches are processed across multiple GPUs, moving away from redundant independent processing towards a collaborative, communication-efficient approach.

GSplit's significance stems from its ability to overcome a fundamental limitation of traditional data parallel training for GNNs: the repeated loading of identical node features through slow interconnects like PCIe. By enabling GPUs to collaboratively sample, load, and train on a single mini-batch, GSplit leverages fast inter-GPU connections like NVLink to drastically reduce data transfer overheads and redundant computations. This work represents a crucial step forward in making large-scale GNN training more practical and efficient, thereby expanding the applicability of GNNs to even larger and more complex datasets.

Background

▶ Watch: Introduction: GNNs and large graph training challenge (0:00)

Graph Neural Networks operate by iteratively aggregating information from a node's local neighborhood to generate a rich, context-aware representation. A GNN layer typically involves transforming node features and then aggregating features from its one-hop neighbors. Stacking multiple GNN layers allows the model to capture information from increasingly distant neighborhoods (e.g., a two-layer GNN aggregates from two-hop neighbors), thereby increasing the model's capacity and expressive power. However, this architectural benefit comes at a substantial computational cost. As layers are added, the number of nodes participating in the forward and backward passes grows exponentially, leading to a rapid increase in both memory consumption and computational complexity.

The challenge is further compounded by the sheer size of real-world graphs, which can contain billions of nodes and edges, often exceeding the memory capacity of a single GPU. To mitigate these scalability issues, mini-batch training has become a standard practice. In this paradigm, instead of processing the entire graph, only a small subgraph is sampled for each training iteration. Specifically, for a set of target vertices, their K-hop neighborhoods are sampled, and an upper bound is often applied to the number of neighbors sampled per node to control memory and computation. The full graph and its features reside in CPU memory, and only the necessary sampled subgraph data is loaded into GPU memory on demand via the PCIe bus.

While mini-batch training effectively reduces per-iteration memory footprint, scaling it to distributed environments introduces new challenges. Data parallel training is the most common distributed strategy, where target vertices are partitioned across multiple GPUs. Each GPU independently samples, loads, and trains on its assigned subset of target vertices. Gradients are then aggregated using an all-reduce operation. The problem with this approach, as identified by the GSplit team, is the frequent duplication of nodes across the K-hop neighborhoods sampled by different GPUs. This leads to the same node features being repeatedly loaded through the slow PCIe bus to multiple GPUs, significantly increasing end-to-end training time and creating a bottleneck for large-scale GNN training. This redundancy in data loading and computation in data-parallel mini-batch GNN training is the core problem GSplit aims to solve.

Key Findings

▶ Watch: Mini-batch training for GNN scalability (4:00)

GSplit introduces a novel split-parallel training paradigm that fundamentally re-architects how GNNs are trained across multiple GPUs, leading to substantial efficiency gains without compromising model accuracy. The key findings and contributions of this work are:

1. Introduction of Split-Parallelism: GSplit proposes a collaborative approach where all GPUs work together to sample, load, and train on a single mini-batch. This stands in stark contrast to data parallel training, where each GPU processes an independent subset of target vertices.
2. Significant Reduction in Redundant Operations: By having GPUs share a mini-batch, GSplit drastically reduces the redundant loading of node features through the slow PCIe bus. It also minimizes duplicated edge aggregations during both forward and backward passes, directly addressing the main bottleneck observed in data parallel GNN training.
3. Leveraging Fast Interconnects: Split-parallelism effectively utilizes high-bandwidth, low-latency interconnects like NVLink to shuffle intermediate representations and partial gradients between GPUs. This enables efficient collaboration and data sharing within a mini-batch.
4. Offline Load-Aware Graph Partitioning Strategy: A critical innovation is GSplit's specialized graph partitioning algorithm. Recognizing that naive partitioning (like Metis) leads to skewed workloads even if node/edge counts are balanced, GSplit introduces an offline load-aware strategy. This strategy estimates the expected node and edge access frequencies based on typical sampling patterns and then partitions the graph to balance the workload (estimated by access frequency) while minimizing expected cross-edges (communication overhead).
5. Consistent and Significant Speedups: Experimental evaluations demonstrate that GSplit provides consistent and substantial speedups across various datasets and graphs. On large graphs like Friendster, GSplit achieves 1.2 times to 3 times speedup compared to state-of-the-art baselines. Against a "naive baseline" that uses traditional graph partitioning, GSplit shows a 1.4 times speedup.
6. Scalability to Multi-Host Settings: GSplit's benefits extend to multi-host environments. In a four-host setting with varying numbers of GPUs, the system still delivers consistent speedups ranging from 1.2 to 2.3 times.
7. Accuracy Preservation: Crucially, all these performance improvements are achieved without any degradation in model accuracy, ensuring that the enhanced efficiency does not come at the cost of GNN model performance.
8. User-Friendly API: GSplit offers easy-to-use APIs, designed as a drop-in replacement for data parallel training, abstracting away the complexities of data shuffling and distributed execution.

Technical Deep Dive

▶ Watch: Introducing GSplit: Split-Parallel training paradigm (8:00)

GSplit's core innovation lies in its split-parallel training paradigm, which fundamentally differs from traditional data parallel approaches for GNNs. In data parallel training, each GPU independently samples a K-hop neighborhood for its assigned target vertices, loads features from CPU memory via PCIe, performs computation, and then uses all-reduce for gradient synchronization. This leads to significant redundancy: multiple GPUs often sample and load the same popular nodes and their features.

GSplit tackles this by having all GPUs collaboratively process a single mini-batch. This means that instead of each GPU working on a separate set of target vertices, they collectively handle the sampling, data loading, and computation for one logical mini-batch. The key is to leverage fast inter-GPU interconnects like NVLink (or similar high-bandwidth, low-latency networks) to exchange intermediate results and gradients efficiently.

The GSplit workflow can be broken down into several stages:

1. Offline Graph Partitioning: This is the most critical preparatory step and where GSplit introduces its unique load-aware graph partitioning strategy.

• Problem with Naive Partitioning: Simply partitioning the graph offline using traditional algorithms like Metis (which aim to balance node/edge counts and minimize edge cuts) proves insufficient. While these algorithms might distribute the graph structure evenly, they fail to account for the runtime access patterns during GNN mini-batch sampling. The result is often a highly skewed workload during training, where some GPU partitions are significantly busier than others, leading to load imbalance and wasted resources.
• GSplit's Load-Aware Approach: GSplit formalizes the partitioning problem with two goals: minimizing communication across partitions and balancing workload. It shows that balancing the expected node access frequency and minimizing the expected cross-edges achieves both goals.
• Estimation of Workload and Communication: To achieve this, GSplit employs a sampling-based algorithm to estimate workload and communication overhead.
• First, a representative set of target vertices is selected.
• Then, the actual K-hop sampling algorithm (the same one used during training) is run multiple times on these target vertices.
• During this sampling process, the access frequency for each node is accumulated. This frequency serves as an estimate of the workload associated with that node.
• Similarly, the access frequency for each edge is estimated, providing a measure of potential cross-partition communication.
• Weighted Graph Partitioning: The result of this estimation is a weighted graph, where nodes and edges are assigned weights corresponding to their estimated access frequencies. A graph partitioning algorithm is then applied to this weighted graph, aiming to balance the total weights across partitions (thereby balancing workload) while minimizing the sum of weighted edge cuts (thereby minimizing communication overhead). This offline step is computationally intensive, taking "a few seconds to minutes," but it's a one-time cost that amortizes over many training epochs.

1. Runtime Mini-Batch Processing:

• Collaborative Sampling: At runtime, when a mini-batch of target vertices is selected, each GPU begins sampling from its own partition. If the sampling process encounters a node that belongs to a different partition (as determined by the offline partitioning), the subsequent sampling task for that branch of the K-hop neighborhood is delegated to the corresponding GPU responsible for that partition. This ensures that each GPU contributes to building a global, shared mini-batch.
• Split Mini-Batch Formation: At the end of the sampling phase, each GPU holds a "split" of the overall mini-batch. This split contains all the necessary nodes, edges, and features that fall within its assigned partition, along with information about any cross-partition dependencies.
• Forward Pass with NVLink: During the GNN layer computation, a target vertex might require features from its one-hop neighbors, some of which might reside on a different GPU's partition. GSplit uses NVLink (or similar fast interconnects) to shuffle these intermediate representations (node features) between GPUs as needed. This ensures that each GPU has access to all necessary information to compute its portion of the target vertex representations.
• Backward Pass with NVLink: The backward propagation of gradients follows a similar inverse pattern. When gradients flow across edges that span partitions, partial gradients are shuffled between GPUs using NVLink. Local gradients are aggregated within each GPU's partition and then backpropagated to update model parameters and earlier layers.

1. API and Integration: GSplit provides user-friendly APIs designed as a "drop-in replacement" for existing data parallel training setups. This simplifies adoption for practitioners, abstracting away the complexities of distributed sampling, data movement, and synchronization.

Comparison with Pipeline Parallelism: During the Q&A, a question arose about the difference between split-parallelism and pipeline parallelism. The speaker clarified that in GNNs, pipeline parallelism (where different layers are assigned to different GPUs) is rarely used due to severe workload imbalance. K-hop sampling causes the number of nodes in early layers to explode rapidly (e.g., 1000 nodes in the first layer, 2000 in the second, 4000 in the third for a fanout of 2). If layers were partitioned, the bottom layers would bear the overwhelming majority of the computation, leading to extreme workload skew and poor utilization of GPUs assigned to upper layers. Split-parallelism, in contrast, partitions the graph and mini-batch across GPUs, allowing for more balanced workload distribution across the entire GNN computation.

Experimental Setup & Results

▶ Watch: How GSplit samples and trains across GPUs (10:00)

GSplit's effectiveness was rigorously evaluated against state-of-the-art baselines across a variety of datasets and configurations, focusing on training efficiency and accuracy.

Baselines: The primary comparison was against a "naive baseline" which likely represents a standard data parallel training setup employing traditional graph partitioning algorithms (like Metis) that aim to balance node/edge counts while minimizing edge cuts, but without considering dynamic access patterns. The talk implies that this baseline suffers from the workload skew issues that GSplit's load-aware partitioning addresses.

Datasets and Graphs: While the talk mentions "various datasets and graphs," it specifically highlights performance on Friendster, a notoriously large social network graph. This emphasis suggests that GSplit provides its most significant benefits on graphs that are challenging to handle due to their immense scale.

Hardware and Environment:

• The experiments were conducted in a multi-GPU setting, implicitly leveraging high-speed inter-GPU interconnects like NVLink for efficient data exchange in split-parallelism.
• The system was also evaluated in a multi-host setting, specifically a "four-host setting," to demonstrate its scalability beyond a single server. This implies the use of high-speed inter-node networking for efficient communication between hosts.

Metrics: The primary performance metric was training speedup, measured as a reduction in end-to-end training time. Crucially, the talk emphasizes that these speedups were achieved without changing the model accuracy, ensuring that performance gains do not come at the cost of model quality.

Headline Results:

• Consistent Speedups on Large Graphs: GSplit consistently delivers significant speedups, particularly on large graphs. For Friendster, speedups ranged from 1.2 times to 3 times compared to existing methods. This indicates that the larger the graph and the more pronounced the data loading redundancy in data-parallel approaches, the greater the benefit from GSplit's collaborative strategy.
• Superiority over Naive Partitioning: Against the "naive baseline" that uses traditional graph partitioning, GSplit achieved a 1.4 times speedup. This specific comparison directly validates the efficacy of GSplit's load-aware partitioning strategy in balancing workloads and reducing communication overhead compared to methods that only consider static graph properties.
• Scalability in Multi-Host Environments: In a four-host setting, GSplit continued to demonstrate robust performance, achieving speedups ranging from 1.2 to 2.3 times. This confirms that the benefits of split-parallelism extend effectively to larger, distributed clusters, making it suitable for enterprise-scale deployments.
• Accuracy Preservation: Throughout all experiments, GSplit maintained the same model accuracy as the baselines. This is a critical finding, ensuring that the performance optimizations do not introduce any compromises in the predictive power of the GNN models.

The experimental results strongly support GSplit's claim of significantly improving GNN training efficiency by intelligently managing data movement and computation across multiple GPUs, especially for large-scale graphs where traditional data parallel approaches hit bottlenecks.

Practical Implications

▶ Watch: GSplit's load-aware graph partitioning strategy (14:00)

GSplit's introduction of split-parallel training and its load-aware partitioning strategy carries significant practical implications for practitioners, infrastructure teams, model builders, and deployers working with Graph Neural Networks.

For practitioners and model builders, GSplit offers the ability to train GNNs on much larger graphs than previously feasible, without being constrained by the memory limits of a single GPU or the communication bottlenecks of traditional data parallelism. This opens up opportunities to leverage GNNs for more complex, real-world datasets in domains like social networks, knowledge graphs, and biological systems. The promise of 1.2x to 3x speedups directly translates to faster iteration cycles for experimentation, quicker deployment of new models, and the ability to explore deeper GNN architectures that might have been prohibitively slow before. The "drop-in replacement" API further lowers the barrier to adoption, allowing existing GNN training pipelines to potentially benefit with minimal code changes.

Infrastructure teams stand to gain from more efficient utilization of their GPU clusters. By reducing redundant data loading over PCIe and balancing workloads more effectively across GPUs, GSplit can lead to better resource utilization and potentially lower overall training costs. The system's ability to scale to multi-host settings (demonstrated with 1.2x to 2.3x speedups on four hosts) means it can be deployed on existing distributed GPU clusters, improving throughput without requiring fundamental changes to hardware. However, it's crucial for infra teams to ensure that the underlying hardware provides fast interconnects like NVLink between GPUs, as GSplit heavily relies on these for efficient intermediate data exchange. Without such high-bandwidth links, the benefits of split-parallelism would be severely diminished.

Deployers will appreciate the consistent model accuracy, ensuring that the accelerated training doesn't compromise the quality of the deployed GNN models. Faster training also means models can be more frequently updated with fresh data, leading to more responsive and accurate real-time applications.

Tradeoffs and Limitations:

• Offline Partitioning Cost: The load-aware graph partitioning is an offline process that can take "a few seconds to minutes." While this is a one-time cost amortized over many training epochs, it adds an initial setup overhead. For highly dynamic graphs where the structure changes frequently, this partitioning might need to be re-run, potentially impacting its overall efficiency.
• Hardware Dependency: GSplit's performance heavily relies on the availability of fast inter-GPU communication channels like NVLink. Systems without such high-bandwidth interconnects would likely not see the same level of performance improvement, as the communication overhead would become a bottleneck.
• Complexity of Sampling: The talk notes that K-hop sampling is the most widely used GNN sampling algorithm. While GSplit improves its efficiency, the fundamental complexity of sampling large neighborhoods (where node counts explode with layers) still contributes to the overall computational burden. GSplit optimizes the distributed execution of this sampling but doesn't inherently change the sampling algorithm itself.

Despite these considerations, GSplit offers a compelling solution to a critical scalability challenge in GNN training, pushing the boundaries of what's possible with large-scale graph data.

Key Takeaways

• Problem Solved: Traditional data parallel GNN training on large graphs suffers from redundant data loading via slow PCIe buses and duplicated computation across GPUs.
• Split-Parallelism: GSplit introduces a novel paradigm where multiple GPUs collaboratively sample, load, and train on a single mini-batch, leveraging fast interconnects like NVLink.
• Load-Aware Partitioning: A critical innovation is GSplit's offline load-aware graph partitioning strategy, which balances expected node access frequency and minimizes expected cross-edges to optimize workload distribution and communication.
• Significant Speedups: GSplit achieves substantial training speedups, ranging from 1.2x to 3x on large graphs like Friendster, and 1.2x to 2.3x in multi-host settings, compared to baselines.
• Accuracy Maintained: All performance improvements are achieved without any degradation in GNN model accuracy.
• Practical Impact: GSplit enables faster training of GNNs on larger graphs, improves GPU utilization, and offers a user-friendly API for adoption by practitioners, provided fast inter-GPU interconnects are available.

About the Speaker(s)

The work on GSplit was a collaborative effort involving Sandeep Polisetty, Juelin Liu, Yi Fung, Hui Guan, and Marco Serafini. The talk at MLSys 2025 was presented by Juelin Liu, who shared insights into the challenges of scaling Graph Neural Network training and the innovative solutions provided by GSplit. The team's research focuses on optimizing systems for machine learning, particularly in the context of large-scale graph processing and distributed computing.