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

Geet Sethi, Bilge Acun, Niket Agarwal, Christos Kozyrakis, Caroline Trippel, Carole Jean Wu

Research output: Chapter in Book/Report/Conference proceedingConference contribution

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.

Original languageEnglish (US)
Title of host publicationASPLOS 2022 - Proceedings of the 27th ACM International Conference on Architectural Support for Programming Languages and Operating Systems
EditorsBabak Falsafi, Michael Ferdman, Shan Lu, Thomas F. Wenisch
PublisherAssociation for Computing Machinery
Pages344-358
Number of pages15
ISBN (Electronic)9781450392051
DOIs
StatePublished - Feb 28 2022
Externally publishedYes
Event27th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, ASPLOS 2022 - Virtual, Online, Switzerland
Duration: Feb 28 2022Mar 4 2022

Publication series

NameInternational Conference on Architectural Support for Programming Languages and Operating Systems - ASPLOS

Conference

Conference27th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, ASPLOS 2022
Country/TerritorySwitzerland
CityVirtual, Online
Period2/28/223/4/22

Keywords

  • AI training systems
  • Deep learning recommendation models
  • Memory optimization
  • Neural networks

ASJC Scopus subject areas

  • Software
  • Information Systems
  • Hardware and Architecture

Fingerprint

Dive into the research topics of 'RecShard: Statistical Feature-Based Memory Optimization for Industry-Scale Neural Recommendation'. Together they form a unique fingerprint.

Cite this