EmbedPart: Embedding-Driven Graph Partitioning for Scalable Graph Neural Network Training

Nikolai Merkel; Ruben Mayer; Volker Markl; Hans-Arno Jacobsen

cs.LGApr 1, 2026

EmbedPart: Embedding-Driven Graph Partitioning for Scalable Graph Neural Network Training

Authors: Nikolai Merkel, Ruben Mayer, Volker Markl, Hans-Arno Jacobsen

Organizations: Technical University of Munich (TUM) · University of Bayreuth · TU Berlin, BIFOLD, DFKI · University of Toronto

Abstract

Graph Neural Networks (GNNs) are widely used for learning on graph-structured data, but scaling GNN training to massive graphs remains challenging. To enable scalable distributed training, graphs are divided into smaller partitions that are distributed across multiple machines such that inter-machine communication is minimized and computational load is balanced. In practice, existing partitioning approaches face a fundamental trade-off between partitioning overhead and partitioning quality. We propose EmbedPart, an embedding-driven partitioning approach that achieves both speed and quality. Instead of operating directly on irregular graph structures, EmbedPart leverages node embeddings produced during the actual GNN training workload and clusters these dense embeddings to derive a partitioning. EmbedPart achieves more than 100x speedup over Metis while maintaining competitive partitioning quality and accelerating distributed GNN training. Moreover, EmbedPart naturally supports graph updates and fast repartitioning, and can be applied to graph reordering to improve data locality and accelerate single-machine GNN training. By shifting partitioning from irregular graph structures to dense embeddings, EmbedPart enables scalable and high-quality graph data optimization.

Figures & tables

Panel of Figure 1 — Figure 1. (a) Traditional approaches perform graph partitioning or reordering as a preprocessing step based solely on the graph structure. (b) Our approach ( EmbedPart ) derives partitions from node embeddings produced during GNN training. These embeddings capture both structural and feature information and enable partitioning or reordering directly in the dense embedding space using scalable clustering.

Figure 2. Vertices are assigned to partitions. Edges connecting vertices of different partitions are cut.

Figure 3. Distributed GNN training pipeline: 1 the input graph is partitioned into $k$ partitions; 2 hyperparameters for the GNN are selected; 3 the model is trained in a distributed fashion; 4 the model is evaluated; repeat 3 - 4 with the next hyperparameter set; 5 once the final model is trained, it is applied to downstream tasks.

Figure 4. EmbedPart overview: The input to EmbedPart is like in GNN training, a graph with features that are attached to nodes. In Phase 1, we train a GNN model (the actual workload) and get node embeddings. At any point, the process can transition to Phase 2, where clustering on the embeddings assigns nodes to clusters. To maintain balanced distributed training, nodes are migrated from overloaded to underloaded clusters such as node 4. The resulting clusters are then used for graph partitioning. Finally, the process returns to Phase 1 to continue GNN training with the improved data layout.

Graph	Vertices	Edges	Mean Deg.
arxiv	169,343	1,166,243	13.77
reddit	232,965	114,615,892	983.98
products	2,449,029	123,718,280	101.03
papers100M	111,059,956	1,615,685,872	29.10

Table 1. Evaluation graphs along with the number of vertices, number of edges, and mean degree.

Table 2. Average edge-cut ratio (lower is better) after different numbers of GNN training epochs.

Figure 5. Loss and accuracy for reddit over epochs.

Table 3. Edge-cut ratio reduction by an untrained GNN compared to random partitioning in percentage. Larger is better.

Panel of Figure 6 — Figure 6. Comparison of partitioning time speedups and edge-cut reductions.

Panel of Figure 7 — Figure 7. Partitioning run-time on different graphs and number of partitions. Lower is better

Panel of Figure 8 — Figure 8. Partitioning quality metrics: Edge-cut ratio on different graphs and number of partitions. Lower is better.

Figure 9. Speedup of partitioners over random partitioning. Larger is better.

Panel of Figure 10 — Figure 10. Achieved speedups for distributed GNN training on different graphs, number of GPUs, and number of GNN layers.

Graph / size (%)	arxiv	papers100M	products	reddit
$G_{1}$ / 10	0.56	0.50	0.30	0.32
$G_{2}$ / 30	0.48	0.50	0.23	0.28
$G_{3}$ / 50	0.46	0.49	0.21	0.27
$G_{4}$ / 70	0.46	0.49	0.20	0.27
$G_{5}$ / 90	0.45	0.49	0.20	0.26
$G$ / 100	0.45	0.49	0.21	0.26

Table 4. Average edge-cut ratio (lower is better) when training on different graph snapshots of

G

