DreamShard: Generalizable Embedding Table Placement for Recommender Systems
Anonymous Author(s)
AffiliationAddress email
Abstract
We study embedding table placement for distributed recommender systems, whichaims to partition and place the tables on multiple hardware devices (e.g., GPUs)to balance the computation and communication costs.
Although prior work hasexplored learning-based approaches for the device placement of computationalgraphs, embedding table placement remains to be a challenging problem because of1) the operation fusion of embedding tables, and 2) the generalizability requirementon unseen placement tasks with different numbers of tables and/or devices.
Tothis end, we present DreamShard, a reinforcement learning (RL) approach forembedding table placement.
DreamShard achieves the reasoning of operationfusion and generalizability with 1) a cost network to directly predict the costsof the fused operation, and 2) a policy network that is efficiently trained on anestimated Markov decision process (MDP) without real GPU execution, wherethe states and the rewards are estimated with the cost network.
Equipped withsum and max representation reductions, the two networks can directly generalizeto any unseen tasks with different numbers of tables and/or devices without finetuning.
Extensive experiments show that DreamShard substantially outperformsthe existing human expert and RNN-based strategies with up to 19% speedup overthe strongest baseline on large-scale synthetic tables and our production tables.
The code is at https://anonymous.4open.science/r/dreamshard-EBE5/ .
Introduction
Embedding learning is a commonly used technique to deal with categorical features in deep recommendation models by mapping sparse features into dense vectors [1, 2, 3, 4, 5].
However, theembedding tables can be extremely large due to the large feature sizes [6].
For example, in the
YouTube recommendation model, a single categorical feature contains tens of millions of video
IDs [7]; the Meta recommendation model demands multi-terabyte memory [8].
Distributed traininghas been adopted to place the tables on multiple hardware devices such as GPUs [3, 6, 9, 10, 11].
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 recommendation model [8].
In our internal production model, which has hundreds of tables, embedding lookupaccounts for 48% and 65% of the total computation and communication costs, respectively.
How the embedding tables are placed can significantly impact the costs.
Figure 1 shows the tracesof different placement strategies on a task of placing 50 tables on 4 devices.
Typically, embeddinglookup consists of four stages.
In the forward pass, the sparse indices are mapped into dense vectors(forward computation), which are then sent to the target devices (forward communication).
In thebackward pass, the gradients of the embedding vectors are sent back from the target devices (backward
[Figure]
v i [Figure] e
[Figure]
communication) and applied to the embedding vectors (backward computation).
The tables will easily lead to imbalances if not carefully partitioned.
The random placement in Figure 1a is bottleneckedby GPU2 with a 56.6 milliseconds latency, whilethe more balanced placements in Figure 1b and 1csignificantly reduce the costs to 42.8 and 35.95 milliseconds, respectively.
This work asks: given a setof embedding tables, how can we identify the bestplacement of the tables to balance the costs?
Device placement is essentially a partition problem,which is one of the classical NP-hard combinatorialoptimization problems [12].
A recent line of research uses reinforcement learning (RL) for deviceplacement of computational graphs [13, 14, 15, 16,17, 18, 19, 20].
For example, [13] proposed to trainan RNN controller with content-based attention topredict the placement.
Other studies advanced [13]in different ways, such as using hierarchical models [14], more sophisticated RL algorithms [15],and graph neural networks [16].
However, embedding table placement remains to bean open and challenging problem due to the operation fusion [21] of tables and the generalizabilityrequirement.
Modern embedding implementations (e.g., FBGEMM [22]), use a single operationto subsume multiple tables for acceleration.
Thespeedup of the fused operation over the sum ofthe single-table operation costs is not constant anddepends on the characteristics of the fused tables(e.g.
, table dimensions).
Our analysis finds that thespeedups vary significantly across different table combinations, ranging from 1X to 3X (Figure 12 in
Appendix A.3.2).
Thus, we not only need to reason about cost balance but also how the tables shouldbe fused to maximize the speedup.
In real-world scenarios, the adopted embedding tables and theavailable devices can change frequently (e.g., machine learning engineers/researchers may conductconcurrent experiments with various table combinations and numbers of devices).
Thus, a practicalalgorithm should generalize to tasks with unseen tables, different numbers of tables, and differentnumbers of devices.
It is non-trivial to achieve this with the existing device placement approaches.
To this end, we introduce DreamShard, an RL approach for embedding table placement.
DreamShardachieves the reasoning of operation fusion and generalizability with two novel ideas.
1) It learns acost network to directly predict the costs of the fused operations.
Specifically, the network takesas input the table features (e.g., table dimension) of each single-table and outputs the computationand communication costs.
It trains a policy network by interacting with an estimated Markovdecision process (MDP) without real GPU execution, where the states and the rewards are estimatedby the predictions of the cost network.
Equipped with sum reductions for the table representationsand max reductions for the device representations, the two networks can directly generalize to unseenplacement tasks with different numbers of tables and/or devices without fine-tuning.
Extensive experiments show that DreamShard outperforms the existing human expert and RNNbased [13] strategies on open-sourced synthetic tables [23] and our production tables, achieving upto 19% speedup over the strongest baseline.
Moreover, it can generalize to unseen tasks that havedifferent 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.
Generalizable Embedding Table Placement Problem
The embedding table placement problem seeks a device placement 1 of all the tables such that the overall cost (in terms of execution time) is minimized (we provide a background for the distributed trainingof recommendation models in Appendix A.1).
Consider M embedding tables { e 1 , e 2 , ..., e M } and Ddevices, where e i ∈ R N 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 a =
[ a 1 , a 2 , ..., a M ] ,where a i ∈ { 1 , 2 , ..., D } , assigns each table to a device.
Let c ( a ) denote the cost measured on
GPUs.
The goal of embedding table placement is to find the a such that c ( a ) is minimized.
Dueto the NP-hardness of the partition problem [12], identifying the exact solution demands extensivecomputational overhead.
Thus, the state-of-the-art algorithms often approximate the optimal partition via sampling with RL [24, 13].
However, sampling remains expensive because obtaining c ( a )requires running operations on GPUs.
Given that the embedding tables and the available devices canfrequently change, we wish to approximate the best a without GPU execution.
Motivated by this, we study the generalizable embedding table placement (GEPT) problem.
Let
E be the space of all the embedding tables.
A placement task can be denoted as T i =
( E i , D i ) ,where E i ⊆
E is a set of tables, and D i is the number of devices.
Given N train training tasks
T train =
{ T 1 , T 2 , ..., T N train } , and N test testing tasks T test =
{ T 1 , T 2 , ..., T N test } , the goal is to train aplacement policy based on T train (GPU execution is allowed during training) such that the learnedpolicy can minimize the costs for the tasks in T test without GPU execution.
[Figure]
DreamShard Framework
[Figure]
We present DreamShard framework
An overview of the framework is shown in Figure 2.
The key idea is to formulate the table placementprocess as an MDP (Section 3.1) and train a cost network to estimate its states and rewards (Section 3.2).
A policy network with a tailored generalizable network architecture is trained by efficiently interacting with the estimated MDP (Section 3.3).
Thetwo networks are updated iteratively to improve thestate/reward estimation and the placement policy.
3.1 MDP Formulation
Given embedding tables { e 1 , e 2 , ..., e M } and D devices, we aim to generate a placement a =[ a 1 , a 2 , ..., a M ] .
The key idea is to place the tables one by one at each step, where the state characterizes the tables that have been placed so far, the action is the device ID, and the reward represents theexecution time on GPUs.
Specifically, at a step t , the state s t = { s t,d } Dd =1 is all the table featuresof the tables placed on all the devices, where s t,d = { e i | i ∈ P d } denotes all the table featurescorresponding to device d
( P d is the set of table IDs that have been placed on device d ).
We furtheraugment the raw features with cost features which are obtained by collecting the operation computation and communication times from GPUs (Appendix A.3 provides a comprehensive analysis of thecost features).
Formally, the augmented state is defined as (cid:101) s t = { s t , { q t,d } Dd =1 } , where q t,d ∈ R 3 hasthree elements representing forward computation time, backward computation time, and backwardcommunication time for the current operation in device d (we provide detailed explanations of whyforward communication time is excluded in Appendix A.4).
We find that the augmented cost featurescan significantly boost the performance, evidenced by the ablations in Table 3.
The action a t ∈
A tis an integer specifying the device ID, where A t is the set of legal actions at step t .
A device ID is
[Figure]
considered legal if placing the current table on the corresponding device does not cause a memoryexplosion.
The reward r t is 0 for all the intermediate steps, and the reward at the final step M is thenegative of the cost, i.e., r M = − c
( a ) , which encourages the agent to achieve lower cost.
The procedure is illustrated for an example task of placing 8 tables { e 1 , e 2 , ..., e 8 } on 2 devicesin Figure 3.
At step 0 , no table has been placed so s
= {{} , {}} and the augmented state(cid:101)
s
= { s 0 , { q 0 , 0 , q 0 , 1 }} , where both q 0 , 0 and q 0 , 1 are zero vectors (i.e., [0 , 0 , 0] ) since all thecomputation and communication times are 0 as well.
Then the action a 0 = 1 makes the MDP transitto the next state s 1 = {{ e 1 } , {}} with its corresponding augmented state (cid:101) s
1 = { s 1 , { q 1 , 0 , q 1 , 1 }} ,where q 1 , 0 becomes a non-zero vector containing the computation and communication costs byrunning { e 1 } and {} on GPUs.
We repeat the above process, and finally at step 8 , we haves 8 = {{ e 1 , e 4 , e 7 , e 8 } , { e 2 , e 3 , e 5 , e 6 }} .
The corresponding q 8 , 1 , and q 8 , 2 are the measured times ofrunning { e 1 , e 4 , e 7 , e 8 } and { e 2 , e 3 , e 5 , e 6 } on two devices.
The action sequence a
=
[ a 1 , a 2 , ..., a 8 ]is the generated placement, which is then evaluated on GPUs to obtain the reward.
Discussion 1.
The MDP enjoys two desirable properties.
The legal action A t can guarantee thatthe generated placement satisfies the memory constraints.
The one-by-one placement enables theagent to be generalized across different numbers of tables.
For example, an agent trained on an MDPwith very few tables can be applied to another MDP with more tables by simply executing more steps.
Discussion 2.
A straightforward idea to solve the MDP is to greedily place the current table on thedevice with the lowest cost at each step, where the cost function can be one of or a combination ofthe 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 a learning-based algorithm to explore variousplacement possibilities and make comprehensive decisions based on all the state features.
3.2 Learning an Estimated MDP
Interacting with the above MDP is computationally expensive since obtaining the cost features andthe 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 f cost denote the costnetwork.
f cost takes as input the raw table features s t , and predicts cost features { q t,d } D d =1 and theoverall cost c ( a ) .
f cost is trained with mean squared error (MSE) loss using the cost data collectedfrom the GPUs.
Once trained, it can predict the cost features or the reward with a single forward passwithout GPU execution.
However, it is non-trivial to design the architecture of f cost because it needsto accommodate different numbers of devices (i.e., s t can have variable sizes), and different numbersof tables in each device (i.e., s t,d can have variable lengths).
The left-hand side of Figure 4 shows DreamShard’s generalizable design of f cost , which is based ontwo key ideas.
First, it uses a shared MLP to map raw table features into table representations.
Forany unseen tables, this MLP can be directly applied to extract table representations.
Second, it enablesa fixed-dimension representation for each device with sum reductions (i.e., the element-wise sum ofthe table representations in the device), and similarly for the overall representation across deviceswith max reductions (Appendix B.3 compares different reduction choices and finds that this sum-maxcombination leads to the most accurate prediction).
The reduced representations are then followed bymultiple MLP heads for cost predictions.
For unseen tasks with different numbers of tables and/ordevices, the reductions will always lead to fixed-dimension device/overall representations, so that theprediction heads can be directly applied.
Appendix B.1 provides more details.
[Figure]
F o r w a r d H ead
B a ck w a r d
H ead
C o mm .
H ead
[Figure]
P r ed i c t ed
B a ck w a r d C o s t
P r ed i c t ed
C o mm .
C o s t
P r ed i c t ed
F o r w a r d C o s t
[Figure]
1: Input: Training tasks T train , the number of data collection steps N collect , the number of cost network update steps N cost , the batch size for updating the cost network N batch , the number RL update steps N RL , the number of considered episodes in each RL update N episode 2: Initialize a cost network, a policy network, and a buffer 3: for iteration = 1, 2, ... until convergence do 4: for step = 1, 2, ...
N collect do ▷ Collect cost data from GPUs 5:
Randomly sample a training task from T train 6:
Generate a 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 step = 1, 2, ...
N cost do ▷ Update the cost network (no GPU execution) 10:
Randomly sample N batch cost data from the buffer 11:
Update the cost network based on MSE loss (Eq. 1 in Appendix B.4.1) 12: end for 13: for step = 1, 2, ... N RL do ▷ the
Randomly sample a training task from T train 15:
Collect N episode 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.4.1) 17: end for 18: end for
3.3 Training the Policy Network on the Estimated MDP
Generalizable policy network architecture.
Let π be the policy network.
π maps the augmentedstate (cid:101) s t
= { s t , { q t,d } Dd =1 } to action a t , i.e., a t = π ( (cid:101) s t ) .
π also adopts a generalizable design, shownin the right-hand side of Figure 4.
Like f cost , π uses a shared MLP and sum reductions to producea fixed-dimension representation, which is then concatenated with the cost features to obtain thedevice representation.
To accommodate the potentially variable action space (i.e., the number ofavailable devices may vary), a shared MLP will process each device representation separately toobtain a confidence score, followed by a Softmax layer to produce action probabilities.
This designallows π to generalize across different numbers of devices.
Appendix B.2 provides more details.
Training and inference.
Algorithm 1 summarizes the training procedure of DreamShard, whichiteratively executes the following: 1) collect cost data from GPUs based on the placements generatedby 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 trainingprocess, the estimated MDP gradually becomes more accurate, and the resultant policy network tendsto generate better placements.
Appendix B.4.2 provides more details of the training procedure.
Forthe inference, the trained cost network and policy network can be directly applied to unseen tasks togenerate placements without GPU execution, which is summarized by Algorithm 2 in Appendix B.4.3.
Experiments
Our experiments aim to answer the following research questions.
RQ1: How does DreamShardcompare with the existing human expert and RL-based placement strategies?
RQ2: Can DreamShardgeneralize to placement tasks with different numbers of tables and/or devices?
RQ3: How efficientis 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 to enable a meaningful evaluationbecause the cost will always be very small no matter how the tables are placed.
Thus, we usetwo industrial-scale datasets.
DLRM 2 is a large-scale synthetic dataset with 856 tables, recentlyreleased by Meta.
It shares memory access reuse patterns similar to those arising in Meta productionworkloads.
Prod is an internal large-scale dataset for production recommendation models.
It has asimilar scale as DLRM.
The main difference is that DLRM only has a fixed dimension for all thetables, while Prod is more challenging with diverse table dimensions, ranging from 4 to 768.
Forreproducibility, we mainly focus on the DLRM dataset since it is open-sourced.
We only report themain results on the Prod dataset for verification purposes.
We provide more details in Appendix C.
Baselines.
We compare DreamShard against human expert strategies from previous work [27, 8, 28],including size-based , dim-based , lookup-based , size-lookup-based greedy balancing strategies.
We also include an RNN-based RL algorithm [13], which uses RNN architecture to map operatorsto devices.
Since the feature extraction layers of the RNN-based method were designed for otheroperations instead of embedding tables, for a fair comparison, we adapt [13] by making it have thesame feature extraction layers as in DreamShard.
We provide more details in Appendix D.
Configurations.
To evaluate the generalizability of DreamShard, we randomly divide the tables intoa training pool E train and a testing pool E test .
The two pools have the same number of tables but theyare not overlapped.
A sharding task T i is constructed by randomly sampling a subset of |E i | tablesfrom a pool, where the number of tables |E i | ∈ { 10 , 20 , 30 , 40 , 50 , 60 , 70 , 80 , 90 , 100 } for the DLRMdataset, and |E i | ∈ { 20 , 40 , 80 } for the Prod dataset.
For all the experiments, we randomly sample 50training and 50 testing tasks from E train and E test , respectively.
DreamShard is trained on the trainingtasks and will be evaluated on unseen tables in the testing tasks.
We denote placement tasks withdifferent numbers of tables and devices using the format of dataset - num_tables ( num_devices ).
For example, DLRM-30 (4) suggests that there are 30 tables sampled from the DLRM dataset in eachtraining/testing task with 4 available devices.
We provide more details in Appendix E.
Implementation Details.
We use the same hyperparameters for all the experiments with N collect = 10 ,
N cost = 300 , N batch = 64 , N RL = 10 , N episode = 10 , and 10 training iterations. We use an entropyterm in the policy gradient loss with a weight of 0 .
001 for exploration.
GeForce RTX 2080 Ti GPUsand V100 GPUs are used for the DLRM and Prod datasets, respectively.
All the experiments are run5 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 quantitativecomparisons of DreamShard against the baselines.
Qualitatively, Figure 1 visualizes the tracesof DreamShard and the baselines on one of the tasks from DLRM-50 (4).
DreamShard achievessignificant better overall cost than the best baseline (35.95 vs. 42.8) with 1) a better balance of forwardand backward commutation workloads, and 2) less communication time, possibly due to a betterbalance of table dimensions.
Quantitatively, Table 1 presents comprehensive evaluations on taskswith 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 generalizabilityon unseen tables, achieving the same level of performance on all the testing and training tasks.
DreamShard appears to be more advantageous on harder tasks.
Specifically, DreamShard achieves
[Table] [Table]
more improvement over the baselines on tasks with more tables/devices on the Prod dataset.
Inparticular, DreamShard achieves 19% improvment over the strongest baseline on Prod-80 (8). 4)
RNN-based method is only better than the random strategy on tasks with few tables/devices, but isworse on harder tasks.
A possible reason is that RNN-based algorithm does not have a cost network,and using RL alone could lead to unstable performance.
5) Lookup-based strategy is the best baselineon the DLRM dataset, while dim-based strategy is better on the Prod dataset.
A potential reason isthat the tables in the Prob dataset have very diverse table dimensions, while the tables in the DLRMdataset have the same dimension.
As such, the dimensions on Prod tasks can more easily becomeimbalanced, leading to poor communication efficiency.
Dim-based strategy can better balance thedimensions, which leads to a better overall performance.
DreamShard outperforms the baselines onboth 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 trainedfrom one task to another task without any fine-tuning (the rightmost column), where the source andthe target tasks have different numbers of tables and/or devices.
DreamShard shows neglectableperformance drop, suggesting that it is generalizable across different numbers of tables and/or devices.
[Figure]
[Figure] [Figure] [Table] [Figure]
[Table] [Figure] [Table] [Figure]
[Table] [Figure] [Table] [Figure]
[Table] [Figure] [Table] [Figure]
[Table]
[Figure]
[Table]
[Figure] [Table] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
Analysis of training efficiency (RQ3).
Figure 5 plots the performance of DreamShard w.r.t.
thenumber of iterations and running time in seconds on four 1080Ti GPUs.
The training of DreamShardis highly efficient.
On DLRM-50 (4), it can achieve strong performance in less than 5 iterations or200 seconds.
Note that we only need to re-train or fine-tune DreamShard when the table pools havesignificant changes.
Once trained, it only needs a forward pass for inference.
Hyperparameter study (RQ4).
We study the impacts of two key hyperparameters: 1) N RL , whichcontrols the RL update frequency, and 2) N cost , which determines the cost network update frequency.
We vary one of them with the other fixed, shown in Figure 6.
Increasing N RL or N cost will bothlead to improvement, suggesting that both the cost network and the policy network need to besufficiently trained.
However, when N RL and N cost are large enough, increasing them will not bringmore improvement.
Considering that larger values will lead to more computational costs, we set
N RL = 10 or N cost =
300 as a trade-off between the performance and the training efficiency.
Ablation study (RQ5).
We study the importance of each table feature, compare DreamShard witha greedy baseline, and check whether the RNN architecture helps, with the following ablations. 1)
We remove each of the features in the state. 2) Instead of using RL to make decisions, we simplyplace the table on the device that has the lowest cost at each step, where the cost is estimated bythe sum of all the cost features predicted by a trained cost network.
We add an RNN upon thedevice 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 proposedaugmented state.
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 computationand communication workloads.
3) Using more features leads to consistently good performance.
While the policy makes decisions sequentially, RNN does not provide clear benefits.
This is why wehave 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 trainan accurate cost network, and how accurate the cost network needs to be to enable a strong policy.
Specifically, we randomly sample 10,000 cost data points from the DLRM-50 (4) dataset.
Thenwe use 20% for testing, and vary the size of the training data to train the cost network.
Further,we fully train a policy network with 100 iterations based on each of the trained cost networks.
Wemake two observations from the results in Figure 7. 1) As expected, more data points lead to amore accurate cost network.
Interestingly, after around 100 data points, the policy network doesnot keep improving even though the cost network becomes more accurate.
Thus, we only need asufficiently (but not perfectly) accurate cost estimation to achieve the best performance.
This alsopartially explains why DreamShard can generalize: even though the cost network could be not veryaccurate 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 costfeatures and rewards directly from GPUs, shown in Figure 8.
Using the estimated MDP can make thetraining 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.
[Figure] [Figure]
Using cost network to estimate the leads orders faster training and
Related Work
Embedding tables.
Embedding tables are commonly used to deal with sparse features in recommendation models [1, 2, 3, 4, 5, 29, 30, 31].
However, the extremely large embedding tables areoften the storage and efficiency bottlenecks [6, 7, 8, 3, 6, 9, 10, 11, 32].
To our knowledge, the onlywork that targets the embedding table placement problem is RecShard [27], which approaches tableplacement from the system view by exploiting the tiered memory hierarchy.
RecShard identifies thehot rows based on the unequal access patterns and places them in GPU memory with the others in
CPU memory.
However, they still use the size-lookup-based strategy to estimate the costs, which isinaccurate.
In contrast, our work develops a neural cost network for cost estimation and an RL-basedoptimization algorithm, which are orthogonal to their efforts in the tiered memory hierarchy.
Anotherline of work focuses on reducing the embedding table sizes [33, 34, 35, 36, 37, 38, 39], which isorthogonal to our work since DreamShard is also applicable to compressed tables.
Device placement optimization.
The existing device placement techniques can be mainly groupedinto two categories: RL-based algorithms, and cost modeling methods. 1) RL-based algorithms treatthe device placement as a black box and optimize the cost objective in a trial-and-error fashion [13, 14,15 , 16, 17, 18, 19, 20].
Unfortunately, these methods are computationally expensive and often requiretraining from scratch on unseen tasks.
While Placeto [16] shows generalizabily on computationalgraphs with graph embeddings, it can not deal with various table combinations since there is no graphstructure, and it can not handle different numbers of devices.
2) Cost modeling methods build a costmodel to reflect the real performance and adopt offline algorithms (e.g., scheduling, and dynamicprogramming) to optimize the placement [40, 41, 42, 43, 44].
However, the cost model could beinaccurate.
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.
Deep RL.
Deep RL has recently made significant progress in games [45, 46, 47, 48, 49, 50].
Ourwork is related to using RL to optimize machine learning model designs, such as neural architecturesearch [51], data augmentation [52], and pipeline search [53].
However, these methods often onlyfocus on one task and can not generalize to unseen tasks.
Our work is also related to solvingcombinational optimization problems with RL [24, 54].
Unlike the above studies, we show that RLcan 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 formulatethe problem as an MDP which places the tables one by one at each step.
Then we leverage RL tosolve the MDP.
To accelerate the training and the inference, we build an estimated MDP by training acost network to approximate the state features (i.e., computation and communication times) and thereward (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 demonstratethe superiority of DreamShard over the existing algorithms.
Moreover, DreamShard shows stronggeneralizability, 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.
Checklist
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 limitationsin the future work.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] Our workdoes 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]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A]
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 datasetsare public.
All the codes will be open-sourced.
(b) Did you specify all the training details (e.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 experiments multiple times)?
[Yes] We have run all the experiments five times.
The meanand the standard deviation are reported.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] We have described the hardware
in the appendix.
If you are 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 in the supplemental material or as a URL?
[No]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[Yes]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[Yes]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
[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 to participants and the total amount
spent on participant compensation?
[N/A]