OMeGa: Boosting Large-scale Graph Embeddings with Heterogeneous Memory Processing
Peng Fang, Siqiang Luo, Fang Wang, Bolong Zheng, Hong Jiang, Dan Feng, Hechang Pan, Xingyu Wan · 2025
Graph embedding, which maps graph nodes to lowdimensional vectors, is a widely used technique for graph representation learning. However, most existing graph embedding models suffer from high memory consumption, limiting their scalability to large graphs. Heterogeneous memory systems that combine DRAM and Persistent Memory (PM) offer new opportunities for scaling up memory capacity. Despite this advantage, the performance gap (on the order of 5x) between DRAM and PM is magnified (by 3.3-4.2x) under non-uniform memory access (NUMA) architecture. Additionally, the inherent sparsity of graphs induces numerous random accesses in the fundamental Sparse Matrix and Dense Matrix Multiplication (SpMM) operations of graph embedding, hindering high-performance heterogeneous memory processing. To address these challenges, this paper presents OMeGa that focuses on Optimizing heterogeneous Memory processing for large-scale Graph embedding. OMeGa leverages an entropy-aware thread allocation, simultaneously achieving workload balancing and tail latency reduction across threads. It also incorporates a workload feature-aware prefetcher to alleviate random accesses during streaming heterogeneous processing. In addition, OMeGa devises a NUMA-aware data placement, aiming to minimize the adverse impact of NUMA on heterogeneous memory. The experiments conducted on billion-scale graphs demonstrate that OMeGa exhibits an average acceleration of 32.03x with strong scalability. This pioneering capability enables the efficient generation of large-scale graph embeddings, free from the memory size constraints and performance disparities typically encountered in heterogeneous memory systems.