DreamShard: Generalizable Embedding Table Placement for Recommender Systems
Authors: Daochen Zha
Louis Feng
Qiaoyu Tan
Zirui Liu
Kwei-Herng Lai
Bhargav Bhushanam
Yuandong Tian
Arun Kejariwal
Xia Hu
Venue: NeurIPS 2022 Accept (2022)
Keywords: Reinforcement Learning
Recommender System
Distributed Training
Embedding Table
Cost Modeling

Abstract: 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.

DreamShard: Generalizable Embedding Table
Placement for Recommender Systems
Daochen Zha'
Louis Feng² Qiaoyu Tan
Zirui Liu'
Kwei-Herng Lai
Bhargav Bhushanam² Yuandong Tian"
Arun Kejariwal Xia Hu¹
Rice University Meta Platforms, Inc. {daochen.. zha, Zirui Liu, khlai, a Aaaii,iiiHHHll.. Haa.Huurricce @rice edu (lofe, bbhushanam, yuandong, akejariwal] @fb. com qytanêtamu. edu
Teexs A&M University
com
Abstract
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 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, present DreamShard, reinforcement learning (RL) approach for embedding table placement. DreamShard achieves the reasoning of operation fusion and generalizability with 1) cost network to directly predict the costs 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 networks 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 production tables. The code is available at pectps://githae com /ww.henzchhhaaaaaaeaaaaaaaaaaaaa batpeehhhddrreasaae dreamshard
Introduction
Embedding learning is a commonly used technique to deal with categorical features in deep rec- ommendation models by mapping sparse features into dense vectors [1] 3 embedding tables can be extremely large due to the large feature sizes 6 For example, in the YouTube recommendation model, single categorical feature contains tens of millions of video IDs 77] the Meta recommendation model demands multi-terabyte memory 8. Distributed training has been adopted place the tables multiple hardware devices such GPUs B 6 9 10 [1]. However, even with distributed training, the embedding tables are often still the efficiency bottlenecks. For instance, embedding lookup is shown to dominate the training throughput in the Meta recommen- dation model [ In our internal production model, which has hundreds of tables, embedding lookup accounts for 48% and 65% of the total computation and communication costs, respectively.
4 5). However, the
How the embedding tables are placed can significantly impact the costs. Figure [ shows the traces of different placement strategies a task of placing 50 tables on 4 devices. Typically, embedding lookup consists of four stages. In the forward pass, the sparse indices are mapped into dense vectors (forward computation), which then sent to the target devices (forward communication).
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

In the backward pass, the gradients of the embed- ding vectors are sent back from the target devices (backward communication) and applied to the em- bedding vectors (backward computation). The ta- bles will easily lead to imbalances if not carefully partitioned. The random placement Figure Ia is bottlenecked by GPU2 with a 56.6 milliseconds latency, while the more balanced placements in Fig- ure Tb and Ic significantly reduce the costs to 42.8 and 35.95 milliseconds, respectively. This work asks: given set of embedding tables, how can identify the best placement of the tables to balance the costs?
Fanert
Ferwar
Bacossid
Racovdd
GPUL
Dovfee
36,8
GPU4
Tanestamp (millicecon))
(a) Random placement
GPUS
Devite
GPUZ
Device placement is essentially a partition problem, which is one of the classical NP-hard combinatorial optimization problems [12]. A recent line of re- search uses reinforcement learning (RL) for device
(millsecond)
placement of computational graphs [ 14 15 16
(b) The existing best human expert strategy
18 19 20). For example, [13 proposed to train an RNN controller with content- content-based attention to
predict the placement. Other studies advanced [[3]
Dene
GPU2
in different ways, such as using hierarchical mod-
GPU2
els [14. more sophisticated RL algorithms [15]
and graph neural networks [16].
GPU4
35.66
10
20
30
50
However, embedding table placement remains to be an open and challenging problem due to the oper- ation fusion [21] of tables and the generalizability requirement. 1) Modern embedding implementa- tions (c.g.. FBGEMM (22)), use single operation to subsume multiple tables for acceleration. The speedup of the fused operation over the sum of the single-table operation costs is not constant and depends the characteristics of the fused tables (e.g., table dimensions). Our analysis finds that the speedups vary significantly across different table combinations, ranging from IX to 3X (Figure 12 in Appendix A.3.2) Thus, we not only need to reason about cost balance but also how the tables should be fused maximize the speedup. 2) In real-world scenarios, the adopted embedding tables and the available devices can change frequently (e.g.. machine learning engineers/researhhers may conduct concurrent experiments with various table combinations and numbers of devices). Thus, a practical algorithm should generalize to tasks with unseen tables, different numbers of tables, and different numbers of devices. It is non-trivial to achieve this with the existing device placement approaches.
Timesture (millisconn))
(c) DreamShard
Figure Visualization of random placement, the existing best human expert strategy, and DreamShard on a task of placing 50 tables on GPUs. The dense computations and commu- nications are omitted the traces because they do not have an imbalance issue. We provide more visualizations in Appendix
To this end, we introduce DreamShard, an RL approach for embedding table placement. DreamShard achieves the reasoning of operation fusion and generalizability with two novel ideas. 1) It learns a cost network directly predict the costs of the fused operations. Specifically, the network takes as input the table features (e.g., table dimension) of each single-table and outputs the computation and communication costs. 2) It trains a policy network by interacting with an estimated Markov decision process (MDP) without real GPU execution, where the states and the rewards are estimated by the predictions of the cost network. Equipped with sum reductions for the table representations and max reductions for the device representations, the two networks can directly generalize to unseen placement tasks with different numbers of tables and/or devices without fine-tuning.
Extensive experiments show that DreamShard outperforms the existing human expert and RNN- based [3] strategies open-sourced synthetic tables [23] and our production tables, achieving up to 19% speedup over the strongest baseline. Moreover, it can generalize to unseen tasks that have different numbers of tables and/or devices with neglectable performance drop (< 0.5 milliseconds). Additionally, its inference is very efficient. It can place hundreds of tables in less than one second.

---

