# DreamShard: Generalizable Embedding Table Placement for Recommender Systems

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

## Key facts

- Cited by 14 papers in the map (data through 14 Oct 2026).
- Builds on 52 of its 77 references that are in the map.
- First posted on arXiv 5 Oct 2022.

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

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Continuous control with deep reinforcement learning](https://mlsanity.com/paper/1509.02971/) (2015)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Neural Architecture Search with Reinforcement Learning](https://mlsanity.com/paper/1611.01578/) (2016)
- [Hindsight Experience Replay](https://mlsanity.com/paper/1707.01495/) (2017)
- [Neural Collaborative Filtering](https://mlsanity.com/paper/1708.05031/) (2017)
- [Efficient Neural Architecture Search via Parameter Sharing](https://mlsanity.com/paper/1802.03268/) (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)
- [Neural Combinatorial Optimization with Reinforcement Learning](https://mlsanity.com/paper/1611.09940/) (2016)
- [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)
- [Beyond Data and Model Parallelism for Deep Neural Networks](https://mlsanity.com/paper/1807.05358/) (2018)
- [Billion-scale Commodity Embedding for E-commerce Recommendation in Alibaba](https://mlsanity.com/paper/1803.02349/) (2018)
- [Device Placement Optimization with Reinforcement Learning](https://mlsanity.com/paper/1706.04972/) (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)
- [DNNFusion: Accelerating Deep Neural Networks Execution with Advanced Operator Fusion](https://mlsanity.com/paper/2108.13342/) (2021)
- [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)
- [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)

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

- [Data-centric Artificial Intelligence: A Survey](https://mlsanity.com/paper/2303.10158/) (2023)
- [Bring Your Own View: Graph Neural Networks for Link Prediction with Personalized Subgraph Selection](https://mlsanity.com/paper/2212.12488/) (2022)
- [SurCo: Learning Linear Surrogates For Combinatorial Nonlinear Optimization Problems](https://mlsanity.com/paper/2210.12547/) (2022)
- [Chakra: Advancing Performance Benchmarking and Co-design using Standardized Execution Traces](https://mlsanity.com/paper/2305.14516/) (2023)
- [RSC: Accelerating Graph Neural Networks Training via Randomized Sparse Computations](https://mlsanity.com/paper/2210.10737/) (2022)
- [MP-Rec: Hardware-Software Co-Design to Enable Multi-Path Recommendation](https://mlsanity.com/paper/2302.10872/) (2023)
- [Collaborative Graph Neural Networks for Attributed Network Embedding](https://mlsanity.com/paper/2307.11981/) (2023)
- [Pre-train and Search: Efficient Embedding Table Sharding with Pre-trained Neural Cost Models](https://mlsanity.com/paper/2305.01868/) (2023)
- [Robust and Minimally Invasive Watermarking for EaaS](https://mlsanity.com/paper/2410.17552/) (2024)
- [Two-dimensional Sparse Parallelism for Large Scale Deep Learning Recommendation Model Training](https://mlsanity.com/paper/2508.03854/) (2025)
- [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)
- [Pushing the Performance Envelope of DNN-based Recommendation Systems Inference on GPUs](https://mlsanity.com/paper/2410.22249/) (2024)

## Abstract (arXiv)

We study embedding table placement for distributed recommender systems, which aims to partition and place the tables on multiple hardware devices (e.g., GPUs) to balance the computation and communication costs. Although prior work has explored learning-based approaches for the device placement of computational graphs, embedding table placement remains to be a challenging problem because of 1) the operation fusion of embedding tables, and 2) the generalizability requirement on unseen placement tasks with different numbers of tables and/or devices. To this end, we present DreamShard, a reinforcement learning (RL) approach for embedding table placement. DreamShard achieves the reasoning of operation fusion and generalizability with 1) a cost network to directly predict the costs of the fused operation, and 2) a policy network that is efficiently trained on an estimated Markov decision process (MDP) without real GPU execution, where the states and the rewards are estimated with the cost network. Equipped with sum and max representation reductions, the two networks can directly generalize to any unseen tasks with different numbers of tables and/or devices without fine-tuning. Extensive experiments show that DreamShard substantially outperforms the existing human expert and RNN-based strategies with up to 19% speedup over the strongest baseline on large-scale synthetic tables and our production tables. The code is available at https://github.com/daochenzha/dreamshard