Single-machine graph neural network optimization method and system based on historical embedding cache

Abstract

A single-machine graph neural network optimization method and system based on historical embedding caching is proposed. This method involves pre-sampling and analyzing the training dataset to collect sampled data. After calculating the size of the historical embedding cache for each level using heuristic methods, the cache is set according to the calculated size, and pre-sampling is performed, i.e., iterative sampling based on subgraphs. During each sampling process, cache replacement is simulated, subgraph pruning is performed, and node access sequences and subgraph topological connections are obtained. After repeating sampling several times to obtain the corresponding number of training subgraphs, several steps of optimal feature cache replacement are calculated and the optimal feature cache replacement information is saved. When loading the training subgraphs for actual training, historical embedding is pushed / pulled in conjunction with the sampled subgraphs during each training round, and the underlying feature cache is replaced in conjunction with the optimal feature cache replacement information, thus achieving single-machine graph neural network optimization.

Description

Technical Field

[0001] This invention relates to a technology in the field of neural networks, specifically a single-machine graph neural network optimization method and system based on historical embedding caching. Background Technology

[0002] Existing neural network technologies require feature caching in memory when training large-scale graph data on a single machine to alleviate the high latency of accessing external memory and optimize training speed. While existing heuristic methods combined with optimal cache replacement algorithms can improve feature cache hit rate to some extent, they still suffer from low cache hit rate, long data loading time, and low computation-to-memory ratio due to the random access characteristics of graph data and the additional randomness introduced during graph neural network training to satisfy the assumption of independent and identically distributed training samples. Summary of the Invention

[0003] This invention addresses the problem that existing static graph-based techniques are not applicable to training scenarios with dynamic sampling and suffer from low training efficiency in the original training scenario. It proposes a single-machine graph neural network optimization method and system based on historical embedding cache. By setting up a historical embedding cache, historical embeddings in the graph neural network training process are reused, which effectively reduces the access to the underlying original feature data during the training process and improves training efficiency without affecting accuracy.

[0004] This invention is achieved through the following technical solution:

[0005] This invention relates to a single-machine graph neural network optimization method based on historical embedding caching. The method involves pre-sampling and analyzing the training dataset to collect sampled data, calculating the size of the historical embedding cache for each level using heuristic methods, setting the cache according to the calculated size, and performing pre-sampling, i.e., iteratively sampling based on subgraphs. During each sampling process, cache replacement is simulated, subgraph pruning is performed, and node access sequences and subgraph topological connections are obtained. After repeating sampling several times to obtain the corresponding number of training subgraphs, several steps of optimal feature cache replacement are calculated and the optimal feature cache replacement information is saved. When loading the training subgraphs for actual training, historical embedding is pushed / pulled in conjunction with the sampled subgraphs during each training round, and the underlying feature cache is replaced using the optimal feature cache replacement information, thus achieving single-machine graph neural network optimization.

[0006] The aforementioned standalone machine refers to a single server architecture consisting of a CPU, GPU, and memory, and an SSD external storage for storing the original graph data. During graph neural network training, data interaction occurs through memory and SSD external storage.

[0007] The aforementioned pre-sampling analysis refers to: performing a sampling process based on the training nodes of the current batch, collecting the sampled subgraph metadata, but without performing subsequent data loading and actual training processes.

[0008] The sampled subgraph metadata includes the number of nodes at the i-th level in the subgraph, NumNode. i .

[0009] The aforementioned pre-sampling process simulates cache replacement through a multi-layer cache interaction replacement algorithm and obtains node access sequences and pruned subgraphs through subgraph pruning. At the same time, it calculates the optimal replacement information for feature caches based on the node access sequences.

[0010] The subgraph pruning process involves sampling and subgraph pruning within the cache directory obtained from the presampled simulation cache replacement.

[0011] The aforementioned multi-level cache interaction and replacement algorithm refers to: information transfer between caches through a sampling subgraph, ensuring that neighbors of higher-level cache nodes do not appear in lower-level caches. Cache replacement within each level is performed using a FIFO (First-In, First-Out) approach; when a higher-level cache hits, it is removed from the sampling subgraph.

[0012] The cache replacement is further combined with the expiration threshold refresh, specifically: when caching, the batch timestamp is recorded; before each round of cache update begins, it is first judged that if the difference between the timestamp of the head element of the cache queue and the current timestamp exceeds the expiration threshold, the expired cache is evicted first, and then the cache is updated; otherwise, the cache is updated directly.

