# AutoShard: Automated Embedding Table Sharding for Recommender Systems

arXiv: https://arxiv.org/abs/2208.06399 · HTML: https://mlsanity.com/paper/2208.06399/

## Key facts

- Cited by 17 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 43 references that are in the map.
- First posted on arXiv 12 Aug 2022.

## Builds on (most cited first, 30 of 34)

- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Neural Architecture Search with Reinforcement Learning](https://mlsanity.com/paper/1611.01578/) (2016)
- [Neural Collaborative Filtering](https://mlsanity.com/paper/1708.05031/) (2017)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018)
- [Wide & Deep Learning for Recommender Systems](https://mlsanity.com/paper/1606.07792/) (2016)
- [AutoAugment: Learning Augmentation Policies from Data](https://mlsanity.com/paper/1805.09501/) (2018)
- [Deep Interest Evolution Network for Click-Through Rate Prediction](https://mlsanity.com/paper/1809.03672/) (2018)
- [Deep Learning Recommendation Model for Personalization and Recommendation Systems](https://mlsanity.com/paper/1906.00091/) (2019)
- [AutoInt: Automatic Feature Interaction Learning via Self-Attentive Neural Networks](https://mlsanity.com/paper/1810.11921/) (2018)
- [Random Search and Reproducibility for Neural Architecture Search](https://mlsanity.com/paper/1902.07638/) (2019)
- [Deep Learning based Recommender System: A Survey and New Perspectives](https://mlsanity.com/paper/1707.07435/) (2017)
- [The Architectural Implications of Facebook's DNN-based Personalized Recommendation](https://mlsanity.com/paper/1906.03109/) (2019)
- [Exploratory Combinatorial Optimization with Reinforcement Learning](https://mlsanity.com/paper/1909.04063/) (2019)
- [Compositional Embeddings Using Complementary Partitions for Memory-Efficient Recommendation Systems](https://mlsanity.com/paper/1909.02107/) (2019)
- [DouZero: Mastering DouDizhu with Self-Play Deep Reinforcement Learning](https://mlsanity.com/paper/2106.06135/) (2021)
- [Distributed Hierarchical GPU Parameter Server for Massive Scale Deep Learning Ads Systems](https://mlsanity.com/paper/2003.05622/) (2020)
- [Experience Replay Optimization](https://mlsanity.com/paper/1906.08387/) (2019)
- [AutoEmb: Automated Embedding Dimensionality Search in Streaming Recommendations](https://mlsanity.com/paper/2002.11252/) (2020)
- [Understanding Training Efficiency of Deep Learning Recommendation Models at Scale](https://mlsanity.com/paper/2011.05497/) (2020)
- [Learnable Embedding Sizes for Recommender Systems](https://mlsanity.com/paper/2101.07577/) (2021)
- [TorchBeast: A PyTorch Platform for Distributed RL](https://mlsanity.com/paper/1910.03552/) (2019)
- [TODS: An Automated Time Series Outlier Detection System](https://mlsanity.com/paper/2009.09822/) (2020)
- [FBGEMM: Enabling High-Performance Low-Precision Deep Learning Inference](https://mlsanity.com/paper/2101.05615/) (2021)
- [Neural Input Search for Large Scale Recommendation Models](https://mlsanity.com/paper/1907.04471/) (2019)
- [Deep Learning Training in Facebook Data Centers: Design of Scale-up and Scale-out Systems](https://mlsanity.com/paper/2003.09518/) (2020)
- [Learning to Embed Categorical Features without Embedding Tables for Recommendation](https://mlsanity.com/paper/2010.10784/) (2020)
- [Meta-AAD: Active Anomaly Detection with Deep Reinforcement Learning](https://mlsanity.com/paper/2009.07415/) (2020)
- [Model Size Reduction Using Frequency Based Double Hashing for Recommender Systems](https://mlsanity.com/paper/2007.14523/) (2020)
- [Security and Privacy on Blockchain](https://mlsanity.com/paper/1903.07602/) (2019)
- [Rank the Episodes: A Simple Approach for Exploration in Procedurally-Generated Environments](https://mlsanity.com/paper/2101.08152/) (2021)

## Built on by (most cited first, 17 of 17)

- [Data-centric Artificial Intelligence: A Survey](https://mlsanity.com/paper/2303.10158/) (2023)
- [SurCo: Learning Linear Surrogates For Combinatorial Nonlinear Optimization Problems](https://mlsanity.com/paper/2210.12547/) (2022)
- [Scaling User Modeling: Large-scale Online User Representations for Ads Personalization in Meta](https://mlsanity.com/paper/2311.09544/) (2023)
- [CoRTX: Contrastive Framework for Real-time Explanation](https://mlsanity.com/paper/2303.02794/) (2023)
- [DreamShard: Generalizable Embedding Table Placement for Recommender Systems](https://mlsanity.com/paper/2210.02023/) (2022)
- [Chakra: Advancing Performance Benchmarking and Co-design using Standardized Execution Traces](https://mlsanity.com/paper/2305.14516/) (2023)
- [Towards Automated Imbalanced Learning with Deep Hierarchical Reinforcement Learning](https://mlsanity.com/paper/2208.12433/) (2022)
- [MP-Rec: Hardware-Software Co-Design to Enable Multi-Path Recommendation](https://mlsanity.com/paper/2302.10872/) (2023)
- [Pre-train and Search: Efficient Embedding Table Sharding with Pre-trained Neural Cost Models](https://mlsanity.com/paper/2305.01868/) (2023)
- [FlexShard: Flexible Sharding for Industry-Scale Sequence Recommendation Models](https://mlsanity.com/paper/2301.02959/) (2023)
- [Towards Personalized Preprocessing Pipeline Search](https://mlsanity.com/paper/2302.14329/) (2023)
- [Towards Universal Performance Modeling for Machine Learning Training on Multi-GPU Platforms](https://mlsanity.com/paper/2404.12674/) (2024)
- [Disaggregating Embedding Recommendation Systems with FlexEMR](https://mlsanity.com/paper/2410.12794/) (2024)
- [SCRec: A Scalable Computational Storage System with Statistical Sharding and Tensor-train Decomposition for Recommendation Models](https://mlsanity.com/paper/2504.00520/) (2025)
- [Deep Recommender Models Inference: Automatic Asymmetric Data Flow Optimization](https://mlsanity.com/paper/2507.01676/) (2025)
- [NestPipe: Large-Scale Recommendation Training on 1,500+ Accelerators via Nested Pipelining](https://mlsanity.com/paper/2604.06956/) (2026)
- [When KV Meets Embeddings: Dynamic GPU Memory Allocation for Accelerating Generative Recommender Serving](https://mlsanity.com/paper/2605.04450/) (2026)

## Abstract (arXiv)

Embedding learning is an important technique in deep recommendation models to map categorical features to dense vectors. However, the embedding tables often demand an extremely large number of parameters, which become the storage and efficiency bottlenecks. Distributed training solutions have been adopted to partition the embedding tables into multiple devices. However, the embedding tables can easily lead to imbalances if not carefully partitioned. This is a significant design challenge of distributed systems named embedding table sharding, i.e., how we should partition the embedding tables to balance the costs across devices, which is a non-trivial task because 1) it is hard to efficiently and precisely measure the cost, and 2) the partition problem is known to be NP-hard. In this work, we introduce our novel practice in Meta, namely AutoShard, which uses a neural cost model to directly predict the multi-table costs and leverages deep reinforcement learning to solve the partition problem. Experimental results on an open-sourced large-scale synthetic dataset and Meta's production dataset demonstrate the superiority of AutoShard over the heuristics. Moreover, the learned policy of AutoShard can transfer to sharding tasks with various numbers of tables and different ratios of the unseen tables without any fine-tuning. Furthermore, AutoShard can efficiently shard hundreds of tables in seconds. The effectiveness, transferability, and efficiency of AutoShard make it desirable for production use. Our algorithms have been deployed in Meta production environment. A prototype is available at https://github.com/daochenzha/autoshard