LGNNIC: Acceleration of Large-Scale GNN Training using SmartNICs

LGNNIC: Acceleration of Large-Scale GNN Training using SmartNICs Liad Gerstman∗

Aditya Dhakal

Technion Israel Institute of Technology Haifa, Israel liadgerstman@campus.technion.ac.il

Hewlett Packard Enterprise Labs Fort Collins, CO, USA aditya.dhakal@hpe.com

Dejan Milojicic

Avi Mendelson

Hewlett Packard Enterprise Labs Milpitas, CA, USA dejan.milojicic@hpe.com

Technion Israel Institute of Technology Haifa, Israel mendlson@technion.ac.il

arXiv:2608.07733v1 [cs.DC] 7 Aug 2026

Abstract Graph Neural Networks (GNNs) are widely used across domains such as natural sciences, social network analysis, chip design, and recommendation systems. However, as graph sizes grow, storing and processing them entirely on a single-node CPU-GPU system becomes increasingly impractical. A promising approach is to distribute the graph across multiple remote memory nodes, though this introduces a major bottleneck: inter-node network congestion during training. To address this, we propose LGNNIC, a novel internode system architecture that leverages SmartNICs co-located with remote memory nodes—a configuration already available in modern systems—to reduce communication overhead in distributed GNN training. LGNNIC offloads key preprocessing tasks to SmartNICs, reducing the volume of data transferred to computational (training) nodes and alleviating network congestion. We introduce two complementary techniques executed on the SmartNICs during the preprocessing phase: Neighbor Sampling, which performs mini-batch sampling, and Quantization of the sampled batches. To evaluate LGNNIC under different communication infrastructures, we designed both an optimized low-overhead DMAbased synchronization mechanism and a high-overhead socketbased alternative used as a benchmark. We evaluate the core SmartNIC offloading mechanisms across standard GNN workloads and sampling hyperparameters using a proof-of-concept (PoC) system comprising one remote-memory node with an NVIDIA BlueField-2 SmartNIC and one compute node with an A100 GPU. Both Neighbor Sampling and Quantization on the remote node demonstrated substantial training speedups in most configurations. Neighbor Sampling achieved up to a 62.4x speedup with Sockets and 17.5x with DOCA-DMA on commonly used hyperparameters across datasets, primarily due to reductions in data transaction time (up to 73.6x and 5.1x, respectively). Quantization contributed additional speedups of up to 3.6x with Sockets and 1.3x with DOCA-DMA, also due to reduced data transfer (up to 3.6x and 2x, respectively).

1

Introduction

Graph Neural Networks (GNNs) have attracted interest from the research community for machine-learning-based relational systems. Several classes of GNNs have been studied: Graph Convolutional ∗ Part of the work was done while the author was an intern at Hewlett Packard Labs

Networks (GCN) [1], Graph Attention Networks (GAT) [2], GraphSAGE [3] and Graph Isomorphism Network (GIN) [4]. These GNN architectures differ in how they aggregate and transform information from neighboring nodes. GNN training produces node embeddings used for tasks such as node classification, graph classification, and link prediction. GNNs learn node embeddings by aggregating a node’s features and the features of its connected neighbors. Computing node embeddings is a core operation in GNN training and is performed recursively over each node and its neighbors. The challenges for GNN training arise from the exponential growth of the computational graphs as the graph size grows. With each GNN layer, neighbors of nodes from previous layers are aggregated, leading to exponential growth in computation. Due to the simultaneous calculation of node embeddings for the same layer across multiple computation graphs, it is necessary to keep the full graph’s data, along with intermediate data from the embedding calculation, on the GPU. For many graphs, e.g., the heterogeneous academic graph MAG240 (167 GB) [5] and the encyclopedic knowledge graph WikiKG90Mv2 (89 GB) [5], these raw graph sizes can easily exceed high-bandwidth memory (HBM) capacity of GPUs (40-100 GB) and their intermediate training data would require more memory. Consequently, it has become impractical to perform full-batch GNN training on a single CPU-GPU system for such datasets. A common solution involves partitioning the large graph into smaller mini-batches (during preprocessing) that can fit within the GPU’s memory. As a result of graph size explosion, sampling methods, which are typically performed on the CPU, are used (during mini-batch preparation) to reduce either the size of the graph or the mini-batches, allowing them to fit into GPU memory while maintaining a reasonable level of accuracy. An example of such a method is Neighbor Sampling, as introduced in GraphSAGE [3]. When scaling up to even larger GNNs beyond the point where they no longer fit even in host memory, sampling on the CPU in a single CPU-GPU node becomes infeasible. Distributing the graph across multiple memory nodes could serve as a potential solution. However, this approach introduces a significant communication bottleneck between the distributed memory and compute nodes (CPU and/or GPU), since the distribution usually requires transferring large volumes of data. To address the challenges of GNN learning on distributed graphs, we present LGNNIC (Large-Scale GNN Training using SmartNICs).

Liad Gerstman, Aditya Dhakal, Dejan Milojicic, and Avi Mendelson

Table 1: Attributes of the examined PyG datasets Dataset Reddit OGBNProducts OGBNMAG

Nodes 232,965

Edges 114,615,892

Features 602

Classes 41

2,449,029

61,859,140

100

47

1,939,743

21,111,007

128

349

Our proposed system consists of multiple remote memory nodes— memory and/or storage devices equipped with SmartNICs, a common configuration in modern systems—and multiple local computational nodes equipped with CPUs and GPUs. LGNNIC aims to alleviate communication bottlenecks through two complementary data manipulation methods: Neighbor Sampling and Tensor Quantization, both executed on the SmartNICs during the preprocessing (mini-batching) phase. These methods are executed before transferring the mini-batches to the computational nodes for the GPU training phase. We use SmartNICs to create graph samples from large graphs and we further quantize the mini-batch feature tensors before sending the sampled graph over the network. In this paper, we first present results from executing mini-batch GNN training with and without in-layer Neighbor Sampling, as introduced by GraphSAGE [3], on a local CPU-GPU system. This setup provides a baseline, motivating the subsequent profiling of different training phases executed across various components of the LGNNIC system. We then compare preprocessing times with and without Neighbor Sampling on a BlueField-2 SmartNIC (referred to as BF-2 hereafter) versus a local CPU. This comparison underscores the challenges of running preprocessing on a less powerful device near memory, while demonstrating the benefits of Neighbor Sampling in significantly reducing preprocessing time. The impact is further amplified by reduced data transfer times through remote Neighbor Sampling. Following this, we analyze end-to-end execution times for GNN training phases. Our analysis highlights the reductions in data transfer and preprocessing times on the LGNNIC proof-of-concept (PoC) system, resulting in an overall acceleration of the training process when Neighbor Sampling is applied. Our findings indicate up to an 8.7x speedup in data transfer time and up to a 17.5x speedup in total training time when using our DOCA-DMA-based [6] mechanism and greater speedups with the socket-based mechanism. In our quantization approach, we evaluate the effect of quantizing mini-batch feature tensors from 32-bit floating-point (FP32) to 16-bit floating-point (FP16) across all datasets and synchronization mechanisms. Quantization is performed on the SmartNIC of the remote node before transferring the tensors to the computational node. Upon arrival, the tensors are dequantized back to FP32 to continue computation. This technique reduces transfer overhead and improves training performance with small changes in test accuracy. Notably, quantization is handled entirely on the SmartNIC—not on the remote CPU—since, in our system (as in modern architectures), data flows directly from memory through the SmartNIC without requiring an additional copy that would be necessary if routed via

the CPU. Quantization achieves up to a 1.3x training speedup with DOCA-DMA and up to 3.6x with the socket-based mechanism. To summarize our key contributions: • We introduce LGNNIC, a novel architecture based on the BlueField-2 SmartNIC designed to accelerate total training time by mitigating network bottlenecks, thereby enhancing the performance of large-scale GNN algorithms. • We conduct an in-depth analysis of the effects of Neighbor Sampling in GNN training, comparing performance on a lower-performance compute device, such as a SmartNIC, with that on a high-performance device, such as the AMD EPYC 7513. • We profile the execution times of key GNN phases across a baseline CPU-GPU local node and across our LGNNIC system using three well-known GNN workloads. • We present a comprehensive implementation of a latencyoptimized synchronization mechanism that integrates PyG and DOCA-DMA between the remote and local nodes and compare it to our simple socket-based mechanism. • We present the impact of quantization as an additional method to reduce data transfers and accelerate total training time.

2

Background

In this section, we provide background on the libraries and techniques used in this work.

2.1

PyTorch Geometric (PyG)

