Machine Learning-Guided Memory Optimization for DLRM Inference on Tiered Memory

2025-03-01

International Symposium on High-Performance Computer Architecture

AI Summary

Deep learning recommendation models (DLRMs) require terabyte-scale embedding memory, and while hierarchical memory offers cost efficiency, its irregular access patterns severely degrade embedding placement and cache efficiency. To address this, we propose RecMG—a novel system that decouples cache admission decisions from prefetching prediction into two independently trainable models for the first time. We introduce a differentiable loss function explicitly modeling long reuse distances and infrequent embedding accesses, significantly improving prefetching accuracy. RecMG integrates vectorized access pattern learning, hierarchical-memory-aware scheduling, and dynamic prefetching policies. Experimental results show that RecMG reduces on-demand embedding loads by 1.5–2.8× over state-of-the-art baselines and achieves up to a 43% reduction in end-to-end inference latency in industrial-scale deployments.

Technology Category

• Data Mining & Knowledge Management: Recommender Systems
• Machine Learning: Hardware-aware ML
• Search and Optimization: Learning to Search

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs

Abstract

Deep learning recommendation models (DLRMs) are widely used in industry, and their memory capacity requirements reach the terabyte scale. Tiered memory architectures provide a cost-effective solution but introduce challenges in embedding-vector placement due to complex embedding-access patterns. We propose RecMG, a machine learning (ML)-guided system for vector caching and prefetching on tiered memory. RecMG accurately predicts accesses to embedding vectors with long reuse distances or few reuses. The design of RecMG focuses on making ML feasible in the context of DLRM inference by addressing unique challenges in data labeling and navigating the search space for embedding-vector placement. By employing separate ML models for caching and prefetching, plus a novel differentiable loss function, RecMG narrows the prefetching search space and minimizes on-demand fetches. Compared to state-of-the-art temporal, spatial, and ML-based prefetchers, RecMG reduces on-demand fetches by $2.2 imes, 2.8 imes$, and $1.5 imes$, respectively. In industrial-scale DLRM inference scenarios, RecMG effectively reduces end-to-end DLRM inference time by up to 43%.

Problem

• Optimizing memory placement for terabyte-scale recommendation models
• Addressing embedding-access challenges in tiered memory architectures
• Reducing inference latency through ML-guided caching and prefetching

Innovation

• ML-guided system for tiered memory vector caching
• Separate ML models optimize caching and prefetching strategies
• Novel differentiable loss minimizes on-demand memory fetches