Unified and Near-optimal Multi-GPU Cache for Embedding-based Deep Learning

This paper presents UGache , a unified multi-GPU cache system designed for embedding-based deep learning (EmbDL). UGache is primarily motivated by the unique characteristics of EmbDL applications, namely read-only and skewed embedding accesses with affinity and predictability. UGache introduces a novel factored extraction mechanism that avoids bandwidth congestion to fully exploit high-speed cross-GPU interconnects (e.g., NVLink and NVSwitch). Based on a hotness metric, UGache also provides a near-optimal cache policy that balances local and remote access to minimize the extraction time for diverse GPU interconnect topologies. We have implemented UGache and integrated it into two representative frameworks, TensorFlow and PyTorch. Evaluation using two typical types of EmbDL applications, namely graph neural network (GNN) training and deep learning recommendation (DLR) inference, shows that UGache outperforms state-of-the-art replication and partition designs by an average of 1.93 × and 1.63 × (up to 5.25 × and 3.45 ×), respectively. Furthermore, we demonstrate the applicability of UGache ’s principle beyond embedding-based deep learning, with an example of text-to-image generation on an inference cluster.