Popular Python libraries for training GNNs include PyG (PyTorch Geometric [7]) and DGL (Deep Graph Library) [8]. We utilized PyG and its class NeighborLoader for training with mini-batch Neighbor Sampling. To establish the baseline mini-batch training sampling performance with a local node CPU-GPU system, we selected three medium-sized datasets (Reddit, OGBN-Products and OGBN-MAG). The datasets were specifically chosen because they are popular in various papers [3, 5, 9, 10] and because, as shown in Table 1, they exhibit varying characteristics in terms of nodes, edges, and features, allowing us to evaluate the mechanism across graphs with different characteristics. In addition, these datasets cover a range of use cases. The Reddit dataset is used to predict the community to which a new post belongs, based on the comments on its previous post [3]. The second dataset, OGBN-Products, involves predicting the category of an Amazon product based on products purchased together [11]. The third dataset, heterogeneous OGBN-MAG, focuses on predicting the conference or journal to which each paper belongs, based on its content, references, authors, and authors’ affiliations [11]. It is crucial to note that the selection of these datasets was influenced by the architecture of the proof-of-concept system, which consists of a single BlueField-2 SmartNIC, one CPU, and one GPU. In our research, we deliberately chose the largest, widely used realworld graphs that could fit within the memory constraints of these devices, enabling us to effectively demonstrate the benefits of the system.

LGNNIC: Acceleration of Large-Scale GNN Training using SmartNICs

2.2

Neighbor Sampling With PyG

With PyG, during each epoch, mini-batch sampling is performed in a batch-iterative manner by calling the next() function on the NeighborLoader iterable data loader, often named train_loader. The train_loader is instantiated with several arguments. Two of the most important arguments are: (1) num_neighbors—a list specifying the maximum number of neighbors to sample at each layer. For example, [25, 10] indicates a limitation of two layers: the first with up to 25 neighbors per node and the second with up to 10 neighbors per node, with any subsequent layers being truncated. In contrast, [−1, −1] also imposes a two-layer limit but without applying in-layer Neighbor Sampling. (2) num_workers— the number of subprocesses to run in parallel across multiple CPU cores using the CPU Affinity feature. To cover the full graph during mini-batch training, each PyG epoch iterates over the NeighborLoader class (a subclass of PyTorch’s DataLoader) to generate sampled mini-batches. For each mini-batch, the sampled nodes, features, and labels are transferred to the GPU, where training operations such as the forward pass, loss computation, and backpropagation occur. While preloading the full graph into GPU memory can reduce sampling overhead by avoiding repeated feature transfers, this is infeasible in scenarios where the GPU lacks sufficient memory for both data and intermediate computations—as is the case in our study. Our optimization of this workflow with DOCA-DMA is detailed in Algorithms 1 and 2. For completeness, we also provide the standard PyG neighbor-sampling training loop in Appendix A.1, which serves as a reference for the baseline algorithm that our system builds upon. PyG optimizes performance through a feature called CPU Affinity, which allows the sampling process to utilize a pre-defined number of CPU cores via the num_workers parameter. It binds worker subprocesses to CPU cores, with each worker maintaining its own data-loader and sampler state, enabling parallel batch preparation. However, since each worker operates in its own memory space, the sampler’s data must be replicated across these spaces. This replication can reduce performance and limit the number of workers that can be used, especially when the sampling device has limited memory. Therefore, it is crucial to carefully assess the benefits of CPU Affinity, ensuring that the overhead introduced by parallelism is minimal compared to the gains in sampling speed. As shown in Figure 4 and discussed later, this is not always the case.

2.3

Tensor Quantization

Quantization aims to reduce data precision in a way that minimizes quantization error and preserves model accuracy. It is widely used in neural network algorithms to lower memory consumption, reduce computational complexity, and improve inference speed by reducing the precision of activations and weights [12] [13] [14] [15]. Quantization techniques have also been successfully applied to GNNs [16][17][18]. In addition to its computational benefits, quantization can alleviate network bottlenecks in distributed systems by reducing the amount of data transmitted across the network.

3

Motivation

Our motivation for proposing LGNNIC arose from observed potential network bottlenecks in mini-batch sampling during GNN training with GraphSAGE using PyG. GraphSAGE mini-batch training with Neighbor Sampling on a local CPU-GPU system typically involves three main phases: (1) mini-batch preprocessing, which includes Neighbor Sampling; (2) transfer of the sampled labels, nodes, edges, and their features to the GPU; and (3) GPU-based model training. To establish a baseline for later comparison with a distributed setup—comprising remote and local (i.e., computational) nodes—we measured the execution time of these three phases for each of the workloads listed in Table 1, using GraphSAGE with Neighbor Sampling as described in Subsection 2.2. Mini-batch Neighbor Sampling was configured with varying numbers of layers and well-established hyperparameters tailored to each dataset. For Reddit, we used up to 25 neighbors in the first layer and 10 in the second. For OGBN-Products, we sampled up to 15 neighbors in the first layer, 10 in the second, and 5 in the third. For OGBN-MAG, the sampling configuration included 10 neighbors in both the first and second layers. In all cases, sampling was performed on the CPU of the local node, after which the sampled mini-batches were transferred to the local GPU over the PCIe interface. We used an AMD EPYC 7513 32-core processor and an NVIDIA A100 GPU. The PyG-based implementation of GraphSAGE minibatch sampling [3] restricts the sampling operation to CPU execution. To accelerate this phase, we leveraged the CPU affinity feature and employed six worker processes via multiprocessing. For profiling, we measured the average execution times of the three main stages in each mini-batch: (1) preprocessing, (2) data transfer following sampling, and (3) GPU training. Figure 1 compares the performance impact of using Neighbor Sampling with optimized, widely used hyperparameters (bottom) to a baseline without in-layer sampling (top) across all three phases. These hyperparameters were chosen because they are commonly used in the GNN community and yield only a minor reduction in test accuracy (as shown in Table 3 and discussed later), despite reducing the amount of data processed per batch. In the top plot, no in-layer Neighbor Sampling was applied for OGBN-Products and OGBN-MAG—i.e., while two layers were used, no neighbors were actually sampled in those layers. For the Reddit dataset, we applied no sampling in the first layer and limited sampling to 256 neighbors in the second layer to avoid an out-of-memory error due to the GPU’s limited memory capacity. Given the CPU’s high processing power, preprocessing time remains relatively low, while data transfer and GPU training account for the majority of the runtime. As shown in Figure 1, applying Neighbor Sampling significantly reduces execution time across all three training phases. This observation motivates our focus on reducing data transfer overhead by offloading Neighbor Sampling (together with Tensor Quantization) to a remote node in order to accelerate GNN training in distributed systems where communication becomes the primary bottleneck.

Liad Gerstman, Aditya Dhakal, Dejan Milojicic, and Avi Mendelson

(a) Mini-batch preprocessing without in-layer Neighbor Sampling

(b) Preprocessing with commonly used Neighbor Sampling hyperparameters.

Figure 1: Average execution times for key training phases on a baseline CPU-GPU node across all datasets.

4

LGNNIC System Architecture

We present our LGNNIC system architecture in Figure 2. Our LGNNIC solution aims to accelerate the total training execution time of GNN algorithms using SmartNICs as intermediate network infrastructure. Figure 2 contains two classes of nodes: compute nodes, which perform GNN training, and remote-memory nodes, which store graph data. We assume that the two are connected by highbandwidth Ethernet or another fabric such as InfiniBand. Each remote-memory node contains a BlueField-2 SmartNIC that accesses graph data from host memory or storage and performs minibatch sampling and quantization before transmission. In the compute nodes (shown on the left in Figure 2), GPUs such as the NVIDIA A100 or H100 reside alongside the NIC and CPU cores. Our LGNNIC architecture assumes the computational node’s NIC can perform peer-to-peer data transfers directly to the host memory or the GPU’s main memory. This peer-to-peer capability is available in many commercial NICs. In remote-memory nodes (shown on the right in Figure 2), we have BlueField-2 SmartNICs reading from host memory or storage to get the graph and performing mini-batch sampling and quantization before sending it over the network. As a result, they save transaction time by reducing the size of data being transferred. This approach mitigates network bottlenecks by distributing the full graph across multiple remote nodes and offloading mini-batch Neighbor Sampling together with Tensor Quantization closer to the remote memory nodes. In our architecture, to generate mini-batches that fit within the computational node’s GPU memory, the remote memory nodes collectively handle the full graph storage, with each node managing its assigned partition. These nodes execute the preprocessing phase— including both the sampling algorithm and Tensor Quantization— on their locally stored data. These preprocessed mini-batches are then transferred directly from the remote memory to the computational node’s GPU, where the remainder of the training phase is conducted. While the proposed architecture is designed to scale across multiple distributed nodes, the PoC evaluated in this paper isolates a single remote-memory-node/compute-node pair to evaluate the per-node SmartNIC offloading and communicationreduction mechanisms, as detailed in Section 5.