[0013] The simulated cache replacement refers to performing a simulated cache replacement through the cache directory during the sampling process. The number of nodes retained in each level of the cache directory is consistent with the historical embedded cache. The cache directory does not retain actual data, but only the access node number. During the sampling process, the cache directory is dynamically updated according to the multi-level cache interaction replacement algorithm.

[0014] The loading of the training subgraph for actual training refers to: for a batch of training nodes, sampling recursively extracting the k-hop neighbors of the training nodes from the adjacency matrix of the graph; after sampling to obtain the node visit sequence, data preparation is performed according to the visit sequence (loading the original features of the nodes from the SSD external storage) and the prepared data is passed to the GPU, and the GPU performs the subsequent graph neural network calculation process.

[0015] The heuristic method mentioned refers to combining the number of nodes at each level, NumNode, obtained from pre-sampling analysis. i (i∈{1,2,…,k}), combined with the sampling fanout parameter FanOut i(i∈{1,2,…,k}), the historical embedding dimension HiddenFeatureDim, and the GPU memory capacity limit B are used to calculate the size of the historical embedding cache at each level, specifically including:

[0016] i) First, configure the number of top-level historical embedded cache nodes (i=1): CacheNode i =NumNode i *CacheBatch, sets the number of nodes cached in the underlying historical embedded cache (CacheNode). i =CacheNode i-1 *FanOut i *α(i∈{2,3,…,k}), finally calculate the expected memory size to be used by the cache. Then compare it with the memory capacity limit B. If MemExpt≤B, the calculation process ends and the historical embedding cache is configured according to the calculation result; otherwise, proceed to step ii), where: α is the sampling neighborhood overlap degree;

[0017] ii) Calculate sequentially And find the first one that meets the requirements MemExpt k′ For allocation levels ≤ B, k′ is assigned. For levels less than or equal to k′, historical embeddings are configured according to the calculation results. The historical embedding cache size for levels greater than k′ is set to 0. The configuration result is returned.

[0018] The subgraph pruning mentioned above refers to the following: during the sampling process, if the current sampling node is hit in the corresponding cache directory, the hit status of the current sampling node is marked, and the current sampling node is removed from the subsequent recursive sampling; the above process is repeated for each level of sampling, and finally a pruned subgraph is obtained.

[0019] The aforementioned historical embedding push / pull refers to the following: When performing the forward propagation process of the neural network, if the embedding of a subgraph node is not calculated in real time from the original features (corresponding to a presampling processing cache hit), the historical embedding of the corresponding node needs to be pulled from the historical embedding cache and concatenated into the corresponding tensor for forward propagation in subsequent layers; for the missing embeddings calculated in real time from the original features in the current round, the calculated embedding value needs to be pushed to the corresponding historical embedding cache to update the historical embedding cache.

[0020] Technical effect

[0021] This invention collects and analyzes sampled data based on pre-sampling analysis, and calculates the size of a multi-level historical embedding cache according to the sampled data and training parameters. Through multi-level cache interaction and replacement, cache hit status is obtained in real time during pre-sampling processing, and the status of hit nodes is marked as hit in a timely manner to avoid the influence of neighboring nodes of hit nodes on the next level of cache. In pre-sampling processing, cache replacement is simulated by maintaining a cache directory, and subgraph pruning is performed in conjunction with a multi-level cache interaction and replacement algorithm, ultimately reducing the number of bottom-level nodes in the subgraph and the number of accesses to external storage. Compared with existing technologies, this invention can effectively reduce the size of the sampled subgraph, thereby reducing memory access during training and significantly reducing the computational and memory access volume during training, thus improving training efficiency. Attached Figure Description

[0022] Figure 1 This is a flowchart of the present invention;

[0023] Figure 2 This is a schematic diagram of the pre-sampling analysis module;

[0024] Figure 3 This is a schematic diagram of the multi-layer cache interaction and replacement process;

[0025] Figure 4 The graph shows the result of the reduced training time for this system. Detailed Implementation