2 Generalizable Embedding Table Placement Problem
The embedding table placement problem seeks a device placemenf all cost (in terms of execution time) is minimized (we provide background for the distributed training of recommendation models in Appendix A.I). Consider M embedding tables {el, e2, .... eM) } and D devices, where e; E IRW denotes the table features that characterize the embedding lookup patterns. In our work, we use 21 table features, including hash size, dimension, table size, pooling factor, and distribution (their definitions are provided in Appendix A.2). A placement
of of all the tables such that the over-
[a1, 02. .... am), where a; E {1,2 2,2..... D), assigns each table to a device. Tett c(a) denote the cost measured on GPUs. The goal of embedding table placement is find the such that c(a) is minimized. Due to the NP-hardness of the partition problem 12), identifying the exact solution demands extensive computational overhead. Thus, the state-of-the-arr algorithms often approximate the optimal parti- tion via sampling with RL [24] [3]. However, sampling remains expensive because obtaining c(a) requires running operations on GPUs. Given that the embedding tables and the available devices can frequently change, we wish to approximate the best a without GPU execution.
Motivated by this, be the space of all the embedding tables. A placement task can be denoted as T, = (E,, D.). where & C E is set of tables, and D; is the number of devices. Given Nirin training tasks Tuaain = (7, T,....,77 T....,
study the generalizable embedding table placement (GETP) problem. Let
Nun... ). and Nies testing tasks TTss (71, T2,-., TNIA J. the goal is to train a placement policy based on Train (GPU execution is allowed during training) such that the learned policy can minimize the costs for the tasks in Trest without GPU execution.
Estimated
3 DreamShard Framework
We present DreamShard, an RL framework based on estimated MDP, tackle the GETP problem. An overview of the framework is shown in Figure 2 The key idea is to formulate the table placement process as an MDP (Section[3.]) and train a cost net- work to estimate its states and rewards (3.2)) A policy network with a tailored generalizable net- work architecture is trained by efficiently interact- ing with the estimated MDP (Section (.3). The two networks are updated iteratively to improve the state/reward estimation and the placement policy.
Final State
Extimated
Sampled
Cost Data
State
Action
Reward
Device
e UB
Policy Network
Placement
L6 MB
Kl. Agont
Real Hardiaare
Figure 2: DreamShard framework. The agent interacts with the estimated MDP, which is trained with the cost data collected from GPUs.
3.1 MDP Formulation
Given embedding tables {c|, C2...., and D devices, we aim to generate a placement a = [ai, Q2,,.., (.s). The key idea is to place the tables one by one at each step, where the state character- izes the tables that have been placed so far, the action is the device ID, and the reward represents the execution time on GPUs. Specifically, at step the state St = {st,a) is all the table features of the tables placed on all the devices, where st.d {eeli E Pa} denotes all the table features corresponding to device d (P& is the set of table IDs that have been placed on device d). We further augment the raw features with cost features which are obtained by collecting the operation computa- tion and communication times from GPUs (Appendix A..3 provides a comprehensive analysis of the cost features). Formally, the augmented state is defined as s = {st, (qt.) where qt.d E TR3 has three elements representing forward computation time, backward computation time, and backward communication time for the current operation in device d (we provide detailed explanations of why forward communication time is excluded in Appendix (A.4). We find that the augmented cost features can significantly boost the performance, evidenced by the ablations in Table 3 The action a, E A an integer specifying the device ID, where A, is the set of legal actions at step l. A device ID is considered legal if placing the current table the corresponding device does not cause memory explosion. The reward re is 0 for all the intermediate steps. and the reward at the final step M is the negative of the cost, i.c., TM c( (a), which encourages the agent to achieve lower cost.
In this work. we focus on GPU devices, where all the GPU devices are identical, which is the most common configuration in production. We defer the mixed scenarios of both GPUs and CPUs future work.
3

---

Unplaced
ds
dg
Device
Device
Device
Step
Step
Step
Step
Step
Figure 3: MDP formulation of embedding table placement.
Predicted Overall (Reward)
Eachard Cost Cost
Prodicted
Overall Cost Head
Softmax
Policy
Max
Comm. Cost
Predicted
Sum
Sum
Stuar
Carl Tousen
Forwardd Cost
Predicid
Festure
Predicted
Device
Device 2
(State)
Device
Device 2
Pollcy Network
Figure 4: DreamShard's cost network (left) and policy network (right).
The procedure illustrated for an example task of placing tables (ei, e2,.., es} on 2 devices in Figure 3 At step 0, no table has been placed so where both qo and qu,1 are zero vectors (i.e., [0. 0, 0)) since all the
{0. (}} and the augmented state
50 = {su. qo 40.1}} qo.0.qo.1}}, computation and communication times are 0 well. Then the action dg makes the MDP transit to the next state $1 ({e)) J. }}} with its corresponding augmented state 81 = {s1. {q1.0- q1,1 }. where qı,0 becomes non-zero vector containing the computation and communication costs by running {ci) and (} on GPUs. We repeat the above process, and finally at step 8, we have 88 fe1,, Ca: C7, 0z) (e2, C3: C5 C.) The corresponding 98,1, 1, and qs,2 the measured times of running (c), C4, C7, Cs and (C2, C32 Cs Co } on two devices. The action sequence a = [a) a2,,,,,.... ag] the generated placement, which is then evaluated on GPUs to obtain the reward.
Discussion 1. The MDP enjoys two desirable properties. 1) The legal action A, can guarantee that the generated placement satisfies the memory constraints. 2) The one-by-one placement enables the agent to be generalized across different numbers of tables. For example, an agent trained on an MDP with very few tables can be applied to another MDP with more tables by simply executing more steps.
Discussion 2. straightforward idea to solve the MDP is to greedily place the current table on the device with the lowest cost at each step, where the cost function can be one of or a combination of the state features (e.g., the sum of the table dimensions, or the sum of all the cost features). However, greedy heuristics are often sub-optimal. Thus, we seek learning-based algorithm to explore various placement possibilities and make comprehensive decisions based on the state features.
3.2 Learning an Estimated MDP
Interacting with the above MDP is computationally expensive since obtaining the cost features and the reward requires GPU execution. Motivated by world models [25] 26). we build an estimated MDP by approximating the cost features and the reward with a cost network. Let feest denote the cost network. Jeos takes as input the raw table features 8t, and predicts cost features {q.a) and the overall cost c(a). fenu is trained with mean squared error (MSE) loss using the cost data collected from the GPUs. Once trained, it can predict the cost features or the reward with single forward pass without GPU execution. However, it is non-trivial to design the architecture of foos because needs

---

Algorithm Training of DreamShard
1: Input:
Training tasks Timin. the number of data collection steps Ncolleat, the number of cost network update steps Noos the batch size for updating the cost network Noath. the number RL update steps NRL, the number of considered episodes in each RL update Nepisode
2: Initialize a cost network, a policy network, and a buffer
3: for iteration = 1. 2, until convergence do
4: for step = 1, 2, Neallect do
Collect cost data from GPUs
5: Randomly sample a training task from Timan
6: Generate placement by interacting with the estimated MDP using the policy network
7: Evaluate the placement on the hardware and store the collected cost data to the buffer
8:
end for
9:
for 1, Neost do
> Update the cost network (no GPU execution) Randomly sample Nbath cost data from the buffer Update the cost network based on MSE loss (Eq. in Appendix B.4.1) end for
10:
11:
12:
13:
for step = 1, 2. NRL do
Update the policy network (no GPU execution)
14:
Randomly sample a training task from Timan
15:
Collect Nepisode episodes by interacting with the estimated MDP using the policy network
16:
Update the policy network based on the policy gradient loss (Eq. 2 in Appendix [B.41])
end for
18:
end for
to accommodate different numbers of devices (i.e., 8t can have variable sizes), and different numbers of tables each device (i.e., $t,d can have variable lengths).
The left-hand side of Figure 4 shows DreamShard's generalizable design of fcos, which is based on two key ideas. First, it uses a shared MIP map raw table features into table representations, For any unseen tables, this MI.P can be directly applied to extract table representations. Second, it enables a fixed-dimension representation for each device with sum reductions (i.c., the clement-wise sum of the table representations in the device), and similarly for the overall representation across devices with max reductions (Appendix [.3 compares different reduction choices and finds that this sum-max combination leads the most accurate prediction). The reduced representations then followed by multiple MLP heads for cost predictions. For unseen tasks with different numbers tables and/or devices, the reductions will always lead to fixed-dimenion device/overall representations, that the prediction heads can be directly applied. Appendix B.T provides more details.
3.3 Training the Policy Network on the Estimated MDP
Generalizable policy network architecture. Let æ be the policy network. æ maps the augmented state = {st, (qraa d=1] to action a,, i.e., ae = x (.t). π also adopts generalizable design, shown in the right-hand side of Figure 4 Like feom: uses shared MLP and sum reductions to produce fixed-dimension representation, which is then concatenated with the cost features to obtain the device representation. To accommodate the potentially variable action space (i.e., the number of available devices may vary), a shared MI. will process each device representation separately to obtain a confidence score, followed by a Softmax layer to produce action probabilities. This design allows to generalize across different numbers of devices. Appendix B.2 provides more details.
Training and inference. Algorithm summarizes the training procedure of DreamShard, which iteratively executes the following: 1) collect cost data from GPUs based on the placements generated by the current policy, 2) update the cost network with the previously collected cost data, and 3) update the policy network by interacting with the current estimated MDP. Throughout the training process, the estimated MDP gradually becomes more accurate, and the resultant policy network tends to generate better placements. Appendix B.4.2 provides more details of the training procedure. For the inference, the trained cost network and policy network can directly applied to unseen tasks to generate placements without GPU execution, which is summarized by Algorithm 2 Appendix B.4.3
5