5

LGNNIC Evaluation

In this section, we begin by presenting the hardware setup and software implementations used to evaluate our proof-of-concept system, including a detailed overview of both the optimized and baseline synchronization mechanisms. We then discuss the challenges associated with offloading Neighbor Sampling from a powerful processor, such as the host CPU, to a lower-performance device like the BlueField-2 DPU. Despite these challenges, we demonstrate that offloading yields substantial benefits. Most notably, it results in a significant overall training speedup by reducing data transfers, which are a major bottleneck in high-overhead network environments. Finally, we show that Tensor Quantization results in relatively small changes in test accuracy across the evaluated datasets and hyperparameters, while yielding a marked improvement in overall training time by further decreasing the volume of transferred data.

5.1

Experimental Hardware

SmartNICs combine local memory with a programmable compute engine to process data in transit across data centers. BlueField-2 [19], NVIDIA’s DPU-based SmartNIC, offloads infrastructure services like networking and storage. It features 8 Armv8 A72 cores, 16 GB/32 GB DDR4 memory, a PCIe Gen 4.0 intra-node connection and supports up to 200 Gb/s connectivity via Ethernet/HDR InfiniBand. BlueField is programmed using NVIDIA’s DOCA SDK [20], which provides a low-level open API for accessing its services. Our focus in this paper is to evaluate the core SmartNIC offloading mechanisms of LGNNIC using a simplified proof-of-concept system. As shown in Figure 3, the PoC consists of one remote-memory node and one compute node and uses relatively large, widely used datasets that fit within the system’s memory constraints. This configuration isolates the effects of SmartNIC-based Neighbor Sampling, Tensor Quantization, and data-transfer optimization without introducing cross-partition sampling, remote-node coordination, or multi-node routing overheads. Evaluation of these multi-node scaling effects is left to future work. The proof-of-concept (PoC) setup used for our measurements, as shown in Figure 3, differs from the proposed system in Figure 2. The PoC uses a single compute and memory node to isolate and highlight the impact of the SmartNIC offloading techniques on

LGNNIC: Acceleration of Large-Scale GNN Training using SmartNICs

Figure 2: LGNNIC, our proposed SmartNIC-based system architecture, comprises multiple computational (local) nodes and remote memory nodes, each equipped with a SmartNIC. The graph is distributed across remote nodes, sampled and quantized on the SmartNICs before transfer.

Figure 3: Our proof-of-concept system was designed to measure the execution times of key training phases using either our DOCA-DMA mechanism or the Sockets benchmark. network bottlenecks. It consists of a BlueField-2 SmartNIC, an EPYC 7513 32-core CPU, and an A100 GPU. We use the A100X, a converged card integrating a BlueField-2 DPU and an NVIDIA A100 GPU, for both components. The A100X supports two modes: one where the GPU is visible to the host system, and another where it is accessible only to the SmartNIC’s OS. Note that while Figures 2 and 3 depict the SmartNIC and GPU as separate discrete blocks communicating over PCIe, these diagrams represent the logical data flow rather than the physical layout of the converged A100X card. By applying both offloading techniques, graph data is stored in the SmartNIC’s DRAM, and after in-network sampling during preprocessing, feature tensors are quantized to FP16. The sampled (and quantized) mini-batches are then transferred to the host CPU using our DOCA-DMA-based or socket-based synchronization mechanism. Finally, the mini-batches are passed to the GPU for training, including both forward and backward passes.

5.2

Software Implementation

To demonstrate the benefits of the architecture, we developed an optimized synchronization mechanism for executing mini-batch training with Neighbor Sampling across the PoC system. This mechanism transfers sampled mini-batches from the remote node to the

computational node, managing data buffers on both ends and enabling remote-side computation during transfers. The goal is to achieve low-overhead, high-speed data transfers via DMA, while avoiding unnecessary OS-level memory copies. To this end, our implementation leverages the low-level, onedirectional, C-based DOCA-DMA API [6], which enables DMAbased transfers between host and BlueField-2 DPU buffers over PCIe. The mechanism avoids socket and network-stack overheads; however, it is not fully zero-copy in the current implementation. Because the DOCA-DMA buffer size is fixed and neither the buffer nor its underlying memory pointer can be changed without reinitializing the DOCA process, received chunks must be copied once from the host-side DOCA buffer into a larger local aggregation buffer. Thus, each transferred chunk incurs a single host-side memory copy before reconstruction of the corresponding tensor. As described later in the paper, we implemented a generalized version of this API that is integrated with PyG to support a bidirectional synchronization protocol between the host (computational node) and the DPU (remote node). This protocol efficiently transfers large tensor buffers using acknowledgment (ACK) messages to ensure correctness and ordering. Integrating the C-based DOCA API with PyG’s Python-based sampling required careful coordination. To maintain buffer synchronization and performance, each sampling batch N had to be

Liad Gerstman, Aditya Dhakal, Dejan Milojicic, and Avi Mendelson

Algorithm 1 BF-2 Side (Mem. node) of DOCA-DMA + PyG Training Pipeline

Algorithm 2 Host Side (Comp. node) of DOCA-DMA + PyG Training Pipeline

1: Initialize: DOCA Source & Destination Buffers 2: Initialize: Python & Python Objects 3: Initialize: Dataset & train_loader Configuration 4: for epoch = 1 to Epoch Number do 5: while batch=next(train_loader) do 6: for Each Batch Element Tensor do 7: Quantize: Tensor 8: Convert: Batch Elements to Bytes 9: for Each Chunk In Element Tensor do 10: Memcopy: Data to DOCA Buffer By Order 11: Wait: For ACK From Host 12: Send: Finish Message To Host

1: Initialize: DOCA Source & Destination Buffers 2: Initialize: Python & Python Objects 3: Initialize: Model 4: Initialize: Optimizer 5: for epoch = 1 to Epoch Number do 6: while Epoch_Finished_Message is False do 7: for Each Batch Element Tensor do 8: for Each Chunk In Element Tensor do 9: Wait: For New Chunk From BlueField-2 10: Memcopy: DOCA Buffer to Local Buffer 11: Send: ACK To BF-2 12: PyMemoryView: Local Buffer To Tensor 13: Dequantize: Tensor 14: if Epoch Finished Message Was Received then 15: Assign: Epoch_Finished_Message = True 16: Break 17: Transfer: Mini-Batch Data to GPU 18: Perform: Forward Pass Through The Model 19: Calculate: Loss Using Predicted and Actual Labels 20: Calculate: Back-Propagate Gradients 21: Update: Model Parameters Using Optimizer

⊲ i.e., Packet ⊲ Optional

coordinated with its predecessor N-1, ensuring data from minibatch N wouldn’t overwrite buffer contents still in use from N-1. Since DOCA buffers are limited to 1 MB, tensor transfers occur in chunks, requiring an ACK-based mechanism to prevent overwriting in-progress data transfers. As we further analyze in Subsection 6.2, this hardware limitation acts as a significant performance bottleneck, which directly necessitated the design of our complex chunking and synchronization protocol. The ACK−based synchronization mechanism sends each chunk with a header comprising the following information: (1) the packet number, (2) the tensor’s number of bytes, (3) the number of the current chunk, and (4) the number of bytes in the current chunk (which varies between minibatches and tensors). A packet is defined as the combination of all the chunks needed to transfer one tensor of data related to a specific sampled batch. For instance, when transferring nodes’ feature tensor to the local node, a packet comprises all the chunks needed to transfer the feature tensor. Both the Reddit and OGBN-Products datasets are homogeneous, meaning they consist of only one kind of node. In addition, neither dataset includes edge features. Hence, in their case, only three packets were sent for each batch: one for the node-feature tensor, one for the batch-node labels, and one for the edges’ adjacency matrix tensor, each composed of multiple chunks. However, OGBN-MAG is a heterogeneous graph, meaning it consists of more than one kind of node and edge. In addition, it does not include edge features. Hence, in OGBN-MAG’s case, for each batch, twelve packets were sent with the following tensors: paper node features, paper nodes labels, author node features, institution node features, field_of_study features, (author, affiliated_with, institution) adjacency matrix, (author, writes, paper) adjacency matrix, (paper, cites, paper) adjacency matrix, (paper, has_topic, field_of_study) adjacency matrix, (institution, rev_affiliated_with, author) adjacency matrix, (paper, rev_writes, author) adjacency matrix, and (field_of_study, rev_has_topic, paper) adjacency matrix. To prevent overwriting the current chunk in the local (computational) node’s DOCA buffer, the remote node waits for an ACK from

⊲ i.e., Packet

⊲ Optional