[0026] This embodiment relates to a graph neural network optimization system based on a single-machine off-core architecture, including: a pre-sampling analysis module, a pre-sampling simulation cache replacement module, and a trainer history embedding update module. The pre-sampling analysis module performs heuristic calculations based on the metadata information of the sampled subgraph obtained from the sampling process to obtain the size of the history embedding cache at each level. The pre-sampling simulation cache replacement module performs a multi-layer cache interaction replacement algorithm based on the pre-sampling subgraph and cache directory status information to obtain the pruned subgraph result to be trained. The trainer history embedding update module performs historical embedding push / pull processing based on the pruned subgraph information to obtain the final training result.

[0027] like Figure 1 As shown, the graph neural network optimization method based on the above system is as follows: After calculating the configuration of the historical embedding cache size for each level through the pre-sampling analysis module, pre-sampling is performed. That is, while performing subgraph-based sampling in a loop, cache replacement is simulated through pre-sampling processing, and pre-sampling processing subgraph pruning is performed in combination with cache information to obtain pruned training subgraphs. After repeating this process several times to obtain the corresponding number of training subgraphs, the optimal feature cache replacement is calculated for several steps, and cache replacement information is calculated. During the training phase, the training subgraphs are loaded for actual training. During each round of training, the historical embedding update module pushes / retrieves historical embeddings and updates the feature cache.

[0028] The aforementioned presampling specifically includes:

[0029] Step 1: Maintain a cache directory for each sampled layer, initially initializing each cache directory to empty. In each training round, based on the training nodes of the current round and the cache directory status of the current layer, perform subgraph-based sampling of the training nodes in conjunction with pre-sampling subgraph pruning.

[0030] The sampling algorithm described herein may employ, but is not limited to, a water storage tank sampling algorithm.

[0031] Step 2: After sampling is completed, for nodes that are not found in the corresponding level of the cache directory, the cache directory is updated by simulating cache replacement through pre-sampling processing. Then, the missing nodes are updated to the corresponding cache directory through a multi-level cache interaction replacement algorithm to simulate the cache replacement situation in the actual training process, so that the newly generated node embeddings can be used in subsequent rounds.

[0032] Step 3: Repeat steps 1 and 2 a total of k times. Finally, sample k node sequences node_idx and k node adjacency matrices adjs. Based on the access information of node_idx in k rounds, run the optimal cache replacement algorithm to sample the future access data obtained in advance, obtain the feature cache replacement information of each round, and save it to SSD for use in the training phase.

[0033] The optimal cache replacement algorithm refers to prioritizing the replacement of data blocks with the furthest future access time during actual execution. Since subsequent access data is obtained in advance through sampling, it can pass the optimal cache.

[0034] The feature cache replacement information includes: the newly replaced feature node number in_id, the replaced feature node number out_id, and the location where the replaced feature is stored in_position.

[0035] The training phase specifically includes:

[0036] Step a: Load the node sequence node_idx and node adjacency matrix adjs obtained during the pre-training process. Prepare training data according to the current feature cache. Specifically, for node n∈node_idx, if the features of n are in the current feature cache, load them directly from the cache; otherwise, load the features of n from SSD. Finally, obtain the original feature tensor batch_input. Input the original feature tensor batch_input and the node adjacency matrix adjs into the neural network to be optimized for forward propagation-based training.

[0037] Step b: Push / fetch historical embeddings during the training process. If the corresponding embedding value is not found in the historical embedding cache, the newly calculated embedding needs to be updated to the historical embedding cache.

[0038] The update involves removing the oldest embedding from the cache when the historical embedding cache is full, according to the FIFO principle, to reduce the lag of the historical embeddings used, and to ensure training convergence and training accuracy.

[0039] The term "historical embedding cache miss" means that the corresponding position of the historical embedding of a node is empty, indicating that the corresponding node has not been accessed or has been evicted and replaced in previous training rounds.

[0040] Step c: After forward propagation, load the feature cache replacement information calculated in the presampling phase. Update the feature cache based on the feature cache replacement information.

[0041] Step d: Repeat steps a, b, and c a total of k times to complete one super-batch training.

[0042] Step e: Delete the node sequence node_idx, node adjacency matrix adjs, and feature cache replacement information stored on the SSD.

[0043] Through specific experiments, a Linux server equipped with a 40GB A100 GPU, where the A100 contains 108 streaming multiprocessors (SMs), each with 64 FP32 cores, was tested. The PCIe type was PCI-E 4.0×16, with a maximum bandwidth of 31.5GB / s. The server was equipped with an Intel(R) Xeon(R) Silver 4310 CPU and 128GB of host memory, along with a 3.84TB Samsung Pm9a3 NVMe SSD. The Pm9a3 achieved an input / output speed of 200K per second (IOPS) and a sequential read speed of 6,900MB / s. Multiple real-world graph training datasets (with node labels, used for performance testing and accuracy measurement) were used as workloads, including OGBN-Papers100M, OGBN-Products, and IGB-Medium, as well as multiple real-world graph datasets (without node labels, used for performance validation), including Twitter and Friendster.

[0044] like Figure 4The figure shows a comparison of the system training performance in this embodiment. This embodiment and a baseline system without historical embedding caching were simultaneously trained on physical machines for different graph datasets, and the training time was measured and compared. Calculations show that all three datasets can achieve significant training acceleration on this system. Compared to the baseline system without historical embedding caching, the ogbn-papers100M dataset in this embodiment achieves an 18.3% reduction in training time while maintaining the same training accuracy as the baseline; the ogbn-products dataset achieves a 23.4% reduction in training time; and the igb-medium dataset achieves a 20% reduction in training time.

[0045] Compared to existing technologies, this method's presampling simulation cache replacement module updates the cache directory by simulating real cache replacement. Combined with presampling subgraph pruning, it obtains pruned subgraphs, reducing the number of original features accessed during training of the pruned subgraphs. This achieves higher training efficiency and shorter end-to-end training latency without compromising training accuracy. Compared to related single-machine core-outer graph neural network training systems, it reduces training time by 15% and end-to-end training latency by 8%.

[0046] The above-described specific implementations can be partially adjusted by those skilled in the art in different ways without departing from the principles and purpose of the present invention. The scope of protection of the present invention is defined by the claims and is not limited to the above-described specific implementations. All implementation schemes within the scope of the claims are bound by the present invention.

Claims

1. A graph neural network optimization method based on a single-machine off-core architecture, characterized in that, By pre-sampling and analyzing the training dataset to collect sampled data, and then calculating the size of the historical embedding cache for each level using heuristic methods, the cache is set according to the calculated size of the historical embedding cache for each level, and pre-sampling is performed, i.e., sampling based on subgraphs is performed cyclically: in each sampling process, cache replacement is simulated, subgraph pruning is performed, and node access sequences and subgraph topological connections are obtained. After repeating the sampling several times to obtain the corresponding number of training subgraphs, the optimal feature cache replacement is calculated for several steps and the optimal feature cache replacement information is saved; when loading the training subgraphs for actual training, historical embedding is pushed / pulled in combination with the sampled subgraphs in each round of training, and the bottom-level feature cache is replaced in combination with the optimal feature cache replacement information, thus realizing the optimization of single-machine graph neural network; The aforementioned pre-sampling analysis refers to: performing a sampling process based on the training nodes of the current batch, collecting the sampled subgraph metadata, but without performing subsequent data loading and the actual training process; The aforementioned pre-sampling process simulates cache replacement through a multi-layer cache interaction replacement algorithm and obtains node access sequences and pruned subgraphs through subgraph pruning, while simultaneously calculating the optimal replacement information for feature caches based on the node access sequences. The subgraph pruning mentioned above involves sampling and subgraph pruning in the cache directory obtained by pre-sampling simulation cache replacement. The aforementioned multi-level cache interaction and replacement algorithm refers to: passing information between caches through a sampling subgraph; avoiding the occurrence of neighbors of higher-level cache nodes in lower-level caches; using a FIFO approach for cache replacement in each level; and removing a cache from the sampling subgraph when a higher-level cache is hit. The simulated cache replacement refers to: during the sampling process, simulated cache replacement is performed through the cache directory. The number of nodes retained in each level of the cache directory is consistent with the historical embedded cache. The cache directory does not retain actual data, but only the access node number. During the sampling process, the cache directory is dynamically updated according to the multi-level cache interaction replacement algorithm. The heuristic method mentioned refers to combining the number of nodes at each level obtained from pre-sampling analysis. ( ), combined with sampling fan-out parameters Historical Embedded Dimension and GPU memory capacity limitations The size of the historical embedding cache at each level is calculated, specifically including: i) First, configure the top-level history embedding cache, i.e. Number of cache nodes * Set the number of nodes cached in the lower-level historical embedded cache. ,Right now Finally, calculate the estimated memory size to be used by the cache. And compare it with the memory capacity limit B, when The calculation process ends, and the historical embedding cache is configured according to the calculation result; otherwise, step ii) is executed, where: This represents the overlap of the sampling neighborhood. ii) Calculate sequentially , And find the first one that meets the requirements. Allocation hierarchy For levels less than or equal to Configure historical embeddings according to the calculation results; the hierarchy is greater than Set the historical embedding cache size to 0; return the configuration result.

2. The graph neural network optimization method based on a single-machine off-core architecture according to claim 1, characterized in that, The cache replacement is further combined with the expiration threshold refresh, specifically: when caching, the batch timestamp is recorded; before each round of cache update begins, it is first judged that if the difference between the timestamp of the head element of the cache queue and the current timestamp exceeds the expiration threshold, the expired cache is evicted first, and then the cache is updated; otherwise, the cache is updated directly.

3. The graph neural network optimization method based on a single-machine off-core architecture according to claim 1, characterized in that, The loading of the training subgraph for actual training refers to: for a batch of training nodes, sampling recursively extracting the k-hop neighbors of the training nodes from the adjacency matrix of the graph; after sampling to obtain the node visit sequence, data preparation is performed according to the visit sequence, that is, loading the original features of the nodes from the SSD external storage and passing the prepared data to the GPU, and the GPU performs the subsequent graph neural network calculation process.

4. The graph neural network optimization method based on a single-machine off-core architecture according to claim 1, characterized in that, The subgraph pruning mentioned above refers to the following: during the sampling process, if the current sampling node is hit in the corresponding cache directory, the hit status of the current sampling node is marked, and the current sampling node is removed from the subsequent recursive sampling; the above process is repeated for each level of sampling, and finally a pruned subgraph is obtained.

5. The graph neural network optimization method based on a single-machine off-core architecture according to claim 1, characterized in that, The aforementioned historical embedding push / pull refers to the following: during the forward propagation process of the neural network, if the embedding of a subgraph node is not calculated in real time from the original features, i.e., the corresponding presampling processing cache is hit, the historical embedding of the corresponding node needs to be pulled from the historical embedding cache and concatenated into the corresponding tensor for forward propagation in subsequent layers; for the missing embeddings calculated in real time from the original features in the current round, the calculated embedding value needs to be pushed to the corresponding historical embedding cache to update the historical embedding cache.

6. The graph neural network optimization method based on a single-machine off-core architecture according to claim 1, characterized in that, The loading of the training subgraph for actual training specifically includes: Step a: Load the node sequence obtained during the pre-training process. and node adjacency matrix The training data is prepared based on the current feature cache, specifically for node n. If the features of n are in the current feature cache, they are loaded directly from the cache; otherwise, the features of n are loaded from the SSD, ultimately yielding the original feature tensor. , the original feature tensor and node adjacency matrix The neural network to be optimized is then trained based on forward propagation. Step b: Push / fetch historical embeddings during the training process. If the corresponding embedding value is not found in the historical embedding cache, the newly calculated embedding needs to be updated to the historical embedding cache. The update involves removing the oldest embedding from the cache when the historical embedding cache is full, according to the FIFO principle, to reduce the lag of the historical embeddings used, and to ensure training convergence and training accuracy. The historical embedding cache miss means that the corresponding position of the node's historical embedding is empty, indicating that the corresponding node has not been accessed in previous training rounds or has been evicted and replaced. Step c: After the forward propagation is completed, load the feature cache replacement information calculated in the presampling stage, and update the feature cache according to the feature cache replacement information; Step d: Repeat steps a, b, and c a total of k times to complete one superbatch training. Step e: Delete the node sequence stored on the SSD Node adjacency matrix Replace information with feature cache.

7. A graph neural network optimization system based on a single-machine off-core architecture for implementing the method of any one of claims 1-6, characterized in that, include: The system comprises a pre-sampling analysis module, a pre-sampling simulation cache replacement module, and a trainer history embedding update module. The pre-sampling analysis module uses heuristic methods to calculate the size of the history embedding cache at each level based on the metadata information of the sampled subgraph obtained from the sampling process. The presampled simulation cache replacement module performs multi-layer cache interaction replacement algorithm processing based on the presampled subgraph and cache directory status information to obtain the pruned subgraph result to be trained; the trainer historical embedding update module performs historical embedding push / pull processing based on the pruned subgraph information to obtain the final training result.

Images