SNI-GNN: SmartNIC-Assisted Full-Graph GNN Training with In-Network Embedding Prediction
TL;DR
SmartNIC-based in-network prediction is a practical complement to partitioning and compression techniques for communicationefficient full-graph GNN training at scale, and results indicate SmartNIC-based in-network prediction is a practical complement to partitioning and compression techniques for communicationefficient full-graph GNN training at scale.
Abstract
Full-graph GNN training delivers high accuracy but scales poorly on multi-server clusters due to heavy, irregular internode embedding exchanges. We present SNI-GNN, a SmartNICassisted full-graph training system that reduces communication while preserving accuracy by predicting remote embeddings innetwork. 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 intermediateresult reuse. We provide error and convergence bounds showing that predictor bias remains controlled under bounded secondorder 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 communicationefficient full-graph GNN training at scale.