Paper deep dive
LGNNIC: Acceleration of Large-Scale GNN Training using SmartNICs
Liad Gerstman, Aditya Dhakal, Dejan Milojicic, Avi Mendelson
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 88%
Last extracted: 8/11/2026, 3:54:58 AM
Summary
The paper introduces LGNNIC, a system architecture that accelerates large-scale Graph Neural Network (GNN) training by offloading preprocessing tasks, specifically Neighbor Sampling and Tensor Quantization, to SmartNICs co-located with remote memory nodes. This approach reduces inter-node network congestion and data transfer volumes, achieving significant training speedups compared to traditional CPU-based preprocessing and socket-based communication.
Entities (12)
Relation Signals (9)
LGNNIC â uses â SmartNIC
confidence 95% ¡ LGNNIC ... leverages SmartNICs co-located with remote memory nodes
LGNNIC â offloads â Neighbor Sampling
confidence 92% ¡ LGNNIC offloads key preprocessing tasks to SmartNICs ... Neighbor Sampling
NVIDIA BlueField-2 â implements â LGNNIC
confidence 90% ¡ comprising one remote-memory node with an NVIDIA BlueField-2 SmartNIC
LGNNIC â offloads â Tensor Quantization
confidence 90% ¡ LGNNIC offloads key preprocessing tasks ... Quantization of the sampled batches
LGNNIC â reduces â network congestion
confidence 90% ¡ reduce communication overhead in distributed GNN training ... alleviating network congestion
Neighbor Sampling â achievesspeedup â LGNNIC
confidence 85% ¡ Neighbor Sampling achieved up to 62.4x and 17.5x speedups
Tensor Quantization â achievesspeedup â LGNNIC
confidence 85% ¡ Quantization provided additional speedups of up to 3.6x
DOCA-DMA â isusedby â LGNNIC
confidence 85% ¡ optimized low-overhead DMA-based synchronization mechanism ... DOCA-DMA
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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 inter-node 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 DMA-based synchronization mechanism and a high-overhead socket-based 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 62.4x and 17.5x speedups with Sockets and DOCA-DMA, respectively, primarily due to reduced data transaction time. Quantization provided additional speedups of up to 3.6x and 1.3x, respectively, by reducing data transfer.
Tags
Links
- Source: https://arxiv.org/abs/2608.07733v1
- Canonical: https://arxiv.org/abs/2608.07733v1
Trouble viewing inline? Open PDF directly â
Full Text
75,147 characters extracted from source content.
Expand or collapse full text
LGNNIC: Acceleration of Large-Scale GNN Training using SmartNICs Liad Gerstman Technion Israel Institute of TechnologyHaifa, Israel liadgerstman@campus.technion.ac.il , Aditya Dhakal Hewlett Packard Enterprise LabsFort Collins, CO, USA aditya.dhakal@hpe.com , Dejan Milojicic Hewlett Packard Enterprise LabsMilpitas, CA, USA dejan.milojicic@hpe.com and Avi Mendelson Technion Israel Institute of TechnologyHaifa, Israel mendlson@technion.ac.il 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 inter-node 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 DMA-based synchronization mechanism and a high-overhead socket-based 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 Networks (GCN) (Kipf and Welling, 2017), Graph Attention Networks (GAT) (VeliÄkoviÄ et al., 2018), GraphSAGE (Hamilton et al., 2017) and Graph Isomorphism Network (GIN) (Xu et al., 2019). 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) (Hu et al., 2021) and the encyclopedic knowledge graph WikiKG90Mv2 (89 GB) (Hu et al., 2021), 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 (Hamilton et al., 2017). 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. Table 1. Attributes of the examined PyG datasets Dataset Nodes Edges Features Classes Reddit 232,965 114,615,892 602 41 OGBN- Products 2,449,029 61,859,140 100 47 OGBN- MAG 1,939,743 21,111,007 128 349 To address the challenges of GNN learning on distributed graphs, we present LGNNIC (Large-Scale GNN Training using SmartNICs). 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 (Hamilton et al., 2017), 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 (NVIDIA, 2024a) 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 latency-optimized 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 (Fey and Lenssen, 2019)) and DGL (Deep Graph Library) (Wang et al., 2020). 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 (Hamilton et al., 2017; Hu et al., 2021; Chiang et al., 2019; Lee et al., 2022) 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 (Hamilton et al., 2017). The second dataset, OGBN-Products, involves predicting the category of an Amazon product based on products purchased together (Hu et al., 2020). 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 (Hu et al., 2020). 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 real-world graphs that could fit within the memory constraints of these devices, enabling us to effectively demonstrate the benefits of the system. 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][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][-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 (Vanhoucke et al., 2011) (Hubara et al., 2016) (Han et al., 2016) (Krishnamoorthi, 2018). Quantization techniques have also been successfully applied to GNNs (Chen et al., 2023)(Novkin et al., 2024)(Wang et al., 2023). In addition to its computational benefits, quantization can alleviate network bottlenecks in distributed systems by reducing the amount of data transmitted across the network. (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. 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 mini-batch sampling (Hamilton et al., 2017) 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. 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 high-bandwidth 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 mini-batch 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 communication-reduction mechanisms, as detailed in Section 5. 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. 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 (NVIDIA, 2021), 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 (NVIDIA, 2024b), 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. 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. 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 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, one-directional, C-based DOCA-DMA API (NVIDIA, 2024a), which enables DMA-based 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. Algorithm 1 BF-2 Side (Mem. 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=1epoch=1 to Epoch Number do 5: while batch=next(train_loader) do 6: for Each Batch Element Tensor do âł i.e., Packet 7: Quantize: Tensor âł Optional 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 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: Model 4:Initialize: Optimizer 5:for epoch=1epoch=1 to Epoch Number do 6: while Epoch_Finished_Message is False do 7: for Each Batch Element Tensor do âł i.e., Packet 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 âł Optional 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 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 coordinated with its predecessor N-1, ensuring data from mini-batch 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 mini-batches 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 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. (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. The mini-batch preprocessing phase (top-left subfigures in Figures 4(a), 4(b), 4(c)) 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 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 4(a): [-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 4(b): [-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 4(c), 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 num\_workers>0) with unconstrained two-hop Neighbor Sampling (num_neighbors=[â1,â1] 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 4(a), 4(b), and 4(c)). Notably, increasing the number of unconstrained neighborhood-expansion layers exhausted the available SmartNIC memory even with num_workers=0 num\_workers=0. The corresponding runs did not complete successfully and were therefore excluded from the performance figures. (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. 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 num\_workers=0, BlueField-2 Neighbor Sampling was approximately 2â10Ă2--10Ă slower than sampling on the EPYC 7513. Nevertheless, as evidenced by Figures 4(a), 4(b), and 4(c), 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 4(a), 4(b), and 4(c), since preprocessing on the SmartNIC without CPU Affinity provided superior performance, we disabled this feature (num_workers=0 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 num\_workers=0). For example, with Reddit (Figure 4(a)), BlueFieldâs preprocessing time dropped from âź17 17 seconds to âź0.3 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-DMA-based synchronization mechanism to transfer the preprocessed 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 5(a), 5(b), and 5(c)) and Sockets (upper subfigure in Figures 5(a), 5(b), and 5(c))â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. Sockets DOCA-DMA Num Neigh. Trans. Speedup Total Speedup Trans. Speedup Total Speedup Reddit Products MAG Reddit Products MAG Reddit Products MAG Reddit Products MAG 8.56 84.49 52.64 8.76 75.12 45.51 5.10 15.36 7.44 17.46 18.43 7.45 21.32 â 73.57 21.73 â 62.39 7.99 â 8.73 29.03 â 8.61 â 19.53 â â 18.48 â â 6.45 â â 7.40 â 5.19 118.12 472.08 5.33 107.17 353.57 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. Figure 5 shows that DOCA-DMA consistently outperforms the socket-based mechanism across all datasets, primarily due to its use 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 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. Num Neighbors Reddit OGBN- Products OGBN- MAG FP32 FP16 FP32 FP16 FP32 FP16 [-1, -1] â â 0.777 0.774 0.445 0.444 [-1, 256] 0.951 0.951 â â â â [15, 10, 5] â â 0.789 0.790 â â [25, 10] 0.952 0.952 0.771 0.776 0.444 0.447 [10, 10] 0.952 0.952 â â 0.453 0.451 [-1] 0.943 0.938 0.685 0.673 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. 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 (DOCA-DMA 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 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). REDDIT Num Neighbors Trans. Speedup Exec. Speedup [â1,256][-1,256] 1.64 1.62 Sockets [25,10][25,10] 3.64 3.56 [10,10][10,10] 3.60 3.48 [â1][-1] 3.65 3.59 [â1,256][-1,256] 1.24 1.03 DOCA -DMA [25,10][25,10] 1.36 1.20 [10,10][10,10] 1.24 1.15 [â1][-1] 1.37 1.24 OGBN- PRODUCTS Num Neighbors Trans. Speedup Exec. Speedup [â1,â1][-1,-1] 2.25 2.16 Sockets [15,10,5][15,10,5] 3.39 2.84 [25,10][25,10] 3.63 2.71 [â1][-1] 4.53 3.31 [â1,â1][-1,-1] 1.24 1.02 DOCA -DMA [15,10,5][15,10,5] 1.41 1.04 [25,10][25,10] 1.18 1.03 [â1][-1] 2.02 1.11 OGBN- MAG Num Neighbors Trans. Speedup Exec. Speedup [â1,â1][-1,-1] 3.03 2.95 Sockets [25,10][25,10] 2.16 2.01 [10,10][10,10] 1.96 1.83 [â1][-1] 1.20 1.20 [â1,â1][-1,-1] 1.33 1.13 DOCA -DMA [25,10][25,10] 1.93 1.25 [10,10][10,10] 1.96 1.31 [â1][-1] 1.12 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. 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 software-based 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. Architectural Improvements Working with BlueField-2 has provided valuable insights into potential architectural improvements. This section presents a selection of the most important ones. 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-to-node 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 DOCA-DMA 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. 7. Related Work Unlike prior work, which generally limits SmartNIC offloading to lightweight or stateless operations, LGNNIC offloads mini-batch 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 mini-batchesâ computational graphs. These sampled mini-batches are then fed into the GPU for training. Various sampling categories exist, including node-wise (e.g., GraphSAGE (Hamilton et al., 2017), VR-GC (Chen et al., 2018a), PinSAGE Sampler (Ying et al., 2018)), layer-wise (e.g., FastGCN (Chen et al., 2018b), AS-GCN (Huang et al., 2018), LADIES (Zou et al., 2019)), and subgraph-wise (e.g., GraphSAINT (Zeng et al., 2019), RWT (Bai et al., 2021)). Among node-wise methods, the most widely used approach is Neighbor Sampling, introduced by Hamilton et al. (Hamilton et al., 2017). 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) (Tong et al., 2006), 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 (Chiang et al., 2019) 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 (Wu et al., 2019) 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 (Korolija et al., 2021), 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 (Cai et al., 2021) also addresses communication overheads by distributing GNN training across multiple GPUs rather than across remote memory nodes. The Sequential Aggregation and Rematerialization paper (Mostafa, 2022) 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 (Lee et al., 2022), e.g., deals with accelerating the same GraphSAGE (Hamilton et al., 2017) sampling algorithm implemented based on the same library (PyG (Fey and Lenssen, 2019)) we use, with an in-storage architecture exploiting a smartSSD composed of NVMe SSD and an FPGA. The GNNear paper (Zhou et al., 2022) 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 (Zhang et al., 2021; Gui et al., 2022) and on multi-FPGA platforms (Lin et al., 2023). While prior work mostly focuses on quantizing activations and weights (Vanhoucke et al., 2011; Hubara et al., 2016; Han et al., 2016; Krishnamoorthi, 2018; Chen et al., 2023; Novkin et al., 2024; Wang et al., 2023), 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 large-scale GNN training using SmartNICs. LGNNIC offloads Neighbor Sampling and Quantization to the SmartNIC, reducing data transfer 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 peer-to-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 J. Bai, Y. Ren, and J. Zhang (2021) Ripple walk training: a subgraph-based training framework for large and deep graph neural network. In 2021 International Joint Conference on Neural Networks (IJCNN), Vol. , p. 1â8. External Links: Document Cited by: §7.1. Z. Cai, X. Yan, Y. Wu, K. Ma, J. Cheng, and F. Yu (2021) DGCL: an efficient communication library for distributed GNN training. In Proceedings of the Sixteenth European Conference on Computer Systems, EuroSys â21, New York, NY, USA, p. 130â144. External Links: ISBN 9781450383349, Link, Document Cited by: §7.2. J. Chen, J. Zhu, and L. Song (2018a) Stochastic training of graph convolutional networks with variance reduction. External Links: 1710.10568, Link Cited by: §7.1. J. Chen, T. Ma, and C. Xiao (2018b) FastGCN: fast learning with graph convolutional networks via importance sampling. CoRR abs/1801.10247. External Links: Link, 1801.10247 Cited by: §7.1. S. Chen, D. Zheng, C. Ding, C. Huan, Y. Ji, and H. Liu (2023) 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, SC â23, New York, NY, USA. External Links: ISBN 9798400701092, Link, Document Cited by: §2.3, §7.2. W. Chiang, X. Liu, S. Si, Y. Li, S. Bengio, and C. Hsieh (2019) Cluster-gcn: an efficient algorithm for training deep and large graph convolutional networks. KDD â19, New York, NY, USA, p. 257â266. External Links: ISBN 9781450362016, Link, Document Cited by: §2.1, §7.1. M. Fey and J. E. Lenssen (2019) Fast graph representation learning with PyTorch Geometric. In ICLR Workshop on Representation Learning on Graphs and Manifolds, Cited by: §2.1, §7.2. Y. Gui, B. Wei, W. Yuan, and X. Jin (2022) Hardware acceleration of sampling algorithms in sample and aggregate graph neural networks. arXiv preprint arXiv:2209.02916. External Links: Link Cited by: §7.2. W. Hamilton, Z. Ying, and J. Leskovec (2017) 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, p. . External Links: Link Cited by: §1, §1, §1, §2.1, §2.1, §3, §7.1, §7.1, §7.2. S. Han, H. Mao, and W. J. Dally (2016) 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.), External Links: Link Cited by: §2.3, §7.2. W. Hu, M. Fey, H. Ren, M. Nakata, Y. Dong, and J. Leskovec (2021) OGB-lsc: a large-scale challenge for machine learning on graphs. arXiv preprint arXiv:2103.09430. Cited by: §1, §2.1. W. Hu, M. Fey, M. Zitnik, Y. Dong, H. Ren, B. Liu, M. Catasta, and J. Leskovec (2020) Open graph benchmark: datasets for machine learning on graphs. arXiv preprint arXiv:2005.00687. Cited by: §2.1. W. Huang, T. Zhang, Y. Rong, and J. Huang (2018) 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, p. . External Links: Link Cited by: §7.1. I. Hubara, M. Courbariaux, D. Soudry, R. El-Yaniv, and Y. Bengio (2016) Binarized neural networks. In Advances in Neural Information Processing Systems, D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett (Eds.), Vol. 29, p. . External Links: Link Cited by: §2.3, §7.2. T. N. Kipf and M. Welling (2017) Semi-supervised classification with graph convolutional networks. External Links: 1609.02907 Cited by: §1. D. Korolija, D. Koutsoukos, K. Keeton, K. Taranov, D. MilojiÄiÄ, and G. Alonso (2021) Farview: disaggregated memory with operator off-loading for database engines. External Links: 2106.07102 Cited by: §7.2. R. Krishnamoorthi (2018) Quantizing deep convolutional networks for efficient inference: A whitepaper. CoRR abs/1806.08342. External Links: Link, 1806.08342 Cited by: §2.3, §7.2. Y. Lee, J. Chung, and M. Rhu (2022) SmartSAGE: training large-scale graph neural networks using in-storage processing architectures. External Links: 2205.04711 Cited by: §2.1, §7.2. Y. Lin, B. Zhang, and V. Prasanna (2023) HitGNN: high-throughput GNN training framework on cpu+multi-fpga heterogeneous platform. arXiv preprint arXiv:2303.01568. External Links: Link Cited by: §7.2. H. Mostafa (2022) Sequential aggregation and rematerialization: distributed full-batch training of graph neural networks on large graphs. arXiv preprint arXiv:2111.06483. External Links: Link Cited by: §7.2. R. Novkin, F. Klemme, and H. Amrouch (2024) Approximation- and quantization-aware training for graph neural networks. IEEE Transactions on Computers 73 (2), p. 599â612. External Links: Document Cited by: §2.3, §7.2. NVIDIA (2021) NVIDIA BLUEFIELD-2 DPU Data Center Infrastructure on a Chip. NVIDIA. Note: https://resources.nvidia.com/en-us-accelerated-networking-resource-library/bluefield-2-dpu-datasheet Cited by: §5.1. NVIDIA (2024a) NVIDIA Docs. Note: https://docs.nvidia.com/doca/sdk/doca+dma/index.html Cited by: §1, §5.2. NVIDIA (2024b) NVIDIA doca. NVIDIA Developer. Note: https://developer.nvidia.com/networking/docaAccessed: 2024-01-11 Cited by: §5.1. H. Tong, C. Faloutsos, and J. Pan (2006) Fast random walk with restart and its applications. In Sixth International Conference on Data Mining (ICDMâ06), Vol. , p. 613â622. External Links: Document Cited by: §7.1. V. Vanhoucke, A. Senior, and M. Z. Mao (2011) Improving the speed of neural networks on cpus. In Deep Learning and Unsupervised Feature Learning Workshop, NIPS 2011, Cited by: §2.3, §7.2. P. VeliÄkoviÄ, G. Cucurull, A. Casanova, A. Romero, P. Liò, and Y. Bengio (2018) Graph attention networks. External Links: 1710.10903 Cited by: §1. 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 (2020) Deep graph library: a graph-centric, highly-performant package for graph neural networks. External Links: 1909.01315 Cited by: §2.1. S. Wang, B. Eravci, R. Guliyev, and H. Ferhatosmanoglu (2023) 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, CIKM â23, New York, NY, USA, p. 2626â2636. External Links: ISBN 9798400701245, Link, Document Cited by: §2.3, §7.2. F. Wu, A. Souza, T. Zhang, C. Fifty, T. Yu, and K. Weinberger (2019) Simplifying graph convolutional networks. In Proceedings of the 36th International Conference on Machine Learning, K. Chaudhuri and R. Salakhutdinov (Eds.), Proceedings of Machine Learning Research, Vol. 97, p. 6861â6871. External Links: Link Cited by: §7.1. K. Xu, W. Hu, J. Leskovec, and S. Jegelka (2019) How powerful are graph neural networks?. External Links: 1810.00826 Cited by: §1. R. Ying, R. He, K. Chen, P. Eksombatchai, W. L. Hamilton, and J. Leskovec (2018) Graph convolutional neural networks for web-scale recommender systems. In Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, KDD â18, New York, NY, USA, p. 974â983. External Links: ISBN 9781450355520, Link, Document Cited by: §7.1. H. Zeng, H. Zhou, A. Srivastava, R. Kannan, and V. K. Prasanna (2019) GraphSAINT: graph sampling based inductive learning method. ArXiv abs/1907.04931. External Links: Link Cited by: §7.1. B. Zhang, S. R. Kuppannagari, R. Kannan, and V. Prasanna (2021) Efficient neighbor-sampling-based gnn training on cpu-fpga heterogeneous platform. In 2021 IEEE High Performance Extreme Computing Conference (HPEC), Vol. , p. 1â7. External Links: Document Cited by: §7.2. Z. Zhou, C. Li, X. Wei, X. Wang, and G. Sun (2022) GNNear: accelerating full-batch training of graph neural networks with near-memory processing. arXiv preprint arXiv:2111.00680. External Links: Link Cited by: §7.2. D. Zou, Z. Hu, Y. Wang, S. Jiang, Y. Sun, and Q. Gu (2019) 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, p. . External Links: Link Cited by: §7.1. Appendix 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=1epoch=1 to Epoch Number do 6: for each mini-batch in Train_Loader do 7: Transfer: Mini-Batch Data to GPU 8: Perform: Forward Pass Through The Model 9: Calculate: Loss Using Predicted and Actual Labels 10: Calculate: Gradients via Backpropagation 11: Update: Model Parameters Using Optimizer A.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. 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. Num Neighbors Tot. Trans. Time [s] (FP32 / FP16) Speed -up Tot. Exec. Time [s] (FP32 / FP16) Speed -up [â1,256][-1,256] 315.52 / 192.87 1.64 325.84 / 201.43 1.62 Sockets [25,10][25,10] 36.84 / 10.11 3.64 37.19 / 10.45 3.56 [10,10][10,10] 14.80 / 4.11 3.60 15.00 / 4.31 3.48 [â1][-1] 60.75 / 16.65 3.65 61.10 / 17.00 3.59 [â1,256][-1,256] 1.73 / 1.39 1.24 11.54 / 11.18 1.03 DOCA -DMA [25,10][25,10] 0.34 / 0.25 1.36 0.66 / 0.55 1.20 [10,10][10,10] 0.22 / 0.17 1.24 0.40 / 0.35 1.15 [â1][-1] 0.44 / 0.33 1.37 0.76 / 0.61 1.24 (a) Reddit Num Neighbors Tot. Trans. Time [s] (FP32 / FP16) Speed -up Tot. Exec. Time [s] (FP32 / FP16) Speed -up [â1,â1][-1,-1] 305.76 / 135.99 2.25 318.23 / 147.03 2.16 Sockets [15,10,5][15,10,5] 15.65 / 4.61 3.39 17.22 / 6.07 2.84 [25,10][25,10] 3.62 / 1.00 3.63 4.24 / 1.56 2.71 [â1][-1] 2.59 / 0.57 4.53 2.97 / 0.90 3.31 [â1,â1][-1,-1] 1.59 / 1.28 1.24 13.81 / 13.60 1.02 DOCA -DMA [15,10,5][15,10,5] 0.25 / 0.17 1.41 1.87 / 1.79 1.04 [25,10][25,10] 0.10 / 0.09 1.18 0.75 / 0.73 1.03 [â1][-1] 0.09 / 0.04 2.02 0.48 / 0.43 1.11 (b) OGBN-Products Num Neighbors Tot. Trans. Time [s] (FP32 / FP16) Speed -up Tot. Exec. Time [s] (FP32 / FP16) Speed -up [â1,â1][-1,-1] 87.02 / 28.75 3.03 176.76 / 59.96 2.95 Sockets [25,10][25,10] 1.65 / 0.76 2.16 3.62 / 1.80 2.01 [10,10][10,10] 1.18 / 0.60 1.96 2.62 / 1.43 1.83 [â1][-1] 0.18 / 0.15 1.20 0.44 / 0.36 1.20 [â1,â1][-1,-1] 0.91 / 0.69 1.33 4.09 / 3.62 1.13 DOCA -DMA [25,10][25,10] 0.12 / 0.06 1.93 0.55 / 0.44 1.25 [10,10][10,10] 0.10 / 0.05 1.96 0.47 / 0.36 1.31 [â1][-1] 0.02 / 0.017 1.12 0.10 / 0.10 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.