. The trained model is applied to the latest graph

G

without retraining.

	CPU				GPU
Strategy	arxiv	products	reddit	papers100M	arxiv	products	reddit	papers100M
Cuttana	1.029	1.149	1.110	0.951	1.001	2.052	1.026	OOM
Ldg	1.037	1.127	1.145	1.022	1.002	2.021	1.007	OOM
Metis	0.989	1.161	1.175	0.935	1.002	2.077	1.011	OOM
Spinner	1.003	1.127	1.136	0.955	1.000	2.018	1.011	OOM
EmbedPart-l-3	1.025	1.147	1.147	0.954	1.001	1.995	1.013	OOM

Table 5. Speedups through reordering for single-machine GNN training. Larger is better. OOM indicates out-of-memory errors.

	Average graph bandwidth				Edge-cut ratio
Strategy	arxiv	products	reddit	papers100M	arxiv	products	reddit	papers100M
EmbedPart-l-3	74K	1022K	154K	57594K	0.471	0.233	0.477	0.564
EmbedPart-l-3 -unbalanced	56K	659K	124K	38407K	0.360	0.162	0.372	0.286
EmbedPart-v-3	79K	1196K	151K	61468K	0.680	0.357	0.401	0.716
EmbedPart-v-3 -unbalanced	72K	875K	123K	44247K	0.619	0.269	0.293	0.271
EmbedPart-l-2	71K	1025K	154K	55970K	0.470	0.233	0.459	0.504

Table 6. Both the average graph bandwidth and edge-cut ratio are better (lower) for the unbalanced EmbedPart .

Explore similar work

SNI-GNN: SmartNIC-Assisted Full-Graph GNN Training with In-Network Embedding Prediction
Aug 6, 2026Guofan Yu, Sitian Chen, Zhenheng Tang +2Graph Neural Networks Communication-Efficient Distributed Training
Hong Kong Baptist University · Hong Kong University of Science and Technology · Hong Kong University of Science and Technology (Guangzhou)
Communication-free Sampling and 4D Hybrid Parallelism for Scalable Mini-batch GNN Training
Apr 3, 2026Cunyang Wei, Siddharth Singh, Aishwarya Sarkar +7Graph Neural Networks Distributed Training
University of Maryland, College Park, USA · NVIDIA, Inc. · Iowa State University, USA +3
FeLoG: Scalable and Efficient Distributed Graph Embedding with Feedback Loop Mechanism
Jun 20, 2026Peng Fang, Arijit Khan, Ziqiang Wu +4Communication-Efficient Distributed Training Graph Embedding
Huazhong University of Science and Technology · Bowling Green State University

Aug 6, 2026J/K move · Enter open · S save

SNI-GNN: SmartNIC-Assisted Full-Graph GNN Training with In-Network Embedding Prediction

Guofan Yu, Sitian Chen, Zhenheng Tang, Xiaowen Chu+1

Hong Kong Baptist University · Hong Kong University of Science and Technology · Hong Kong University of Science and Technology (Guangzhou)

Graph Neural Networks Communication-Efficient Distributed Training

Open paper AskPDF arXiv

Full-graph GNN training delivers high accuracy but scales poorly on multi-server clusters due to heavy, irregular inter-node embedding exchanges. We present SNI-GNN, a SmartNIC-assisted full-graph training system that reduces communication while preserving accuracy by predicting remote embeddings in-network. SNI-GNN deploys a lightweight linear-trend predictor on SmartNICs to refine cached historical embeddings, coupled with an importance-based boundary-node sampling policy and an asynchronous DPU--GPU data pipeline with intermediate-result reuse. We provide error and convergence bounds showing that predictor bias remains controlled under bounded second-order dynamics and yields standard non-convex convergence with inexact gradients. Implemented on NVIDIA BlueField-3, SNI-GNN integrates with state-of-the-art full-graph systems, cuts communication by 21--45%, achieves 1.3--3.6

\times

end-to-end speedups over BNS-GCN and up to 1.29

\times

over baseline SANCUS, with accuracy loss

\leq 0.01

, and scales efficiently to 16 GPUs on graphs with up to tens of millions of edges. These results indicate SmartNIC-based in-network prediction is a practical complement to partitioning and compression techniques for communication-efficient full-graph GNN training at scale.

Abstract

Explore similar work

SNI-GNN: SmartNIC-Assisted Full-Graph GNN Training with In-Network Embedding Prediction

Communication-free Sampling and 4D Hybrid Parallelism for Scalable Mini-batch GNN Training

FeLoG: Scalable and Efficient Distributed Graph Embedding with Feedback Loop Mechanism