Graph Neural Network Training With Memory Expansion Batching
Find Innovative SolutionsGenerate Solutions
Solution Overview
Problem
The training of graph neural networks on large-scale graphs is hindered by GPU memory limitations, leading to slow training speeds due to the GPU waiting for data from disk-based systems, which consume significant CPU time in reading and processing.
Innovation Solution
A method involving a GPU that determines batch information and transmits it to memory expansion devices for feature data acquisition, allowing parallel processing and reducing data movement, using devices like CXL Memory Modules with DRAM and NAND flash for efficient data handling.
Engineering Contradictions & Design Principles
Engineering Contradiction Analysis
1Quantity of substance
If disk-based systems are used for training large-scale graph neural networks, then GPU memory limitations are addressed, but training speed deteriorates due to GPU waiting time
Solution Approach 1:
The patent segments the training data into multiple batches and processes them in parallel across multiple GPUs. Each GPU receives a portion of the batch information and processes corresponding feature data independently, eliminating the sequential processing bottleneck and reducing overall training time while maintaining the ability to handle large-scale graphs
Solution Approach 2:
The patent implements preliminary action by pre-fetching and caching feature data in high-speed memory before the GPU needs it. The system prepares batch information and corresponding feature data in advance, so when the GPU is ready to process, the data is already available, eliminating waiting time and improving training speed
2Reliability
If CPU reads and processes training data from disk, then data is made available for training, but time consumption increases significantly
Solution Approach 1:
The patent introduces an intermediary component (such as a memory management unit or cache system) that sits between the disk storage and the GPU. This intermediary pre-loads and caches feature data in high-speed memory, acting as a buffer that eliminates the slow disk I/O operations from the critical training path while ensuring data availability for the GPU
Solution Approach 2:
The system performs preliminary data preparation by pre-fetching feature data from disk storage into high-speed cache memory before the GPU processing stage. This advance preparation eliminates the time-consuming disk read operations during actual training, maintaining data availability while dramatically reducing processing time
Data Source
AI summary
A method and a device for training a graph neural network are provided. The method may be performed by a graphics processing unit (GPU), and may include determining at least one batch of training data; transmitting batch information corresponding to the determined at least one batch to at least one memory expansion device, so that the at least one memory expansion device acquires feature data for one or more data blocks of the at least one batch based on the batch information, receiving the feature data from the at least one memory expansion device; and training the graph neural network based on the feature data.