the local node confirming the data has been copied to a larger aggregation buffer, where all incoming tensor chunks are aggregated. Once the data is copied from its local DOCA buffer, the remote node proceeds without waiting—either handling the next chunk or sampling the next batch if the last chunk was already sent. The local node identifies new data in its DOCA buffer by detecting a change in the "current chunk number" field in the header, which is the last field to be updated by the remote node. The other header parameters are used later during training for tensor-related calculations. The synchronization flow is shown in Algorithm 1 for the BlueField-2 remote node and Algorithm 2 for the local node. A corresponding synchronization mechanism was also built as a benchmark with sockets to demonstrate the importance of using a low-overhead DMA-based mechanism with faster transactions and to underscore the significance of sampling, particularly when using a high-latency protocol such as Ethernet.

5.3

Offloading Neighbor Sampling to the SmartNIC

5.3.1 Performance Comparison of Neighbor Sampling on Host CPU vs. BlueField-2 DPU. To evaluate the computational disparities between our BlueField-2 SmartNIC and the powerful local node EPYC CPU—two devices whose performance difference reflects realistic system architectures—we compared the execution time of preprocessing two layers in the computational trees with and without in-layer Neighbor Sampling (and other configurations) on the host CPU versus the remote BlueField-2 node. Quantization was not applied in this part of the evaluation. The measurements were taken as a function of the number of workers for each dataset: Reddit, OGBN-Products, and OGBN-MAG, as illustrated in Figure 4. The mini-batch preprocessing phase (top-left subfigures in Figures 4a, 4b, 4c) entails constructing two-layer computational trees for each node in the mini-batch without applying sampling at these layers (i.e., no in-layer sampling). Additionally, for each dataset, three mini-batch preprocessing steps including extensive sampling with varying hyperparameters were evaluated. For all datasets, not

LGNNIC: Acceleration of Large-Scale GNN Training using SmartNICs

(a) Reddit

(b) OGBN-Products

(c) OGBN-MAG

Figure 4: Impact of in-layer Neighbor Sampling on preprocessing performance, comparing execution on the host CPU versus the BlueField-2 SmartNIC across the three datasets under various sampling hyperparameters. only were the commonly used hyperparameters employed, but additional ones were tested to explore the effects of Neighbor Sampling with varying numbers of nodes and layers. Each hyperparameter is enclosed in brackets, [], where each number represents the maximum number of nodes to sample at a given layer. Reading from left to right, the first number corresponds to the sampling size for the first neighbor hop from the root in its computational tree, continuing sequentially for subsequent layers. If L numbers are specified, this indicates that the computational tree consists of L layers (hops). When a value of -1 is used, it indicates that no sampling was performed at that layer, meaning all neighbors in that layer were retained. For the Reddit dataset, the following sampling hyperparameters were selected and are shown in Figure 4a: [-1, -1] (no in-layer sampling, top left), [25, 10] (commonly used hyperparameters, top right), [10, 10] (more extensive sampling, bottom left), and [-1] (a single layer with no in-layer sampling, bottom right). For the OGBN-Products dataset, the sampling hyperparameters were as follows and are presented in Figure 4b: [-1, -1] (no in-layer sampling, top left), [25, 10] (more extensive sampling, top right), [15, 10, 5] (commonly used hyperparameters, bottom left), and [-1] (a single layer with no in-layer sampling, bottom right). For the OGBN-MAG dataset, the selected sampling hyperparameters, shown in Figure 4c, were: [-1, -1] (no in-layer sampling, top left), [25, 10] (more extensive sampling, top right), [10, 10] (commonly used hyperparameters, bottom left), and [-1] (a single layer with no in-layer sampling, bottom right). Specifically, on the SmartNIC, configurations combining CPU Affinity (num_workers > 0) with unconstrained two-hop Neighbor Sampling (num_neighbors = [−1, −1], i.e., all neighbors were retained at both hops) exhausted the available DRAM capacity across all evaluated datasets. Memory monitoring showed a rapid increase in memory consumption until the device’s available memory was depleted. Because each worker executes in a separate process and replicates the data loader and sampler state, multi-worker execution substantially increases the aggregate memory footprint on the BlueField-2 (Figures 4a, 4b, and 4c). Notably, increasing the number of unconstrained neighborhood-expansion layers exhausted the

available SmartNIC memory even with num_workers = 0. The corresponding runs did not complete successfully and were therefore excluded from the performance figures. Increasing the number of workers did not improve preprocessing performance on the SmartNIC and degraded it in several configurations. Even under the most favorable configurations with extensive sampling, preprocessing on the BlueField-2 remained slower than on the EPYC 7513 CPU because of its lower computational capability and more limited memory resources. With num_workers = 0, BlueField-2 Neighbor Sampling was approximately 2–10× slower than sampling on the EPYC 7513. Nevertheless, as evidenced by Figures 4a, 4b, and 4c, extensive sampling significantly reduces preprocessing time for both the CPU and SmartNIC compared to the no-in-layer scenario, especially for the SmartNIC when CPU Affinity is disabled, as it reduces memory access demands. Therefore, as demonstrated in Figures 4a, 4b, and 4c, since preprocessing on the SmartNIC without CPU Affinity provided superior performance, we disabled this feature (num_workers = 0). As shown in Figure 4, the SmartNIC demonstrates a significant preprocessing time reduction when utilizing in-layer Neighbor Sampling (comparing the top-left subfigures to the commonly used hyperparameters in each dataset with num_workers = 0). For example, with Reddit (Figure 4a), BlueField’s preprocessing time dropped from ∼ 17 seconds to ∼ 0.3 seconds. Other datasets show similar results. 5.3.2 Overall Training Speedup through Offloaded Neighbor Sampling. While Figure 4 highlights the difficulty of sampling on BlueField-2, Figure 5 shows that network bottlenecks can dominate execution time, making data transfer time (from SmartNIC to host) a key factor. As shown later, performing sampling on the SmartNIC can significantly reduce transfer time—even when sampling itself is slow—underscoring the SmartNIC’s value in accelerating overall training. As illustrated in Figure 5, performing remote preprocessing with sampling not only decreases sampling time but also significantly reduces the transfer time for the preprocessed batches. To showcase the benefits of the architecture, we implemented a DOCA-DMAbased synchronization mechanism to transfer the preprocessed

Liad Gerstman, Aditya Dhakal, Dejan Milojicic, and Avi Mendelson

(a) Reddit

(b) OGBN-Products

(c) OGBN-MAG

Figure 5: Average execution time comparison across all datasets. Key mini-batch training phases are analyzed using various sampling hyperparameters with our mechanisms to demonstrate the impact of Neighbor Sampling. batches from the remote node to the local node. This mechanism enables low-overhead, high-speed transactions via DMA, eliminating unnecessary OS memory copies and allowing concurrent preprocessing of the next batch while the current one is in transit. For comparison, we also developed a socket-based synchronization mechanism to expose the network bottleneck that arises when using Ethernet instead of PCIe. Figure 5 contrasts the performance of the two synchronization mechanisms—DOCA-DMA (lower subfigure in Figures 5a, 5b, and 5c) and Sockets (upper subfigure in Figures 5a, 5b, and 5c)—on our proof-of-concept system. The subfigures present execution time measurements for various training phases across all datasets using different sampling hyperparameters mentioned in each subfigure’s legend. Notably, for Reddit, we selected the [-1, 256] configuration (which represents no sampling in the first layer and sampling up to 256 neighbors in the second layer) instead of [-1, -1] due to the memory constraints of the A100 GPU (80 GB), as removing the neighbor limit would exceed GPU capacity. The [-1, 256] setting provided the minimal sampling necessary to stay within the GPU’s memory limits. For each dataset, we measured three training phases: preprocessing on the SmartNIC, transfer of feature, label, and adjacency matrix tensors (at least one each) from the SmartNIC to the local node, and training on the local node. Each transaction consists of a single packet divided into multiple chunks (each chunk includes the relevant tensor bytes and four header fields). This segmentation is due to the 1 MB buffer limit of the DOCA-DMA library, and for consistency, we set the Sockets mechanism’s buffer size to 1 MB as well, which limits the performance of the Sockets mechanism. Using our synchronization mechanisms, tensor chunks are transferred from the SmartNIC to the local CPU for GPU training. Since each dataset consists of multiple tensors representing different features, labels, and adjacency matrices, only non-negligible data transfer times are shown in Figure 5 due to space constraints. However, every individual tensor transfer time was measured and included in the calculation of both the total training time speedup and total transfer time speedup, as shown in Table 2. As illustrated in Figure 5, there is a significant reduction in the execution times of the three phases across all datasets. For example, for the Reddit dataset with the Sockets mechanism, the average

