RecShard: Statistical Feature-Based Memory Optimization for Industry-Scale Neural Recommendation

TL;DR

RecShard introduces a memory optimization technique for recommendation models that intelligently partitions and places embedding tables based on access patterns, significantly improving training throughput and memory utilization.

Contribution

It presents a novel embedding partitioning and placement method tailored for tiered memory hierarchies in deep learning recommendation models, leveraging data distribution and model insights.

Findings

Over 6x increase in training throughput on average.

More than 12x better load balancing of embeddings.

87x reduction in slow memory accesses.

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. RecShard 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.