---

4 Experiments
Our experiments aim to answer the following research questions. RQI: How does DheamShard compare with the existing human expert and RL-based placement strategies? RQ2: Can DheamShard generalize to placement tasks with different numbers of tables and/or devices? RQ3: How efficient is the training of DreamShard? RQ4: How do the hyperparameters influence the performance of DreamShard? RQ5: How does each component of DreamShard contribute to the performance? RQ6: How accurate is the estimated MDP and to what extend can it accelerate the training and inference?
4.1 Experimental Setup
Datasets. Academic recommendation datasets are often too small enable meaningful evaluation because the cost will always be very small no matter how the tables are placed. Thus, we use two industrial-scale datasets. DLRM is a large-scale synthetic dataset with 856 tables, recently released by Meta. It shares memory access reuse patterns similar to those arising in Meta production workloads. Prod is an internal large-scale dataset for production recommendation models. It has a similar scale as DLRM. The main difference is that DI.M only has a fixed dimension for all the tables, while Prod is more challenging with diverse table dimensions, ranging from 4 to 768. For reproducibility, we mainly focus on the DLRM dataset since it is open-sourced. We only report the main results on the Prod dataset for verification purposes. We provide more details in Appendix C
Baselines. We compare Shaad against human expert strategies from previous work (27 8, 28 including size-based, dim-based, lookup-based, size-lookup- We also include an RNN-based RL algorithm [13], which uses RNN architecture to map operators to devices. Since the feature extraction layers of the RNN-based method were designed for other operations instead of embedding tables, for a fair comparison, we adapt [3] by making it have the same feature extraction layers as in DreamShard. We provide more details in Appendix D
greedy balancing strategies.
Configurations,
To evaluate the generalizability of DreamShard, we randomly divide the tables into training pool grain and testing pool Erstt The two pools have the same number of tables but they are not overlapped. A sharding task T; is constructed by randomly sampling subset of [& tables from a pool, where the number of tables |& E {10, 20, 30, 40, 50, 60, 70, 80, 90, 100) for the DLRM dataset, and |&&] E {20, 40, 80} for the Prod dataset. For all the experiments, we randomly sample 50 training and 50 testing tasks from 8 Erran and grest, respectively. DreamShard is trained on the training tasks and will be evaluated on unseen tables in the testing tasks. We denote placement tasks with different numbers of tables and devices using the format of dataset-nur num_tables tables (num_ dev_ddiiiesss For example, DLRM-30 (4) suggests that there are 30 tables sampled from the DLRM dataset in each training/testing task with 4 available devices. We provide more details in Appendix
Implementation Details. We use the same hyperparameters for all the experiments with Newlllett = 10. Noost 300, Nbatch 64, NRR. = 10, Nepisode- = 10, 10 training iterations, and an entropy weight of 0.001 in the policy gradient. 2080 Ti GPUs and V V100 GPUs are used for the DLRM (except that we use V100 for experiments with GPUs) and Prod datasets, respectively. All the experiments are run 5 times, and we report the mean and the standard deviation. Appendix B provides more details.
4.2 Results and Analysis
Evaluation of DreamShard against baselines (RQ1). We perform qualitative and quantitative comparisons of DreamShard against the baselines. Qualitatively, Figure D visualizes the traces of DreamShard and the baselines on one of the tasks from DLRM-50 (4). DreamShard achieves significant better overall cost than the best baseline (35.95 vs. 42.8) with 1) better balance of forward and backward commutation workloads, and 2) less communication time, possibly due better balance of table dimensions. Quantitatively, Table 1 presents comprehensive evaluations on tasks with different numbers of tables and devices on the DLRM and the Prod datasets. Observations: 1) DreamShard outperforms the baselines on all the tasks. 2) DreamShard shows strong generalizability on unseen tables, achieving the same level of performance on all the testing and training tasks. 3) DreamShard appears to be more ad advatagcouus on harder tasks. Specifically, DreamShard achieves more improvement over the baselines on tasks with more tables/devices on the Prod dataset. In particular, DreamShard achieves 19% improvment over the strongest baseline on Prod-80 (8). 4)

---

