AutoShard: Automated Embedding Table Sharding for Recommender Systems

by   Daochen Zha, et al.

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


page 1

page 2

page 3

page 4


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

Sharding a large machine learning model across multiple devices to balan...

Learning Compact Compositional Embeddings via Regularized Pruning for Recommendation

Latent factor models are the dominant backbones of contemporary recommen...

Compositional Embeddings Using Complementary Partitions for Memory-Efficient Recommendation Systems

Modern deep learning-based recommendation systems exploit hundreds to th...

Learnable Embedding Sizes for Recommender Systems

The embedding-based representation learning is commonly used in deep lea...

Learning Compressed Embeddings for On-Device Inference

In deep learning, embeddings are widely used to represent categorical en...

LST-Bench: Benchmarking Log-Structured Tables in the Cloud

Log-Structured Tables (LSTs), also commonly referred to as table formats...

Binary Code based Hash Embedding for Web-scale Applications

Nowadays, deep learning models are widely adopted in web-scale applicati...

Please sign up or login with your details

Forgot password? Click here to reset