RecShard: statistical feature-based memory optimization for industry-scale neural recommendation

Abstract

We propose RecShard, a fine-grained embedding table (EMB) partitioning and placement technique for deep learning recommendation models (DLRMs). RecShard is designed based on two key observations. First, not all EMBs are equal, nor all rows within an EMB are equal in terms of access patterns. EMBs exhibit distinct memory characteristics, providing performance optimization opportunities for intelligent EMB partitioning and placement across a tiered memory hierarchy. Second, in modern DLRMs, EMBs function as hash tables. As a result, EMBs display interesting phenomena, such as the birthday paradox, leaving EMBs severely under-utilized. Rec-Shard determines an optimal EMB sharding strategy for a set of EMBs based on training data distributions and model characteristics, along with the bandwidth characteristics of the underlying tiered memory hierarchy. In doing so, RecShard achieves over 6 times higher EMB training throughput on average for capacity constrained DLRMs. The throughput increase comes from improved EMB load balance by over 12 times and from the reduced access to the slower memory by over 87 times.

AutoShard: Automated Embedding Table Sharding for Recommender Systems

“…Baselines. We compare DreamShard against human expert strategies from previous work [27,8,28], including size-based, dim-based, lookup-based, size-lookup-based greedy balancing strategies. We also include an RNN-based RL algorithm [13], which uses RNN architecture to map operators to devices.…”

“…Starting from the table with the highest cost, greedy algorithm will assign tables one-by-one to the device with the lowest sum of the costs so far, so that each device will have roughly an equal sum of the costs in the end. We consider the heuristics with the following cost functions, which have been proven to show strong performance in prior work [26]: the size (the product of dimension and hash size) of the table (size-greedy), the dimension of the table (dim-greedy), the product of the dimension and mean pooling factor of the table (lookup-greedy). We further include a random sharding baseline (rand).…”

“…Our work is orthogonal to these methods since AutoShard can be applied to compressed tables as well. A related work explored using the tiered memory hierarchy to store the embedding tables [26]. They exploited the unequal access patterns of embedding tables to improve the efficiency by placing hot rows in the GPU memory.…”

“…They exploited the unequal access patterns of embedding tables to improve the efficiency by placing hot rows in the GPU memory. Our efforts complement [26] by providing an end-to-end learning-based framework for cost approximation and partitioning optimization.…”

Understanding data storage and ingestion for large-scale deep recommendation model training

“…Deep-RecSys [35] and RecSSD [79] optimized inference requests across CPUs and GPUs, and SSDs, respectively. Acun et al [18] characterized DLRM architectures on GPU trainers, Sethi et al [70] presented an optimized embedding sharding strategy for DLRM training, Maeng et al [51] explored checkpointing trainer state, and AIBox [86] optimized training using hierarchical memory for parameters. However, these prior works do not discuss the DSI pipeline, a critical part of ML training.…”