feature tensor transfer time decreases from 174.072 seconds (for the [-1, 256] hyperparameter) to 36.537 seconds (for the commonly used hyperparameter [25, 10]). Other datasets show similar results. Figure 5 shows that as sampling becomes more extensive, transfer times drop significantly due to reduced data volume. However, while sampling lowers preprocessing time, it is not negligible and can become a bottleneck as transfer costs diminish. As shown in Table 1, the Reddit dataset has far more edges—about 1.85x more than OGBN-Products and 5.5x more than OGBN-MAG—despite having far fewer nodes (roughly 10.5x fewer than OGBN-Products and 8.3x fewer than OGBN-MAG). Thus, sampling batches of 1024 nodes yields longer transfer times for Reddit, as each computational tree includes more nodes on average. This results in extended transfer times for Reddit’s feature and adjacency tensors under identical sampling hyperparameters for both DOCA-DMA and Socket mechanisms. Additionally, the commonly used OGBN-Products hyperparameter ([15, 10, 5]) employs three sampling layers, causing higher preprocessing time than other extensive sampling settings—both within OGBN-Products and compared to other datasets. To gain deeper insights into the results, we calculated the speedup of extensive sampling—using various hyperparameters, including those commonly used to achieve optimal accuracy—compared to the baseline without in-layer sampling (i.e., [-1, -1] or [-1, 256]) for each dataset. The speedup in transaction time and total training time for both DOCA-DMA and Sockets is substantial across all datasets and hyperparameters, as shown in Table 2. For each dataset, the transaction speedup was calculated by dividing the average total transfer time for batches with minimal sampling ([-1, 256] or [-1, -1], depending on the dataset) by the average total transfer time for batches with extensive sampling ([25, 10], [10, 10]/[15, 10, 5], [-1]). The results for each dataset are illustrated in Table 2. Total training time speedup was calculated in a similar manner, taking into account preprocessing and training phases. The results, summarized in Table 2 alongside Figure 5, highlight that graph preprocessing through sampling can mitigate network bottlenecks and significantly reduce data transfer times across all datasets. Since not all hyperparameter combinations are applicable to all datasets, the corresponding cells are left empty and marked with a ‘-’ sign. Figure 5 shows that DOCA-DMA consistently outperforms the socket-based mechanism across all datasets, primarily due to its use

LGNNIC: Acceleration of Large-Scale GNN Training using SmartNICs

Num Neigh. [25, 10] [10, 10] [15, 10, 5] [-1]

Sockets Trans. Speedup Total Speedup Reddit Products MAG Reddit Products MAG 8.56 84.49 52.64 8.76 75.12 45.51 21.32 – 73.57 21.73 – 62.39 – 19.53 – – 18.48 – 5.19 118.12 472.08 5.33 107.17 353.57

DOCA-DMA Trans. Speedup Total Speedup Reddit Products MAG Reddit Products MAG 5.10 15.36 7.44 17.46 18.43 7.45 7.99 – 8.73 29.03 – 8.61 – 6.45 – – 7.40 – 3.88 18.18 47.65 15.20 28.89 39.40

Table 2: Total transaction and training speedups under different sampling hyperparameters and synchronization mechanisms. Highlighted values correspond to the most common configuration per dataset.

of PCIe instead of Ethernet. Table 2 further highlights this contrast: while DOCA-DMA is faster overall, the relative transaction speedup is greater for the socket-based method, since Ethernet’s higher overhead makes reductions in data transfer time more impactful. For Reddit, the reported speedup is smaller because it is measured relative to the minimal sampling setting ([-1, 256]) rather than a potential full no in-layer sampling configuration ([-1, -1]), as discussed earlier. Additionally, for the Sockets mechanism, when data transfers are the dominant bottleneck, as seen in the Reddit Sockets measurements, the total time acceleration closely aligns with the transaction speedup. In contrast, when preprocessing and training phases significantly contribute to the overall training time, the influence of transaction speedup on total training time acceleration diminishes, as seen with Reddit’s DOCA-DMA results. In such cases, optimizing the preprocessing phase becomes crucial. As shown in Figure 4, employing a more powerful CPU can substantially lower preprocessing times, highlighting the need for further exploration of hardware-based optimizations. Regardless of the underlying reason, reducing the preprocessing or training phases through sampling, or minimizing transaction times due to sampling, ultimately results in a shorter total execution time. In summary, by synthesizing the findings from Figure 5 and Table 2, we demonstrate that remote sampling effectively reduces transaction times, alleviates network bottlenecks, and thereby enhances overall system performance.

5.4

Tensor Quantization on the SmartNIC

As shown in Figure 5, transaction time increases with the amount of data transferred over the network. In the previous section, we demonstrated that offloading Neighbor Sampling to the SmartNIC significantly accelerates data transfers by reducing the data volume. Similarly, quantization can further reduce data size and is a lightweight task for the SmartNIC, whether implemented in software or hardware. As illustrated in Figure 5, the feature tensor consistently emerges as a primary bottleneck, accounting for a substantial portion of the transfer time across all datasets, sampling hyperparameters, and synchronization mechanisms. To address this, we utilize the SmartNIC to quantize the feature tensor from FP32 to FP16 in software, thereby reducing its size, lowering transfer time, and shortening total training time (see Table 4). 5.4.1 Impact of Feature Tensor Quantization on Test Accuracy. While quantization reduces tensor size, it can also introduce error that may degrade test accuracy. Since FP32 is the default precision in many

Num Neighbors [-1, -1] [-1, 256] [15, 10, 5] [25, 10] [10, 10] [-1]

Reddit FP32 FP16 – – 0.951 0.951 – – 0.952 0.952 0.952 0.952 0.943 0.938

OGBNProducts FP32 FP16 0.777 0.774 – – 0.789 0.790 0.771 0.776 – – 0.685 0.673

OGBNMAG FP32 FP16 0.445 0.444 – – – – 0.444 0.447 0.453 0.451 0.394 0.381

Table 3: Average test accuracy results across all datasets and sampling hyperparameters with the feature tensor quantized from FP32 to FP16. Highlighted values represent common configurations per dataset.

PyG datasets and models, FP16 offers a practical trade-off by minimizing quantization error while enabling considerable speedups. As shown in Table 3, this precision reduction leads to only minor accuracy degradation, if any, while significantly accelerating data movement and improving overall training time as illustrated in Table 4. In Table 3, we report the average test accuracy for each dataset and each Num Neighbors hyperparameter (with commonly used values highlighted in yellow), comparing FP32 and FP16 precisions. Since not all hyperparameter combinations are applicable to all datasets, the corresponding cells are left empty and marked with a ‘-’ sign. As shown, the accuracy degradation from quantization is relatively small across the evaluated configurations and well justified by the resulting performance gains as illustrated in Table 4. Each training run was repeated four times. During each run, test accuracy was evaluated using the model configuration that achieved the best validation accuracy, computed over multiple epochs until convergence. The final reported test accuracy represents the average across all four runs. 5.4.2 Overall Training Speedups through Offloading Tensor Quantization. In Table 4, we compare FP32 and FP16 precision when quantizing the feature tensor to highlight the impact of quantization on total transfer time and, consequently, total training (execution) time. Average total transaction time, average total execution time, and the resulting speedups were calculated across all mechanisms (DOCADMA and Sockets), datasets, and sampling hyperparameters (with common hyperparameters highlighted in yellow). Regarding Table 4, total transaction time is defined as the sum of all data transfers from the SmartNIC to the host across all mini-batches, averaged over multiple runs. Total execution time is the sum of the three

Liad Gerstman, Aditya Dhakal, Dejan Milojicic, and Avi Mendelson

main phases: (1) preprocessing, (2) all data transfers from the SmartNIC to the host, and (3) GPU-based training, also averaged across multiple runs. Speedup values for transfer and training times were computed by dividing FP32 results by their corresponding FP16 results and are shown in Table 4. The different hyperparameters used in the measurements affect the size of the feature tensor transferred across the network and thus allow for evaluation under a range of data transfer conditions. Hence, since speedups are computed with the sampling hyperparameter held constant (and only the precision varied), the relative gains are solely attributed to the change in precision. This differs from Table 2, where speedups were computed for each FP32 configuration relative to the baseline hyperparameter with minimal sampling (i.e., [-1, 256] or [-1, -1], depending on the dataset). We evaluate performance under both DOCA-DMA and Sockets synchronization mechanisms to assess how different infrastructures and communication protocols impact end-to-end training. As shown in Table 4, quantization consistently delivers significant speedups in both total transfer time and total training time across all datasets, sampling configurations, and mechanisms. Notably, in settings where transaction time dominates overall training time— such as with the high-overhead Sockets mechanism—the execution speedup closely tracks the transfer speedup. This effect is particularly pronounced in the Reddit dataset, which has the highest edge-to-node ratio among those evaluated. Detailed results are provided in Appendix A.2 (Table 5). The high edge-to-node ratio also explains why Reddit exhibits the highest total training time among the evaluated datasets with Sockets, while preprocessing time remains moderate. In contrast, when transaction time constitutes a smaller fraction of total training time—as is typically the case with the low-overhead DOCA-DMA mechanism, where preprocessing is also a significant bottleneck— the execution and transfer speedups diverge. Since DOCA-DMA reduces communication overhead substantially, the impact of quantization is less pronounced than with Sockets, where communication overhead remains the primary bottleneck. In summary, this section demonstrates that quantizing the feature tensor to FP16 is an effective strategy for accelerating training while incurring minimal test accuracy degradation. Even though our quantization was deliberately simple, the SmartNIC is a natural platform for deploying more sophisticated hardware- and softwarebased quantization techniques. Quantizing additional tensors—such as labels and adjacency matrices—and using lower-precision formats may yield even greater speedups. Our goal was to showcase the potential of this approach to accelerate overall training, which was clearly demonstrated. As with the sampling-based method discussed earlier, the benefits of quantization are most pronounced for datasets with high edge-to-node ratios and in systems with significant communication overhead.

