# Pre-train and Search: Efficient Embedding Table Sharding with Pre-trained Neural Cost Models

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 58 of its 88 references that are in the map.
- First posted on arXiv 3 May 2023.

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

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [ALBERT: A Lite BERT for Self-supervised Learning of Language Representations](https://mlsanity.com/paper/1909.11942/) (2019)
- [Pre-train, Prompt, and Predict: A Systematic Survey of Prompting Methods in Natural Language Processing](https://mlsanity.com/paper/2107.13586/) (2021)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [Neural Collaborative Filtering](https://mlsanity.com/paper/1708.05031/) (2017)
- [Deep Reinforcement Learning that Matters](https://mlsanity.com/paper/1709.06560/) (2017)
- [Wide & Deep Learning for Recommender Systems](https://mlsanity.com/paper/1606.07792/) (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)
- [Hierarchical Deep Reinforcement Learning: Integrating Temporal Abstraction and Intrinsic Motivation](https://mlsanity.com/paper/1604.06057/) (2016)
- [Data-centric Artificial Intelligence: A Survey](https://mlsanity.com/paper/2303.10158/) (2023)
- [Beyond Data and Model Parallelism for Deep Neural Networks](https://mlsanity.com/paper/1807.05358/) (2018)
- [Compressing Neural Networks with the Hashing Trick](https://mlsanity.com/paper/1504.04788/) (2015)
- [Self-Imitation Learning](https://mlsanity.com/paper/1806.05635/) (2018)
- [The Science of Detecting LLM-Generated Texts](https://mlsanity.com/paper/2303.07205/) (2023)
- [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)
- [Software-Hardware Co-design for Fast and Scalable Training of Deep Learning Recommendation Models](https://mlsanity.com/paper/2104.05158/) (2021)
- [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)
- [Experience Replay Optimization](https://mlsanity.com/paper/1906.08387/) (2019)
- [Policy-GNN: Aggregation Optimization for Graph Neural Networks](https://mlsanity.com/paper/2006.15097/) (2020)
- [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)
- [Data-centric AI: Perspectives and Challenges](https://mlsanity.com/paper/2301.04819/) (2023)
- [Learnable Embedding Sizes for Recommender Systems](https://mlsanity.com/paper/2101.07577/) (2021)
- [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)

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

- [Wukong: Towards a Scaling Law for Large-Scale Recommendation](https://mlsanity.com/paper/2403.02545/) (2024)
- [Disaggregated Multi-Tower: Topology-aware Modeling Technique for Efficient Large-Scale Recommendation](https://mlsanity.com/paper/2403.00877/) (2024)
- [Two-dimensional Sparse Parallelism for Large Scale Deep Learning Recommendation Model Training](https://mlsanity.com/paper/2508.03854/) (2025)
- [LoKA: Low-precision Kernel Applications for Recommendation Models At Scale](https://mlsanity.com/paper/2605.10886/) (2026)
- [NestPipe: Large-Scale Recommendation Training on 1,500+ Accelerators via Nested Pipelining](https://mlsanity.com/paper/2604.06956/) (2026)

## Abstract (arXiv)

Sharding a large machine learning model across multiple devices to balance the costs is important in distributed training. This is challenging because partitioning is NP-hard, and estimating the costs accurately and efficiently is difficult. In this work, we explore a "pre-train, and search" paradigm for efficient sharding. The idea is to pre-train a universal and once-for-all neural network to predict the costs of all the possible shards, which serves as an efficient sharding simulator. Built upon this pre-trained cost model, we then perform an online search to identify the best sharding plans given any specific sharding task. We instantiate this idea in deep learning recommendation models (DLRMs) and propose NeuroShard for embedding table sharding. NeuroShard pre-trains neural cost models on augmented tables to cover various sharding scenarios. Then it identifies the best column-wise and table-wise sharding plans with beam search and greedy grid search, respectively. Experiments show that NeuroShard significantly and consistently outperforms the state-of-the-art on the benchmark sharding dataset, achieving up to 23.8% improvement. When deployed in an ultra-large production DLRM with multi-terabyte embedding tables, NeuroShard achieves 11.6% improvement in embedding costs over the state-of-the-art, which translates to 6.6% end-to-end training throughput improvement. To facilitate future research of the "pre-train, and search" paradigm in ML for Systems, we open-source our code at https://github.com/daochenzha/neuroshard