AutoShard: Automated Embedding Table Sharding for Recommender Systems

Abstract

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 CCS CONCEPTS• Computing methodologies → Reinforcement learning; Machine learning approaches.

FlexShard: Flexible Sharding for Industry-Scale Sequence Recommendation Models

“…ScratchPipe (Kwon & Rhu, 2022) and RecShard (Sethi et al, 2022) tackle the problem of embedding access latency in hybrid CPU-GPU training systems ScratchPipe addresses the problem through the use of a run-ahead GPUside cache to attempt to have all embedding accesses hit in local GPU HBM, while RecShard uses a mixed-integer linear program and the per embedding table distributions to statically place the most frequently accessed rows in GPU HBM. AutoShard (Zha et al, 2022) focuses on the sharding of embedding tables in a multi-GPU only training system, and uses deep reinforcement-learning and a neural network based cost model to perform its placement decisions.…”

“…Furthermore, while the sharding problem has been increasingly explored in recent works due to its importance (Adnan et al, 2021;Lui et al, 2021;Sethi et al, 2022;Zha et al, 2022), they all, to our knowledge, assume that the embedding tables to be sharded are either one-hot-meaning at most one embedding row per table will be accessed per training sample-or sum-pooled-meaning all embedding row accessed within a table by a training sample will be aggregated via summation before proceeding through the model.…”

“…Embedding Pooling. A simple and common way to perform pooling is to do element-wise summation across all embeddings gathered to generate the output pooled embedding; known as sum-pooling (Gupta et al, 2020a;Ke et al, 2020;Kwon & Rhu, 2022;Mudigere et al, 2021;Sethi et al, 2022;Zha et al, 2022;Zhou et al, 2019b;Wilkening et al, 2021). The use of sum-pooling to perform embedding aggregation comes with beneficial properties with respect to the sharding problem.…”

“…To our knowledge however, prior works which have developed novel solutions to the DLRM sharding problem have done so under the assumption that the embeddings gathered will all be sum-pooled before being consumed by the following model layers (Adnan et al, 2021;Lui et al, 2021;Sethi et al, 2022;Zha et al, 2022). Unfortunately, these benefits do not directly apply to such sequence-based pooling methods where each embedding interacts with every other embedding in the sequence, requiring the full sequence to be present on the requesting GPU to perform the operation.…”

Towards Personalized Preprocessing Pipeline Search

“…Reinforcement Learning. Reinforcement learning has shown strong performance in many rewarddriven tasks [38,69,45,34,46,47,19,64,67,23,65,63,61,58,59,62,60,26,10]. It has also been applied to AutoML search [71].…”