6

6.1

High-overhead Networks

It is important to highlight our expectations regarding larger graphs (relative to our evaluated datasets) that demonstrate good accuracy with robust sampling. In particular, those with a high edge-tonode ratio are likely to show even greater acceleration in transfer times with Neighbor Sampling offloaded to the SmartNIC. This is especially true when data is transferred across high-overhead networks such as Ethernet or InfiniBand. This is attributed to the fact that larger graphs, when sampled with significant in-layer sampling hyperparameters (e.g., [25, 10]), can substantially reduce transaction times by minimizing the number of chunks transmitted over the network without significantly compromising accuracy, as demonstrated in this paper. Furthermore, quantizing such graphs prior to transmission can lead to additional reductions in total training time. Specifically, when large graphs with high edge-to-node ratios—and relatively low preprocessing overhead—are quantized on the SmartNIC before being transferred over high-overhead networks, the resulting speedups are particularly pronounced. We anticipate that further optimizing preprocessing latency on the SmartNIC—by quantizing additional tensors (e.g., label and adjacency matrix tensors) and employing lower-precision formats—will yield even greater performance gains.

6.2

NVIDIA’s BlueField-2 and DOCA

As discussed, PyG supports CPU Affinity, which binds specific cores to preprocessing tasks and allocates dedicated memory per core to enable parallel batch preprocessing. However, as shown in Figure 4, this strategy degraded performance on the SmartNIC due to limited memory, since each worker replicates the data loader’s memory requirements. Consequently, we disabled CPU Affinity during training and limited experiments to small- and medium-sized graphs with few layers—a common setup, since additional layers rarely improve accuracy. Additionally, the SmartNIC’s preprocessing was 2-10x slower than the CPU due to its lower computational power. To address this, we propose increasing BlueField-2’s core count, optimizing individual cores, and expanding memory to reduce preprocessing bottlenecks. As discussed earlier, we identified a 1 MB limitation in the DOCADMA buffer size, which prevented the DOCA-DMA synchronization mechanism from reaching its full performance potential. Additionally, the DOCA-DMA buffers could not be reallocated without reinitializing the DOCA process, which is time-consuming. This constraint necessitated allocating a separate local buffer and copying data from the DOCA buffer, resulting in a performance penalty. Implementing dynamic reallocation for DOCA buffers would allow devices to work directly with DOCA buffers with zero-copy. Therefore, we propose that allowing the buffer size to be increased, along with supporting the reallocation of both remote and local DOCA buffers, could significantly optimize DOCA-DMA data transfers by avoiding an extra memory copy for each chunk.

Architectural Improvements

Working with BlueField-2 has provided valuable insights into potential architectural improvements. This section presents a selection of the most important ones.

7

Related Work

Unlike prior work, which generally limits SmartNIC offloading to lightweight or stateless operations, LGNNIC offloads mini-batch

LGNNIC: Acceleration of Large-Scale GNN Training using SmartNICs

Sockets

Num Neighbors [−1, 256] [25, 10] [10, 10] [−1] [−1, 256]

Trans. Speedup 1.64 3.64 3.60 3.65 1.24

Exec. Speedup 1.62 3.56 3.48 3.59 1.03

DOCA -DMA

[25, 10]

1.36

1.20

[10, 10] [−1]

1.24 1.37

1.15 1.24

REDDIT

OGBNPRODUCTS Sockets

Num Neighbors [−1, −1] [15, 10, 5] [25, 10] [−1] [−1, −1]

Trans. Speedup 2.25 3.39 3.63 4.53 1.24

Exec. Speedup 2.16 2.84 2.71 3.31 1.02

DOCA -DMA

[15, 10, 5]

1.41

1.04

[25, 10] [−1]

1.18 2.02

1.03 1.11

OGBNMAG Sockets

Num Neighbors [−1, −1] [25, 10] [10, 10] [−1] [−1, −1]

Trans. Speedup 3.03 2.16 1.96 1.20 1.33

Exec. Speedup 2.95 2.01 1.83 1.20 1.13

DOCA -DMA

[25, 10]

1.93

1.25

[10, 10] [−1]

1.96 1.12

1.31 1.05

Table 4: Speedups achieved using FP32 vs. FP16 precision: total transaction and execution time improvements across all mechanisms, datasets, and sampling hyperparameters. For more data, see Appendix A.2.

Neighbor Sampling—a non-standard, compute-intensive operation— to the SmartNIC, achieving substantial reductions in data movement. We also offload Tensor Quantization, a more conventional optimization, to further reduce transfer overhead.

7.1

Sampling Large Graphs

When generating GNN embeddings, the computational graphs grow exponentially as more layers and neighbors per layer are added. Consequently, when handling large graphs, a preprocessing phase known as sampling becomes essential. Sampling, a critical component of the preprocessing phase, reduces the size of the minibatches’ computational graphs. These sampled mini-batches are then fed into the GPU for training. Various sampling categories exist, including node-wise (e.g., GraphSAGE [3], VR-GC [21], PinSAGE Sampler [22]), layer-wise (e.g., FastGCN [23], AS-GCN [24], LADIES [25]), and subgraph-wise (e.g., GraphSAINT [26], RWT [27]). Among node-wise methods, the most widely used approach is Neighbor Sampling, introduced by Hamilton et al. [3]. The Neighbor Sampling algorithm limits the number of layers in the computational trees and constrains the number of nodes per layer, thereby creating smaller computational trees for mini-batches. Within the concept of Neighbor Sampling, various strategies for selecting nodes to sample exist. These include random sampling, which may unintentionally prune important nodes, and random walk with restarts (RWR) [28], which ensures the inclusion of significant nodes. The Neighbor Sampling algorithm, which is used in our system, iteratively samples nodes based on the number of GNN layers and a specified maximum number of neighbors per layer. It is important to note that even after sampling, the size of the computational graph still grows exponentially with the number of layers. To ensure stable training with minimal variance in neighbor aggregation—and to prevent a significant drop in accuracy—the average number of neighbors per node after sampling must not decrease substantially. Therefore, overly aggressive sampling of the computational graph is not feasible; it is crucial to strike a careful balance between meeting the memory constraints of GPUs or CPUs and maintaining high accuracy. Other large-graph sampling methods exist to enable large-scale GNN training. For example, Advanced Cluster-GCN [9] constructs small sets of node groups that are aggregated to form larger node

groups, offering better resource utilization compared to Neighbor Sampling. Another example is the Simplified GNN [29] architecture, where a linear matrix is used instead of computing nonlinear activations during the forward pass, enhancing scalability at the expense of expressiveness.

7.2

Communication Overhead

In Farview [30], the authors alleviate data-transfer bottlenecks with a remote buffer cache that is capable of processing data streams before being transmitted over the network. It supports a few defined offloading operators, specifically for SQL queries, using RoCE v.2 RDMA transfers. In our research, as we also explore the potential of offloading to SmartNICs, we focus on evaluating the impact of rebalancing the entire algorithm across the compute and memory components of the training process, rather than simple SQL operations and intra-node optimizations. The DGCL paper [31] also addresses communication overheads by distributing GNN training across multiple GPUs rather than across remote memory nodes. The Sequential Aggregation and Rematerialization paper [32] addresses smart graph partitioning using aggregation and rematerialization, freeing memory across multiple processing workers during the backward pass, enabling the scaling of large GNNs. Other works deal with GNN acceleration on a single computational node (Memory-CPU-GPU and not inter-node networks). SmartSAGE [10], e.g., deals with accelerating the same GraphSAGE [3] sampling algorithm implemented based on the same library (PyG [7]) we use, with an in-storage architecture exploiting a smartSSD composed of NVMe SSD and an FPGA. The GNNear paper [33] also focuses on near-memory processing to accelerate memory-intensive reduce & update operations. There are other works focusing on accelerating the sampling phase on CPU-FPGA systems, including pure sampling acceleration on FPGAs [34, 35] and on multi-FPGA platforms [36]. While prior work mostly focuses on quantizing activations and weights [12–18], our approach quantizes data before transmission and dequantizes it at the receiver prior to training.

8

Conclusions and Next Steps

We introduced LGNNIC, a novel architecture for accelerating largescale GNN training using SmartNICs. LGNNIC offloads Neighbor Sampling and Quantization to the SmartNIC, reducing data transfer

Liad Gerstman, Aditya Dhakal, Dejan Milojicic, and Avi Mendelson

overhead and improving training time. To evaluate this architecture on our PoC system, we developed a low-overhead DOCA-DMA synchronization mechanism for efficient buffer management and compared it to a high-overhead socket-based benchmark. Offloading Neighbor Sampling consistently yielded end-to-end training speedups, particularly in transaction time, while quantization reduced transfer time, especially under network congestion. Even modest FP32-to-FP16 quantization provided substantial benefits, and additional gains are possible by quantizing more tensors or using lower precisions. Performing these operations on the SmartNIC avoids unnecessary data movements between remote and host, making it a natural and efficient location for this task. The results across three datasets, multiple sampling configurations, and two communication mechanisms demonstrate the feasibility of LGNNIC’s per-node SmartNIC offloading approach. Evaluating LGNNIC in a full multi-node deployment, including cross-partition sampling, inter-node coordination, routing, and scalability, is left to future work. Moreover, we plan to explore peerto-peer DOCA-RDMA over Ethernet for higher-latency distributed setups and for zero-copy transfers from remote SmartNIC DRAM to the local GPU. We will profile the system, identify bottlenecks, and assess additional acceleration on BlueField-2, as well as evaluate more advanced BlueField-3 and BlueField-4 devices, which offer greater memory capacity and bandwidth and may further reduce preprocessing and communication overhead, thereby improving overall training performance.

References [1] T. N. Kipf and M. Welling, “Semi-supervised classification with graph convolutional networks,” 2017. [2] P. Veličković, G. Cucurull, A. Casanova, A. Romero, P. Liò, and Y. Bengio, “Graph attention networks,” 2018. [3] W. Hamilton, Z. Ying, and J. Leskovec, “Inductive representation learning on large graphs,” in Advances in Neural Information Processing Systems, I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, Eds., vol. 30. Curran Associates, Inc., 2017. [Online]. Available: https://proceedings.neurips.cc/paper_files/paper/2017/file/ 5dd9db5e033da9c6fb5ba83c7a7ebea9-Paper.pdf [4] K. Xu, W. Hu, J. Leskovec, and S. Jegelka, “How powerful are graph neural networks?” 2019. [5] W. Hu, M. Fey, H. Ren, M. Nakata, Y. Dong, and J. Leskovec, “Ogb-lsc: A largescale challenge for machine learning on graphs,” arXiv preprint arXiv:2103.09430, 2021. [6] NVIDIA, https://docs.nvidia.com/doca/sdk/doca+dma/index.html, 2024. [7] M. Fey and J. E. Lenssen, “Fast graph representation learning with PyTorch Geometric,” in ICLR Workshop on Representation Learning on Graphs and Manifolds, 2019. [8] M. Wang, D. Zheng, Z. Ye, Q. Gan, M. Li, X. Song, J. Zhou, C. Ma, L. Yu, Y. Gai, T. Xiao, T. He, G. Karypis, J. Li, and Z. Zhang, “Deep graph library: A graphcentric, highly-performant package for graph neural networks,” 2020. [9] W.-L. Chiang, X. Liu, S. Si, Y. Li, S. Bengio, and C.-J. Hsieh, “Cluster-gcn: An efficient algorithm for training deep and large graph convolutional networks,” ser. KDD ’19. New York, NY, USA: Association for Computing Machinery, 2019, p. 257–266. [Online]. Available: https://doi.org/10.1145/3292500.3330925 [10] Y. Lee, J. Chung, and M. Rhu, “Smartsage: Training large-scale graph neural networks using in-storage processing architectures,” 2022. [11] W. Hu, M. Fey, M. Zitnik, Y. Dong, H. Ren, B. Liu, M. Catasta, and J. Leskovec, “Open graph benchmark: Datasets for machine learning on graphs,” arXiv preprint arXiv:2005.00687, 2020. [12] V. Vanhoucke, A. Senior, and M. Z. Mao, “Improving the speed of neural networks on cpus,” in Deep Learning and Unsupervised Feature Learning Workshop, NIPS 2011, 2011. [13] I. Hubara, M. Courbariaux, D. Soudry, R. El-Yaniv, and Y. Bengio, “Binarized neural networks,” in Advances in Neural Information Processing Systems, D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, Eds., vol. 29. Curran Associates, Inc., 2016. [Online]. Available: https://proceedings.neurips.cc/paper_ files/paper/2016/file/d8330f857a17c53d217014ee776bfd50-Paper.pdf

[14] S. Han, H. Mao, and W. J. Dally, “Deep compression: Compressing deep neural network with pruning, trained quantization and huffman coding,” in 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, Y. Bengio and Y. LeCun, Eds., 2016. [Online]. Available: http://arxiv.org/abs/1510.00149 [15] R. Krishnamoorthi, “Quantizing deep convolutional networks for efficient inference: A whitepaper,” CoRR, vol. abs/1806.08342, 2018. [Online]. Available: http://arxiv.org/abs/1806.08342 [16] S. Chen, D. Zheng, C. Ding, C. Huan, Y. Ji, and H. Liu, “Tango: re-thinking quantization for graph neural network training on gpus,” in Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, ser. SC ’23. New York, NY, USA: Association for Computing Machinery, 2023. [Online]. Available: https://doi.org/10.1145/3581784.3607037 [17] R. Novkin, F. Klemme, and H. Amrouch, “Approximation- and quantizationaware training for graph neural networks,” IEEE Transactions on Computers, vol. 73, no. 2, pp. 599–612, 2024. [18] S. Wang, B. Eravci, R. Guliyev, and H. Ferhatosmanoglu, “Low-bit quantization for deep graph neural networks with smoothness-aware message propagation,” in Proceedings of the 32nd ACM International Conference on Information and Knowledge Management, ser. CIKM ’23. New York, NY, USA: Association for Computing Machinery, 2023, p. 2626–2636. [Online]. Available: https: //doi.org/10.1145/3583780.3614955 [19] NVIDIA, “NVIDIA BLUEFIELD-2 DPU Data Center Infrastructure on a Chip,” https://resources.nvidia.com/en-us-accelerated-networking-resourcelibrary/bluefield-2-dpu-datasheet, NVIDIA, 2021. [20] ——, “Nvidia doca,” https://developer.nvidia.com/networking/doca, NVIDIA Developer, 2024, accessed: 2024-01-11. [21] J. Chen, J. Zhu, and L. Song, “Stochastic training of graph convolutional networks with variance reduction,” 2018. [Online]. Available: https://arxiv.org/ abs/1710.10568 [22] R. Ying, R. He, K. Chen, P. Eksombatchai, W. L. Hamilton, and J. Leskovec, “Graph convolutional neural networks for web-scale recommender systems,” in Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, ser. KDD ’18. New York, NY, USA: Association for Computing Machinery, 2018, p. 974–983. [Online]. Available: https://doi.org/10.1145/3219819.3219890 [23] J. Chen, T. Ma, and C. Xiao, “Fastgcn: Fast learning with graph convolutional networks via importance sampling,” CoRR, vol. abs/1801.10247, 2018. [Online]. Available: http://arxiv.org/abs/1801.10247 [24] W. Huang, T. Zhang, Y. Rong, and J. Huang, “Adaptive sampling towards fast graph representation learning,” in Advances in Neural Information Processing Systems, S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, Eds., vol. 31. Curran Associates, Inc., 2018. [Online]. Available: https://proceedings.neurips.cc/paper_files/paper/2018/file/ 01eee509ee2f68dc6014898c309e86bf-Paper.pdf [25] D. Zou, Z. Hu, Y. Wang, S. Jiang, Y. Sun, and Q. Gu, “Layer-dependent importance sampling for training deep and large graph convolutional networks,” in Advances in Neural Information Processing Systems, H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, Eds., vol. 32. Curran Associates, Inc., 2019. [Online]. Available: https://proceedings.neurips.cc/paper_ files/paper/2019/file/91ba4a4478a66bee9812b0804b6f9d1b-Paper.pdf [26] H. Zeng, H. Zhou, A. Srivastava, R. Kannan, and V. K. Prasanna, “Graphsaint: Graph sampling based inductive learning method,” ArXiv, vol. abs/1907.04931, 2019. [Online]. Available: https://api.semanticscholar.org/CorpusID:195886159 [27] J. Bai, Y. Ren, and J. Zhang, “Ripple walk training: A subgraph-based training framework for large and deep graph neural network,” in 2021 International Joint Conference on Neural Networks (IJCNN), 2021, pp. 1–8. [28] H. Tong, C. Faloutsos, and J.-y. Pan, “Fast random walk with restart and its applications,” in Sixth International Conference on Data Mining (ICDM’06), 2006, pp. 613–622. [29] F. Wu, A. Souza, T. Zhang, C. Fifty, T. Yu, and K. Weinberger, “Simplifying graph convolutional networks,” in Proceedings of the 36th International Conference on Machine Learning, ser. Proceedings of Machine Learning Research, K. Chaudhuri and R. Salakhutdinov, Eds., vol. 97. PMLR, 09–15 Jun 2019, pp. 6861–6871. [Online]. Available: https://proceedings.mlr.press/v97/wu19e.html [30] D. Korolija, D. Koutsoukos, K. Keeton, K. Taranov, D. Milojičić, and G. Alonso, “Farview: Disaggregated memory with operator off-loading for database engines,” 2021. [31] Z. Cai, X. Yan, Y. Wu, K. Ma, J. Cheng, and F. Yu, “DGCL: An efficient communication library for distributed GNN training,” in Proceedings of the Sixteenth European Conference on Computer Systems, ser. EuroSys ’21. New York, NY, USA: Association for Computing Machinery, 2021, pp. 130–144. [Online]. Available: https://doi.org/10.1145/3447786.3456233 [32] H. Mostafa, “Sequential aggregation and rematerialization: Distributed fullbatch training of graph neural networks on large graphs,” arXiv preprint arXiv:2111.06483, 2022. [Online]. Available: https://arxiv.org/abs/2111.06483 [33] Z. Zhou, C. Li, X. Wei, X. Wang, and G. Sun, “GNNear: Accelerating full-batch training of graph neural networks with near-memory processing,” arXiv preprint