Table 1: Overall cost comparison in milliseconds and relative speedups over random placement, averaged over 50 different randomly sampled tasks. The placement tasks are denoted as dataset- num_ tables (num_devicel _evicces.. For example, DIRM-30 (4) suggests that 30 tables are sampled from the DLRM dataset for each task with 4 GPUs. More results are provided in Appendix F
Tisk
No Stateeyyyy
Random
Sior-basee
Dim-
Human Experts
Lookup-Ssseed
Size-lookup-- boser
RNN-
RL
DreamShard
Train 240106 Test 23.0::05
22.72200
21.3 (-12
19.1000
758
22.4::0.5 (47.1
18.610.1 (+29,0%)
DLRM-20 (4)
21.7:00.0
19.9:00.0
(+15 66
18.3::000 790
1.42:000+255000
20.91:033
17.6.10.3 (+30.7%
DLRM-10 (4)
Train 41.34-02
Test [41. 105
39.6--0000 (+4.36))
40.3:200
(+2.0%)
33744-0.1 (05
37.
0.0
(+ 10
33.6+:0.0 (x22.99)
33.020:1
(424.5%)
$3.6+=0.1 (+22.9%)
33.02:0000
(12550 (1524407)6)4445)
59.0 1.1 (44.8%)
32.420..3 (+26.9%
DIRM-60 (4)
Train 557.7400
Test
58.1
56.6220.1
59.6 (:2,5%)
(+1.9%)
(49.1%) 52.99
49.24:0.1
48.75002
(+17.3%%
49.1
(+:18.35)
55.5-5009 (14.09)
56.0--00777 (63.8196
(7.9
(426.22
Train 75.7
76.0:00.0
70,0
(+8.19)
6653--00.. 0.00000
(+15.9%))
22.24:22.7
DLKM-80 (4)
Test 74.5:::0.8
77.7:00. (-4.15)
69.92:0:4
(+6.6%)
64.1 (+16.25)
(5.1::0.000414.055
72:9:12.4 (+2.255)
(2.7:10.33
Train 1.7
14.1 (-2.4%)
81:2::0444
(-13
94.52-1076--.995
78.41-664177 (411.1%)
DLRM-100 (4)
Test 94.55:6655
05.42.00 (-0.9%%)
84.72:0 (+11 68%
79.520.33
30.8::0.3
94 04.853330000.85)
77.8:10.8 (421.5%))
DLKM-40
Thain 15.62:0
Test 15:24:02
14.1 (-10699))
1..52000 (14.89))
13.2 (+15
13.400011
(-16.4%)
0.0 (-59.2%%)
(-60.09))
9.9:0000
9.5::000
(457 6%%
(460.09)
16 16:24:0..8 (-3.79)
1600::111
(4.9.2%))
DLRM 8044-80000
Train 25.0+:0.2
Text 25
24.0:00000 (44.286)
25.6 (-1.69))
2%8 (4
20.8 (431 25%
17.1
(146.28)))
16.9:0011
(149.1
5144+:396551445)
16.1::0.44 (456.5%)
16.1:
DIRM-1100 (8)
Jest
$3.5
34.00::03
(4..:10))
35.0:+0.0
29.2: (+ )
29.8:2000 (-14.100
25.75:0:0
24.5::0.0
(+41..4%
24.5
(+ (+36.75))
5k.66:2.7 (-42.00))
58.75:5.1 (-42.9%)
23.3:10.. (+45.%%)
22.8 2.88:-0.2 (+46.9%
DLKM-160 (8)
Train
Test 41.11::000
42.8:10..3
41.6:10.0 (+2.9%%
(.3.11)00
339,0:80.00 (+9.5%)
06.4:2000 (+12.9%)
32.00::00
3..8:1000 (+33.496)
(+33, 796)
32.7.10000 (+20.9%)
31.6:10.000
(+20011
59.3-1.544 (-30.79))
(5.3:4-3.5(226669))
29.6.1.0.2.0 (+36.9%
10.3-10.. (+41.3%
Train 51
48. 28.2::0.000
(46 8%)
08.0020000 (+7 (47.5%)
0..0 (-320000333400
(- 1%)
66
(%9)
(4.3.44)
DLKM-200 (8)
Test 50.7200
50.80000
(-0.25)
04.8:20.0 (-13 25%
18.0
(-33.498)
70.44::28
(-788.99)
36.44:00.3 (4.393%%
Train (41.110077
(-..2000-4488))
337.0:0000 (+11 669)
44.22:000
(5.8::00000
5800::003 (48.799)
366.9:10.3 (+13.8%)
Prod-20 (2)
Test 4222004
46.12000 (-7.2%)
39.5020. (48.4%)
45.9 (-6.8%8)
45.7::0.0 (-6.395)
(18.99)))
(7..5::0.2(4444111
Prud-10 (4)
Test 33.32203
Train| 35.1+00.3
39.4 (-1..9%)
31.33+0.0 (2 198 33.52200 (1..2000-12222))
36.4-:000
37.42:00 (42.44%
40.12:000 (44.59)
$3.99:2255
(672228844448))
[2.3++-0.3
50.41:0.77 (+26.0%)
Train 43.2:::02
41.3:10.00
33.00:00.0
(%)
43.7
(66:166
33.6:1..9 (+22..6))
Prod-80
Test 47777+04
(3..9:00.001111996))
41.9:+0.0
13.8
46.1
(43.5%%
29.6-000000 (-3.8%%
62.544:422 (-23.7%)
35.2-+-0.8 (+35.5%%
Table 2: Generalization performance of DreamShard on target tasks w.r.t. to different numbers of tables (top) and devices (bottom), averaged over 50 randomly sampled tasks. A provides the results of more challenging tasks, such as mixed table-wise and device-wise transferring scenarios.
Source Task-
Target Task
Rundom Best baseline strategy| Dreammhhard (trained target task) DreamShard (Imined
task)
DLRM-20 (4) + D00 (4)55.666665
79.5++0.3 (+18.9%)
77.8+0..8 (+21.5%)
77.9+0.44 (+21.3))))))3))
DLRM-10 (4) +DLMM-0 (4)
74.5=00.8
64.1-:0.2 (+16.2%)
62,7=0.33 (+18.8%)
62.7==0.5 (+18.8%)
DLRM (4) + DLLMM-00 (4)
41.1
33.0:1:.1 (+24.5%)
3..4::0.3 (+26.9%)
32.4::0.2 (+26.9%%
DIRM- 100 (4) DI.RM 20 (4)
23.0.:0.5
18.3:10.0
(+25,7%)
17.6110.2 (+330%%%
17.7.00.3 (+29.9%%
DLRM-20 (4) + DLRM-20 (2)
29.9::0.4
26.0::0.0
(+15.00%
25.8:00. (+18.9%)
25.8-00.1 (+15.9%)
DLRM-10 (4) DLRMM-0 (2)
58.6=0077
52.4#000
(+11.8%)
51.9=:011 (+12.9%)
52,0==033 (+12.7%)
DLRM-20 (2) + DLLMM-20 (4)
23.0:20.5
18.3::0.0
17.6::0.2 (+30.7%)
17.8+:0.3 (+29.2%)
DLRM 40 (2) -DLRM-MA 40
41.1
33.0.10.1 (+24.5%)
32.4 (+26.9%)
32.6.0.33 (+26.1%)
RNN-based method is only better than the random strategy on tasks with few tables/devices, but is worse on harder tasks. A possible reason that RNN-based algorithm does not have a cost network, and using RI. alone could lead to unstable performance. 5) Lookup-based strategy is the best baseline on the DLRM dataset, while dim-based strategy is better on the Prod datasct. A potential reason is that the tables in the Prob dataset have very diverse table dimensions, while the tables in the DLRM dataset have the same dimension. As such, the dimensions on Prod tasks can more casily become imbalanced, leading to poor communication efficiency. Dim-based strategy can better balance the dimensions, which leads to a better overall performance. DreamShard outperforms the baselines on both the DLRM and the Prod datasets, showing its flexibility in dealing with different scenarios.
Analysis of generalizability (RQ2). In Table 2 we directly apply a DreamShard model trained from one task to another task without any any fine-tuning (the rightmost column), where the source and the target tasks have different numbers of tables and/or devices. DreamShard shows neglectable performance drop, suggesting that it is generalizable across different numbers of tables and/or devices.
7

---

Cost on hets
haars
41.0
20
3 40.5
60
60
40.5
50
50
40.0
80
40
2
50 100 150
200 250 300 350 Time Seconds
50:000
200
500
Iteration
RL Update Stops Nya
Cost Network Update
Figure 5: Performance (1) of DreamShard the DLRM-50 (4) datasets w.r.t the numbers of iterations (left) and rumning time (right). Curves for other tasks are in Appendix H
Figure 6: Impacts of the number of RL update steps NRL (left) and cost network update steps Noost (right) on the DLRM-50 (4) dataset. Curves for other tasks are in Appendix [
Table 3: Ablation study of DreamShard, Results on other tasks are Appendix J
Task
wio dim |wo bash size/wlo pooling factor wwo table size w/o distribution] wlo cost
w/ RNN
DreamShard
DLRM-50 (4)
Train 40.8::0.4|
40.7+0.11
46,3=0.3
40,8=00.4 40.6.10.7
40.6+10.2 40.51.0.
47.5
40,5
40.4::0.5
Test 40.9.10.66
40.6.10.3
47.2.1.11
06.3.10.1
0..5
40.4.10.6
Analysis of training efficiency (RQ3). Figure ploos the performance of DreamShard w.r.t. the number of iterations and running time seconds on four 1080Ti GPUs. The training of DreamShard is highly efficient. On DLRM-50 (4), it can achieve strong performance in less than 5 iterations or 200 seconds. Note that we only need to re-train or fine-tune DreamShard when the table pools have significant changes. Once trained, it only needs a forward pass for inference.
Hyperparameter study (RQ4). We study the impacts of two key hyperparameters: 1) NRI. which controls the RL update frequency, and 2) Nesa, which determines the cost network update frequency. We vary one of them with the other fixed, shown in Figure 6 Increasing NRL. Nosst will both lead to improvement, suggesting that both the cost network and the policy network need to be sufficiently trained. However, when NRL and Noost are large enough, increasing them will not bring more improvement. Considering that larger values will lead to more computational costs, we set NRL. 10 or Na 300 trade-off between the performance and the training efficiency.
Ablation study (RQ5). We study the importance of each table feature and check whether the RNN architecture helps, with the following ablations. 1) We remove each of the features in the state. 2) We add an RNN upon the device representation in the policy network. We makes several observations from the results in Table 3 1) Cost features play a significant role, which demonstrates the effectiveness of our proposed augmented state. 2) The most contributing table features are the pooling factor and the dimension, which aligns with our intuitions since these two feature are the determining factors of computation and communication workloads. 3) Using more features leads to consistently good performance. 4) While the policy makes decisions sequentially, RNN does not provide clear benefits. This is why we have kept the architecture simple with only MLP in DreamShard.
Study of the estimated MDP (RQ6). First, we study how many data points are required to train an accurate cost network, and how accurate the cost network needs to be enable strong policy. Specifically, we randomly sample 10,000 cost data points from the DLRM-50 (4) dataset. Then use 20% for testing, and vary the size of the training data train the cost network. Further, we fully train policy network with 100 iterations based on each of the trained cost networks. We make two observations from the results in Figure 7 1) As expected, more data points lead more accurate cost network. 2) Interestingly, after around 100 data points, the policy network does not keep improving even though the cost network becomes more accurate. Thus, we only need sufficiently (but not perfectly) accurate cost estimation to achieve the best performance. This also partially explains why DreamShard can generalize: even though the cost network could be not very accurate on unseen tables, DreamShard can still find strong placements with the policy network. Second, we study the necessity of the estimated MDP. We consider a variant that obtains the cost features and rewards directly from GPUs, shown in Figure 3 Using the estimated MDP can make the training and the inference orders of magnitudes faster, while achieving the same level of performance. In particular, the inference time is less than one second even with a hundred tables.
8

---

200
120
80
Infeence Time in Se53305
Edi Ensirated HDW
Terting MSE
1.5
100
70
Entruned NoP
300
A/a Estivated MD2
200
60
200
50
Cos on 50
40
36"
10a
100
Number Training Data Poincs
Number Faiiing Data Points
Figure 7: MSE of the cost network w.r.t. the num- ber of training data points collected from GPUs (left), and the corresponding performance of fully trained policy network based on the trained cost network (right) on the DLRM-50 (4) dataset.
Figure 8: Training curves w/ or w/o the estimated MDP on DLRM-50 (4) dataset (left), and their inference times w.r.t the number of tables (right). Using cost network to estimate the MDP leads orders of magnitudes faster training and inference.
5 Related Work
Embedding tables. Embedding tables are commonly used to deal with sparse features in recommen- dation models 2 3
30 31. However, the extremely large embedding tables are often the storage and efficiency bottlenecks 66 7 3 6
3
29
9 I0 In 32]. To our knowledge, the only two studies that target the embedding table placement problem are RecShard (27] and our previous work AutoShard [33] RecShard approaches table placement at per-row gramularity through exploitation of the underlying feature distributions. RecShard leverages these distributions, along with charac- teristics of the training system, to shard embedding tables across tiered memory hierarchy using a mixed integer linear program, with more frequent rows placed in GPU HBM and the remaining placed in CPU DRAM. In contrast, our work develops neural cost network for cost estimation and an RL-based optimization algorithm, and focuses on sharding across a single memory layer as opposed a tiered memory hierarchy. AutoShard also leverages RL for embedding table sharding. However, it only balances the computational costs, which will lead to sub-optimal solutions since communication also account for significant costs. In addition, our DreamShard is more efficient in training than AutoShard due to the design of the estimated MDP, Another line of work focuses on which is orthogonal to our work since
reducing the embedding table sizes [34 35 36 37 38 39 40
DreamShard is also applicable to compressed tables.
Device placement optimization. The existing device placement techniques can be mainly grouped into two categories: RL-based algorithms, and cost modeling methods. 1) RL-based algorithms treat the device placement black box and optimize the cost objective in a trial-and-error fashion [13 14
15 16
19 20). Unfortunately, these methods are computationally expensive and often require training from scratch unseen tasks. While Placeto [16] shows generalizabily on computational graphs with graph embeddings, it can not deal with various table combinations since there is no graph structure, and it not handle different numbers of devices. 2) Cost modeling methods build cost model to reflect the real performance and adopt offline algorithms (c.g., scheduling, and dynamic 45). However, the cost model could be inaccurate. In particular, they can not deal with the operation fusion of embedding tables. Whereas, DreamShard combines the advances of RL with a neural cost model for accurate cost prediction.
programming) optimize the placement [4] 42 43
Deep RL. Deep RL has recently made significant progress in games 46
48
50
51
52 [33] (54)
Our work is related to using RL optimize machine learning model designs, such neural architecture search [55]
[9]. data augmentation [60]. data sampling [61] 62,, pipeline 65 However, these methods often only focus on one task and can not generalize to
56
57
58
search 63
unseen tasks. Our work is related to meta-learning [66] [7]66]) 68). Instead of performing meta-learning for machine learning tasks, we focus on machine learning system design. Our work also related to solving combinational optimization problems with RL [24] 669. Unlike the above studies, we show that RL can tackle a practical problem of embedding table placement and the learned policy is generalizable.
Conclusions and Future Work
We present DreamShard for embedding table placement in recommender systems. We formulate the problem an MDP which places the tables one by one each step. Then we leverage RL to

---

solve the MDP. To accelerate the training and the inference, we build an estimated MDP by training cost network to approximate the state features (i.e., computation and communication times) and the reward (i.e., the overall cost), leading to orders of magnitudes faster training and inference speeds. Extensive experiments on the open-sourced DLRM dataset and our production dataset demonstrate the superiority of DreamShard over the existing algorithms. Moreover, DreamShard shows strong generalizability, making it a desirable choice in real-world applications. In the future, we will extend DreamShard to tiered memory hierarchy and large-scale training clusters with complex topologies.
Acknowledgements
The work is, in part, supported by NSF (#IIS-2224843.. The views and conclusions in this paper those of the authors and should not be interpreted as representing any funding agencies. We would also like thank the helpful feedback from the anonymous reviewers.
References
[1] Shuai Zhang, Lina Yao, Aixin Sun, and Yi Tay. Deep learning based recommender system: A survey and new perspectives. ACM Computing Surveys (CSUR), 52(1):1-38, 2019.
[2] Heng-Tze Cheng, Levent Koc, Jeremiah Harmsen, Tal Shaked, Tushar Chandra, Hrishi Aradhye Glen Anderson, Greg Corrado, Wei Chai, Mustafa Ispir, et al. Wide & deep learning for recommender systems. In DLRS Workshop. 2016.
[3] Maxim Naumov, Dheevatsa Mudigere, Hao-Jun Michael Shi, Jianyu Huang, Narayanan Sun- daraman, Jongsoo Park, Xiaodong Wang, Udit Gupta, Carole-Jean Wu, Alisson G Azzolini, et al. Deep learning recommendation model for personalization and recommendation systems. arXiv preprint arXiv: 1906: 00091, 2019.
[4] Xiangnan He, Lizi Liao, Hanwang Zhang, Liqiang Nie, Xia Hu, and Tat-Seng Chua. Neural collaborative filtering. In WWW, 2017.
[5] Weiping Song, Chence Shi, Zhiping Xiao, Zhijian Duan, Yewen Xu, Ming Zhang, and Jian Tang. Autoint: Automatic feature interaction learning via self-attentive neural networks. In CIKM, 2019.
[6] Weijie Zhao, Deping Xie, Ronglai Jia, Yulei Qian, Ruiquan Ding, Mingming Sun, and Ping Li. Distributed hierarchical gpu parameter server for massive scale deep learning ads systems. In MLSys, 2020.
[7] Paul Covington, Jay Adams, and Emre Sargin. Deep neural networks for youtube recommenda- tions. In RecSys, 2016.
[8] Bilge Acun, Matthew Murphy, Xiaodong Wang, Jade Nie, Carole-Jean Wu, and Kim Hazelwood. Understanding training efficiency of deep learning recommendation models at scale. In HPCA, 2021.
[9] Amazon dsstne: Deep scalable sparse tensor network engine.
https://gi http//ggtttuuuuu.t com/
[10] Udit Gupta, Carole-Jean Wu, Xiaodong Wang, Maxim Naumov, Brandon Reagen, David Brooks, Bradford Cottel, Kim Hazelwood, Mark Hempstead, Bill Jia, et al. The architectural implications of facebook's dnn-based personalized recommendation. In HPCA, 2020.
[11] Maxim Naumov, John Kim, Dheevatsa Mudigere, Srinivas Sridharan, Xiaodong Wang, Whitney Zhao, Serhat Yilmaz, Changkyu Kim, Hector Yuen, Mustafa Ozdal, et al. Deep learning training in facebook data centers: Design of scale-up and scale-out systems. arXiv preprint arXiv:2003.09518 2020.
[12] Wikipedia. Partition problem- org/n/index, htt.a/acwe..aoiit..ooororrrooo.. wttppp/ww....eeerooooooooooooaa. accessed 05-May-2022).
Wikipedia, the free encyclopedia. http::/ttthp//ee/ //op::///.. en. wikipedia 2022. [Online;
[13] Azalia Mirhoseini, Hieu Pham, Quoc V Le, Benoit Steiner, Rasmus Larsen, Yuefeng Zhou, Naveen Kumar, Mohammad Norouzi, Samy Bengio, and Jeff Dean. Device placement opti- mization with reinforcement learning. In ICML, 2017.
10

---

[14] Azalia Mirhoseini, Anna Goldie, Hieu Pham, Benoit Steiner, Quoe V Le, and Jeff Dean. A hierarchical model for device placement. In ICLR, 2018.
[15]
Yuanxiang Gao, Li Chen, and Baochun Li. Spollghtt: Optimizing device placement for training deep neural networks. In ICML, 2018.
[16] Ravichandra Addanki. Shaileshh Bojja Venkatakrishnan. Shreyan Gupta, Hongzi Mao, and Mo- hammad Alizadeh, Placeto: learning generalizable device placement algorithms for distributed machine learning. In NeurlPS, 2019.
[17] Yanqi Zhou, Sudip Roy, Amirali Abdolrashidi, Daniel Wong, Peter C Ma, Qiumin Xu, Ming Zhong. Hanxiao Liu, Anna Goldie, Azalia Mirhoseini, al. Gdp: Generalized device placement for dataflow graphs. arXiv preprint arXiv:1910.01576 2019.
[18] Aditya Paliwal. Felix Gimeno, Vinod Nair, Yujia Li, Miles Lubin, Pushmeet Kohli, and Oriol Vinyals. Reinforced genetic algorithm learning for optimizing computation graphs. In ICLR, 2019.
[19] Yuanxiang Gao, Li Chen, and Baochun Li. Post: Device placement with cross-entropy mini- mization and proximal policy optimization. In NeurIPS, 2018.
[20] Anna Goldie and Azalia Mirhoseini. Placement optimization with deep reinforcement learning. In ISPD, 2020.
[21] Wei Niu, Jiexiong Guan, Yanzhi Wang, Gagan Agrawal, and Bin Ren. Dnnfusion: accelerating
deep neural networks execution with advanced operator fusion. In PLDI, 2021.
[22] Daya Khudia, Jianyu Huang, Protonu Basu, Summer Deng, Haixin Liu, Jongsoo Park, and Mikhail Smelyanskiy. Fbgemm: Enabling high-performance low-precision deep learning inference. arXiv preprint arXiv:2101 05015, 2021.
[23] github. Embedding lookup synthetic dataset, 2021.
[24] Irwan Bello, Hieu Pham, Quoc Le, Mohammad Norouzi, and Samy Bengio. Neural combina- torial optimization with reinforcement learning. arXiv preprint arXiv: 1,009940 2016.
[25] David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 10122, 2018.
[26]
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control:
Learning behaviors by latent imagination. In ICLR, 2019.
[27] Geet Sethi, Bilge Acun, Niket Agarwal, Christos Kozyrakis, Caroline Trippel, and Carole- Jean Wu. Recshard: Statistical bbsteree-baeddddd memory optimization for industry-scuu neural recommendation. In ASPLOS, 2022.
[28] Michael Lui, Yavuz Yetim, Özgür Özkan, Zhuoran Zhao, Shin- Shin-Yeh Tsai, Carole-Jean Wu, and Mark Hempstead. Understanding capacity driven scale-out neural recommendation inference. In ISPASS, 2021.
[29] Jizhe Wang, Pipei Huang, Huan Zhao, Zhibo Zhang, Binqiang Zhao, and Dik Lun Lee. Billion- scale commodity embedding for e-commerce recommendation in alibaba. In KDD, 2018.
[30] Fuzheng Zhang, Nicholas Jing Yuan, Defu Lian, Xing Xie, and Wei-Ying Ma. Collaborative knowledge base embedding for recommender systems. In KDD, 2016.
[31] Weijie Zhao, Jingyuan Zhang, Deping Xie, Yulei Qian, Ronglai Jia, and Ping Li. Aibox: Ctr prediction model training on a single node. In CIKM, 2019.
[32] Xiangru Lian, Binhang Yuan, Xuefeng Zhu, Yulong Wang, Yongjun He, Honghuan Wu, Lei Sun, Haodong Lyu, Chengjun Liu, Xing Dong, et al. Persia: A hybrid system scaling deep learning based recommenders up to 100 trillion parameters. arXiv preprint arXiv:2111.05897, 2021.
[33] Daochen Zha, Louis Feng, Bhargav Bhushanam, Dhruv Choudhary, Jade Nie, Yuandong Tian, Jay Chae, Yinbin Ma, Arun Kejariwal, and Xia Hu. Autoshard: Automated embedding table sharding for recommender systems. In KDD, 2022.
[34] Caojin Zhang, Yicun Liu, Yuanpu Xie, Sofia Ira Ktena, Alykhan Tejani, Akshay Gupta, Pranay Kumar Myana, Deepak Dilipkumar, Suvadip Paul, Ikuhiro Thara, et al. Model size reduction using frequency based double hashing for recommender systems. In RecSys, 2020.
11

---

[35] Hao-Jun Michael Shi, Dheevatsa Mudigere, Maxim Naumov, and Jiyan Yang. Compositional embeddings using complementary partitions for memory -effrceeefffiiettttt recommendation systems. In KDD, 2020.
[36] Xiangyu Zhao, Chong Wang, Ming Chen, Xudong Zheng, Xiaobing Liu, and Jiliang Tang. Autoemb: Automated embedding dimensionality search in streaming recommendations, In SIGIR, 2020.
[37] Manas R Joglekar, Cong Li, Mei Chen, Taibai Xu, Xiaoming Wang, Jay K Adams, Pranav Khaitan, Jiahui Liu, and Quoe V Le. Neural input search for large scale recommendation models. In KDD, 2020.
[38] Siyi Liu, Chen Gao, Yihong Chen, Depeng Jin, and Yong Li. Learnable embedding sizes for recommender systems. In ICLR, 2021.
[39] Wang-Cheng Kang. Derek Zhiyuan Cheng, Ting Chen, Xinyang Yi, Dong Lin, Lichan Hong, and Ed H Chi. Learning multi-granular quantized embeddings for large-vocab categorical features in recommender systems. In WWW, 2020.
[40] Wang-Cheng Kang. Derek Zhiyuan Cheng, Tiansheng Yao, Xinyang Yi, Ting Chen, Lichan Hong, and Ed H Chi. Learning to embed categorical features without embedding tables for recommendation. In KDD, 2021.
[41] Eugene L Lawler, Jan Karel Lenstra, Alexander HG Rinnooy Kan, and David B Shmoys. Sequencing and scheduling: Algorithms and complexity. Handbooks in operations research and management science, 4:445-522, 1993.
[42] Zhihao Jia, Matei Zaharia, and Alex Aiken. Beyond data and model parallelism for deep neural networks. In MLSys, 2019.
[43] Zhihao Jia, Sina Lin, Charles R Qi, and Alex Aiken. Exploring hidden dimensions in paralleliz- ing convolutional neural networks. In ICML, 2018.
[44] Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R Devanur, Gregory R Ganger, Phillip Gibbons, and Matei Zaharia. Pipedream: generalized pipeline parallelism for dnn training. In SOSP. 2019.
[45] Jakub M Tarnawski, Amar Phanishayee, Nikhil Devanur, Divya Mahajan, and Fanny Nina Par- avecino. Efficient algorithms for device placement of dnn graph operators. In NeurlPS, 2020.
[46] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint 13 2013.
[47] David Silver, Julian Schritwieserr, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, 550(7676):354 359, 2017.
[48] Daochen Zha, Jingru Xie, Wenye Ma. Sheng Zhang, Xiangru Lian, Xia Hu, and Ji Liu. Douzero: Mastering doudizhu with self-play deep reinforcement learning. In ICML, 2021.
[49] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707 nee::07770644 06347, 2017.
[50] Daochen Zha, Kwei-Herng Lai, Kaixiong Zhou, and Xia Hu. Simplifying deep reinforcement learning via self-supervision. arXiv preprint arXiv:2106.05522 2021.
[51] Timothy P Lillicrap, Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforeement learning. In ICLR, 2016.
[52] Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAl Pieter Abbeel, and Wojeiech Zaremba. Hindsight experience replay. NeurIPS, 2017.
[53] Daochen Zha, Wenye Ma, Lei Yuan, Xia Hu, and Liu. Rank the episodes: A simple approach for exploration in procedurallly
environments. In 2021.
[54] Daochen Zha, Kwei-Herng Lai, Songyi Huang, Yuanpu Cao, Keerthana Reddy, Juan Vargas, Alex Nguyen, Ruzhe Wei, Junyu Guo, and Xia Hu. Rleard: a platform for reinforcement learning in card games. In IJCAI, 2021.
12

---

[55] Barret Zoph and Quoe V Le. Neural architecture search with reinforcement learning. In ICLR, 2017.
[56] Yicheng Wang, Xiaotian Han, Chia- Chang, Daochen Zha, Ulisses Braga-Neto, and Xia Hu. Auto-pinn: Understanding and optimizing physics-informet neural architecture. arXiv preprint arXiv:2205. 13748, 2022.
[57] Yuening Li, Zhengzhang Chen, Daochen Zha, Kaixiong Zhou, Haifeng Jin, Haifeng Chen, and Xia Hu. Automated anomaly detection via curiosity-
search and self- aufffiitttiion learning. IEEE Transactions on Neural Networks and Learning Systems, 2021.
(58) Hieu Pham, Melody Guan, Barret Zoph, Quoe Le, and Jeff Dean. Efficient neural architecture search via parameters sharing. In ICMI., 2018.
[59]
Yuening Li, Zhengzhang Chen, Daochen Zha, Kaixiong Zhou, Haifeng Jin, Haifeng Chen, and Xia Hu. Autood: Neural architecture search for outlier detection. In ICDE, 2021.
[60] Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoe Le. Autoaugment: Learning augmentation policies from data. arXiv preprint arXiv:1805.09500 2018.
[61] Daochen Zha, Kwei- Lai, Mingyang Wan, and Xia Hu. Meta-aad: Active anomaly detection with deep reinforcement learning. In ICDM, 2020.
[62] Daochen Zha, Kwei- Sevi--Heng&& Lai, Qiaoyu Tan, Sirui Ding, Na Zou, and Xia Hu. Towards automated imbalanced learning with deep hierarchical reinforcement learning, In CIKM, 2022.
[63] Iddo Drori, Yamuna Krishnamurthy, Remi Rampin, Raoni de Paula Lourenco, Jorge Piazentin Ono, Kyunghyun Cho. Claudio Silva, and Juliana Freire. Alphad3m: Machine learning pipeline synthesis. arXiv preprint arXiv:2 //
1. 1.02508, 2021.
[64] Daochen Zha, Zaid Pervaiz Bhat, Yi-Wei Chen, Yicheng Wang, Sirui Ding, Anmoll Kumar Jain, Mohammad Qazim Bhat, Kwei-Herng Lai, Jiaben Chen, et al. Autovideo: An automated video action recognition system. In IJCAI, 2022.
[65] Kwei-Herng Lai, Daochen Zha, Guanchu Wang, Junjie Xu, Yue Zhao, Devesh Kumar, Yile Chen, Puray Zumkhawaka, Minyang Wan, Diego Martinez, et al. Tods: An automated time series outlier detection system. AAAI, 2021.
[66] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-leaming for fast adapta- tion of deep networks. In ICML. 2017
[67] Yue Zhao, Ryan Rossi, and Leman Akoglu. Automatic unsupervised outlier model selection. NeurIPS, 2021.
[68] Yue Zhao and Leman Akoglu. Towards unsupervised hpo for outlier detection. arXiv preprint A208 / 1727, 2022.
[69] Thomas Barrett, William Clements, Jakob Foerster, and Alex Lvovsky. Exploratory combinato- rial optimization with reinforcement learning. In AAAI, 2020.
[70] Guorui Zhou, Na Mou, Ying Fan, Qi Pi, Weijie Bian, Chang Zhou, Xiaoqiang Zhu, and Kun Gai. Deep interest evolution network for thick-through rate prediction. In AAAI, 2019.
[71] David C Liu, Stephanie Rogers, Raymond Shiau, Dmitry Kislyuk, Kevin C Ma, Zhigang Zhong, Jenny Liu, and Yushi Jing. Related pins at pinterest: The evolution of a real-world recommender system. In WWW, 2017.
[72] Carlos A Gomez-Uribe and Neil Hunt. The netflix recommender system: Algorithms, business value, and innovation. ACM Transactions on Management Information Systems (TMIS), 6(4):1- 19, 2015.
[73] Ronald J Williams. Simple statistical gralicnt- algorithms for connectionist reinforce- ment learning. Machine learning, 8(3):229-22) 256,, 1992.
13

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] We have described the limitations in the future work.
(c) Did you discuss any potential negative societal impacts your work? [Yes] Our work does have any potential negative societal impacts to our knowledge.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full of assumptions all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)? [Yes] The datasets public. All the codes will be open-sourced.
(b) Did you specify all the training details (c.g.. data splits, hyperparameters, how they were chosen)? [Yes] We provide all the details in the appendix.
(c) Did you report error bars (e.g.. with respect to the random seed after running experi- ments multiple times)? [Yes] We have run all the experiments five times. The mean and the standard deviation are reported.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] We have described the hardware in the appendix.
4. If you using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [No] But it is mentioned in the link we
provided.
(c) Did you include any new assets either the supplemental material or as URL? [No]
(d) Did you discuss whether and how consent was obtained from people whose data you'l 're using/curating? [Yes]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information offensive content? [Yes]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "r3k6hKa0HyP", "reviewer": "Reviewer_a7Cx", "summary": "This paper studies the problem of placing embedding tables of a recommender system model in a distributed training environment. How the embedding tables are placed can have significant impact on latency. Embedding lookup consists of 4 stages - In the forward pass, the embeddings are calculated and communicated of dense vector to target devices, while in the backward pass, the gradients are communicated to the device with the embedding table and the gradients are then calculated for the embedding table. The tables can easily lead to imbalances if not placed correctly and thus longer latencies during training. Thus, given a set of embedding tables and devices, identifying the best partitioning strategy is important to minimize costs. However, device placement is a NP-hard combinatorial problem. Further,  operation fusion and generalizability requirement further complicate the problem.\n\nThis paper proposes DreamShard, a Reinforcement Learning based approach for embedding table placement. DreamShard solves the problem of operation fusion and generalizability using two key ideas - (1) Learning a cost network to directly predict the cost of fused networks (2) Training a policy network that interacts with an estimated Markov decision process (MDP) without real GPU execution. ", "questions": "Questions -\n\n1. Can a table be divided and be placed on > 1 devices in this method? Are there scenarios where a table does not fit in a single device?\n2. Can the authors also discuss results comparing a simple greedy based approach? \n3. Is the shared feature extraction MLP and sum reduction in the cost network and policy network sharing weights?\n4. Given the placement of the tables is decided by the RL model and can vary across samples during training, how is operator fusion carried out? Is there a compiler pass after placement to determine what operations on a device can be fused?\n", "limitations": "Have the authors adequately addressed the limitations and potential negative societal impact of their work? Yes", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths \n\n1. The results of the paper are competitive with recent state of the art methods and evaluated on strong baselines\n2. The idea to capture cost of communication and operator fusion using a learned network to generalize to unseen tables, different number of devices etc can help provide a practical solution to this problem \n\n\nWeakness -\nNone that I can think of.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["Responsible Research Practice (e.g., IRB, documentation, research ethics)"]}, "extra_text": {}}, {"review_id": "i8xvomYq742", "reviewer": "Reviewer_AKvA", "summary": "This paper focuses on the embedding table placement problem for distributed recommender systems, where the the operation fusion and the generalizability is important and challenging. For solving the problem, the paper formulates the table placement process as an MDP, and proposes an RL framework named as DreamShard which consists of a cost network and a policy network. Extensive experiments show that DreamShard performs well.\n", "questions": "1.  Does the MDP formulation is usual in the related research? ", "limitations": "I have  no other specific suggestions about the limitations.", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The paper is well-written and easy to follow.\n2. The embedding table placement problem is practical.\n\nWeaknesses:\n1. The introduction of the related work is not enough. I am not sure about the novelty of this paper.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Nww6WIzpPUT", "reviewer": "Reviewer_ypFj", "summary": "This paper presents DreamShard for embedding table placement in recommender systems to balance the computation and communication costs. The authors formulate the table placement process as an MDP and train a cost network to estimate its states and rewards. Then they update the policy network by interacting with the estimated MDP. The results show speedup over the existing algorithms and good generalizability.\n", "questions": "Pleae see comments above. ", "limitations": "The authors have discussed the limitations and potential negative societal impact of their work.  ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths:\n1. The paper is clearly written and very structural. The idea of formulating the table placement process as a MDP is interesting.\n2. This paper considers the different speedups across different table combinations and proposes a generalizable embedding table placement problem.\n3. This paper addresses the embedding table placement problem with two networks, achieving good performance in speedup and generalization. The experiments are well-designed and thorough.\n\n### Weaknesses:\n1. It would be nice to have more results on DLRM with 8GPUs in table 1.\n2. Why the author choose 21 table features such as hash size? Can it influence the performance of the model?\n3. Lack of comparison with related works, e.g., RecShard.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "6jf1tCdm-qP", "reviewer": "Reviewer_Lt2a", "summary": "The paper proposes a RL-based approach for embedding table placement on multiple devices (mostly GPUs), which outperforms human expert and RNN-based baselines.", "questions": "* Could the method scale to larger settings? Currently 4 gpus are being used.\n* In a real world setting, the vocab distribution constantly changes (e.g. new popular items) which may affect the load balance. Could the method adapt to this dynamic setting?\n \n \n", "limitations": "no.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n* The problem is important and rarely exploreded before\n* The proposed method seems reasonable to me, with lots of reasonable heuristic baselines\n* The writing is clear and the paper is easy to follow ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept