  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