LGNNIC: Acceleration of Large-Scale GNN Training using SmartNICs

arXiv:2111.00680, 2022. [Online]. Available: https://arxiv.org/abs/2111.00680 [34] B. Zhang, S. R. Kuppannagari, R. Kannan, and V. Prasanna, “Efficient neighborsampling-based gnn training on cpu-fpga heterogeneous platform,” in 2021 IEEE High Performance Extreme Computing Conference (HPEC), 2021, pp. 1–7. [35] Y. Gui, B. Wei, W. Yuan, and X. Jin, “Hardware acceleration of sampling algorithms in sample and aggregate graph neural networks,” arXiv preprint arXiv:2209.02916, 2022. [Online]. Available: https://arxiv.org/abs/2209.02916 [36] Y.-C. Lin, B. Zhang, and V. Prasanna, “HitGNN: High-throughput GNN training framework on cpu+multi-fpga heterogeneous platform,” arXiv preprint arXiv:2303.01568, 2023. [Online]. Available: https://arxiv.org/abs/2303.01568

A Appendix A.1 Neighbor Sampling With PyG In PyG’s mini-batch training with Neighbor Sampling, each epoch iterates over the NeighborLoader class (a subclass of PyTorch’s DataLoader), typically instantiated as train_loader, to generate sampled mini-batches. The tensors from these batches—including node indices, features, and labels—are transferred to the GPU, where the remaining training steps, such as the forward pass, loss computation, backpropagation, and parameter updates, are executed. Algorithm 3 PyG mini-batch training with Neighbor Sampling 1: Initialize: Dataset 2: Initialize: Train_Loader Configuration 3: Initialize: Model 4: Initialize: Optimizer 5: for epoch = 1 to Epoch Number do 6: 7: 8: 9: 10: 11:

A.2

for each mini-batch in Train_Loader do Transfer: Mini-Batch Data to GPU Perform: Forward Pass Through The Model Calculate: Loss Using Predicted and Actual Labels Calculate: Gradients via Backpropagation Update: Model Parameters Using Optimizer

Overall Training Speedups through Offloading Tensor Quantization

In Table 4, we compare FP32 and FP16 precision when quantizing the feature tensor to highlight the impact of quantization on total transfer time and, consequently, total training (execution) time. Here, in Table 5, we present the detailed results of the measured average total transaction time, average total execution time, and the resulting speedups that were derived from them. The measurements were conducted across all mechanisms (DOCA-DMA and Sockets), datasets, and sampling hyperparameters (with common hyperparameters highlighted in yellow). Notably, in settings where transaction time dominates overall training time—such as with the high-overhead Sockets mechanism— the execution speedup closely tracks the transfer speedup. This effect is particularly pronounced in the Reddit dataset, which has the highest edge-to-node ratio among those evaluated. The high edge-to-node ratio also explains why Reddit exhibits the highest total training time among the evaluated datasets with Sockets, while preprocessing time remains moderate. In contrast, when transaction time constitutes a smaller fraction of total training time—as is typically the case with the low-overhead DOCA-DMA mechanism, where preprocessing is also a significant bottleneck—the execution

and transfer speedups diverge. Since DOCA-DMA reduces communication overhead substantially, the impact of quantization is less pronounced compared to Sockets, where communication overhead remains the primary bottleneck.

Liad Gerstman, Aditya Dhakal, Dejan Milojicic, and Avi Mendelson Num Neighbors Sockets

DOCA -DMA

[ −1, 256] [25, 10] [10, 10] [ −1] [ −1, 256]

Tot. Trans. Time [s] (FP32 / FP16) 315.52 / 192.87 36.84 / 10.11 14.80 / 4.11 60.75 / 16.65 1.73 / 1.39

Speed -up 1.64 3.64 3.60 3.65 1.24

Tot. Exec. Time [s] (FP32 / FP16) 325.84 / 201.43 37.19 / 10.45 15.00 / 4.31 61.10 / 17.00 11.54 / 11.18

Speed -up 1.62 3.56 3.48 3.59 1.03

[25, 10]

0.34 / 0.25

1.36

0.66 / 0.55

1.20

[10, 10] [ −1]

0.22 / 0.17 0.44 / 0.33

1.24 1.37

0.40 / 0.35 0.76 / 0.61

1.15 1.24

(a) Reddit Num Neighbors Sockets

DOCA -DMA

[ −1, −1] [15, 10, 5] [25, 10] [ −1] [ −1, −1]

Tot. Trans. Time [s] (FP32 / FP16) 305.76 / 135.99 15.65 / 4.61 3.62 / 1.00 2.59 / 0.57 1.59 / 1.28

Speed -up 2.25 3.39 3.63 4.53 1.24

Tot. Exec. Time [s] (FP32 / FP16) 318.23 / 147.03 17.22 / 6.07 4.24 / 1.56 2.97 / 0.90 13.81 / 13.60

Speed -up 2.16 2.84 2.71 3.31 1.02

[15, 10, 5]

0.25 / 0.17

1.41

1.87 / 1.79

1.04

[25, 10] [ −1]

0.10 / 0.09 0.09 / 0.04

1.18 2.02

0.75 / 0.73 0.48 / 0.43

1.03 1.11

(b) OGBN-Products Num Neighbors Sockets

DOCA -DMA

[ −1, −1] [25, 10] [10, 10] [ −1] [ −1, −1]

Tot. Trans. Time [s] (FP32 / FP16) 87.02 / 28.75 1.65 / 0.76 1.18 / 0.60 0.18 / 0.15 0.91 / 0.69

Speed -up 3.03 2.16 1.96 1.20 1.33

Tot. Exec. Time [s] (FP32 / FP16) 176.76 / 59.96 3.62 / 1.80 2.62 / 1.43 0.44 / 0.36 4.09 / 3.62

Speed -up 2.95 2.01 1.83 1.20 1.13

[25, 10]

0.12 / 0.06

1.93

0.55 / 0.44

1.25

[10, 10] [ −1]

0.10 / 0.05 0.02 / 0.017

1.96 1.12

0.47 / 0.36 0.10 / 0.10

1.31 1.05

(c) OGBN-MAG

Table 5: Comparison of FP32 and FP16 precision: average total transaction time, average total execution time, and corresponding speedups across all mechanisms, datasets, and sampling hyperparameters.

Related documents

• LLN: Learnable Lens Networks for Parameter-Efficient Long-Horizon Dynamical Prediction#1128762
• Very Deep Convolutional Networks for Large-Scale Image Recognition#6258
• Context-Dependent Pre-Trained Deep Neural Networks for Large-Vocabulary Speech Recognition#303010
• CosmoFlow: Using Deep Learning to Learn the Universe at Scale#620154
• A survey of the recent architectures of deep convolutional neural networks#19729