Learning Efficient Parameter Server Synchronization Policies for Distributed SGD
Authors: Rong Zhu
Sheng Yang
Andreas Pfadler
Zhengping Qian
Jingren Zhou
Venue:  (2020)
Keywords: Distributed SGD
Paramter-Server
Synchronization Policy
Reinforcement Learning

Abstract: We apply a reinforcement learning (RL) based approach to learning optimal synchronization policies used for Parameter Server-based distributed training of machine learning models with Stochastic Gradient Descent (SGD). Utilizing a formal synchronization policy description in the PS-setting, we are able to derive a suitable and compact description of states and actions, allowing us to efficiently use the standard off-the-shelf deep Q-learning algorithm. As a result, we are able to learn synchronization policies which generalize to different cluster environments, different training datasets and small model variations and (most importantly) lead to considerable decreases in training time when compared to standard policies such as bulk synchronous parallel (BSP), asynchronous parallel (ASP), or stale synchronous parallel (SSP). To support our claims we present extensive numerical results obtained from experiments performed in simulated cluster environments. In our experiments training time is reduced by 44 on average and learned policies generalize to multiple unseen circumstances.

ICLR 2020
LEARNING EFFICIENT PARAMETER SERVER SYNCHRONIZATION POLICIES FOR DISTRIBUTED SGD
Rong Zhu", Sheng Yang, Andreas Pfadler, Zhengping Qian, Jingren Zhou Alibaba Group
Corresponding Author
ABSTRACT
We apply a reinforcement learning (RL) based approach to learning optimal syn- chronization policies used for Parameter Server-based distributed training of ma- chine learning models with Stochastic Gradient Descent (SGD). Utilizing a formal synchronization policy description in the PS-setting, we able to derive suit- able and compact description of states and actions, allowing us to efficiently use the standard off-the-shelf deep Q-learning algorithm. As result, are able learn synchronization policies which generalize to different cluster environments, different training datasets and small model variations and (most importantly) lead to considerable decreases in training time when compared to standard policies such as bulk synchronous parallel (BSP), asynchronous parallel (ASP), or stale synchronous parallel (SSP). To support our claims we present extensive numerical results obtained from experiments performed in simulated cluster environments. In our experiments training time is reduced by 44% average and learned poli- cies generalize to multiple unseen circumstances.
INTRODUCTION
In recent years, Stochastic gradient descent (SGD) Bottou (012)) and its variants Kingma & Ba 2014):[Chen al, (2016), have been adopted the main work horse for training machine learning (ML) models. To be able to train large models, which are both computationally demanding require very large training datasets, SGD is often parallelized across several machines, with the well-known parameter-server (PS) framework being the PS setting, there commonly exist one (or several) parameter servers and multiple worker nodes. The parameter server maintains the globally shared model parameters and aggregate updates from workers. Each worker node pulls the latest model parameters from the server, computes all gradients and pushes them back for updating. As this approach generally reduces the amount of inter-node communication. it may provide for considerably reduced training time.
of the most widely adopted distribution strategies. In
Challenges of Distributed SGD in PS. In the PS setting, central task design synchroniza- tion policy, which coordinates the execution progress of all workers. This synchronization policy determines in each step, i.e. whenever a gradient is pushed by some worker, the state ("run" or "wait"" of each worker, until the next update arrives the parameter server. Thus, it directly deter- mines the overall training time. However, finding a good synchronization policy is difficult, as this will least depend on the properties of the underlying optimization problem and the nature of the cluster used for training.
We briefly review some existing policies and discuss their limitations: In the simple bulk syn- chronous parallel (BSP) policy Valiant (1990), the parameter server waits for all workers to push their updated gradients, and then lets them pull the same latest model parameters for the next step. However, in heterogeneous cluster, it is common that some workers, also referred to stragglers, run much slower than others. Waiting for these straggler workers certainly decreases the number of SGD iterations per unit of time and leads less optimal usage of available computational resources.
The asynchronous parallel (ASP) policy alleviates the straggling problem by allowing each worker to run immediately after it pushes its gradients. Therefore, fast workers can move ahead without waiting for others. However, worker may push stale gradients that are evaluated on an older version of the model parameters, which may have negative impact the overall convergence speed Chen phen jet al.] (2016); Cui et al. (2016). Thus, neither BSP nor ASP consistently outperform the other on different models and datasets Zinkevich et al. (2009): Dutta et al. 2018).

---

ICLR 2020
A better policy is, for instance, given by the stale synchronous parallel (SSP) policy Ho et al. 2013 Dai et al. (2015), which can be considered as a trade off between BSP and AP. For SSP one defines fixed threshold s, such that the fastest worker is allowed to outpace the slowest one by at most steps ("bounded staleness"). However, an optimal value of & is hard to specify without enough prior knowledge. Moreover, the optimal value of may change during the entirety of the training process. Other synchronization policies (see in Section 2 may either degenerate to the above mentioned three policies may need extensive manual tuning of hyper-pppe ppper-parameters.. Now, the key problem addressed in this work is the following: how can we design a synchronization policy to automatically and adaptively optimize the overall training time of SGD in PS? Note that the following, whenever we speak of "overall training time" or "time until convergence", we specifically mean the amount of time needed until model reaches a certain pre-defined aceuracy on a validation dataset.
Our Key Contributions. To tackle this problem, we focus on designing automatic synchronization policies for distributed SGD in a PS setting. Instead of relying on expert experience we try learn a better synchronization policy using a reinforcement learning (RL) approach with training data being obtained by observing the execution process of PS-based distributed SGD.
To this end, we first represent the synchronization problem in the distributed SGD training in unified framework, allowing us to formally describe both existing (BSP, ASP and SSP) and learned policies. Based on this framework, we formulate the synchronization policy design problem as an RL problem, where we reward those policies requiring less training time. To train an RL-based synchronization policy (RLP), we carefully design the state and action space of the RL model, such that it is able to generalize to different training data, models and cluster environments, while still ensuring efficient policy learning process.
Our model is trained using the off-the-shelf deep Q-Learning algorithm Mnih al. (2013) Van Has- selt et al. (0166). Furthermore, we design a pre-training process to speed up the convergence. Em pirical results demonstrate the validity of our approach and the advantages of our learned RLP policy terms of training efficiency and generalization ability. In our experiments, RLP improves overall training time by 44% on average in comparison to the best existing policy. More over, RLP is able to generalize to multiple unseen circumstances.
Limitations. We note the following limitations of this work: First, our experiments are based on "plain" SGD in its simplest form, e.g. without momentum or other adjustments. While our frame- work may in principle be used with any first-order optimization scheme, we have chosen to refrain from using more commonly used methods, such as Adam Kingma & Ba (201)) in order to reduce the number of hhper-parametter parameters and allow for clearer differentiation between different policies.
Second, we note that all of our experiments were performed in a simulated cluster environment, allowing us to easily change the number of workers used and artificially create training instances characterized by different levels of straggling workers, resp. gradient staleness.
Third, due to resource constraints, we have not applied our approach to the training of very demand- ing models, such as deep Convolutional Neural Networks Krizhevsky et al. (2012) or transformer- based models like BERT Devlin et al.) (2018). which are now considered the "state of the art" in Computer Vision and NLP, respectively. We leave experiments on this scale for future work.
Despite the above limitations, we argue that results clearly indicates the applicability of our approach in practice. In our experiments, we were not only able to demonstrate the mere existence of optimal synchronization policies for an individual pair of underlying model and training data, but also show that learned policies are able to provide significant speedups, even when applied to training slightly different models on different training datasets and different cluster environments.
2 RELATED WORK
PS-based training with distributed SGD can be considered a standard method for training large ML models. Finding proper synchronization policies coordinating all workers to reduce training time has been a long-standing problem. For example, the classic BSP policy Valiant (1990) has been di- rectly applied for distributed SGD in PS PS-setting. As the performance of BSP is heavily diminished by the straggling problem, other variants, such as the K-syne BSP and K K-batch-syne BSP have been proposed in Dutta al. (2018). They alleviate the straggling problem by slightly relaxing the
2

---

ICLR 2020
synchronization condition. However, the global synchronization barrier still exists, workers tend to spend a considerable amount of time in an idle state. Other ways to speed up the BSP policy include reassigning data Harlap et al. (2016) and adding backup workers Chen et al. (2016). How- ever, they may case additional, non-negligible communication overhead Jiang et al. (211)) Another extreme is to run distributed SGD using the ASP policy, which avoids the stragg inn problem but typically forces one to consider the staleness problem. The negative effects of staleness have been shown in Zinkevich et al. (000))(Chenn al. (2016); Cui al. (2016); Duttta al. (2018): Dai et al. (2019) on typical ML models, where Dai et al. (2019) performed comprehensive experimental analysis. To correct bias caused by staleness, Hadjis et al. (2016): Mitliagkas al. (2016); ZZeng et al. (2117) proposed some methods to compensate for delayed gradients or tune the momentum parameter, which also increases the overall computational costs.
To combine the advantages of BSP and ASP together, Ho et al. (2013): Dai et al. (2015) proposed the SSP policy as a trade-off. SSP may often lead to faster convergence when compared to both BSP and ASP. However, the threshold in SSP which bounds the staleness is hard to tune and
fixed during the training process. To overcome this, et al. (017); Zhao et al. (2019) proposed dynamic SSP policy where is tuned during the training process. Fan et al. (2018) proposed more flexible adaptive asynchronous parallel policy to allow different values
for each worker.
Although more adaptive, they need manual tuning of hyper-parammters. As result, until now, there appears to exist no synchronization policy that is both fully adaptive and automatic.
This paper, for the best of knowledge, marks the first instance where RL is used find synchro- nization policy for distributed SGD in completely data-driven fashion. RL has been successful applied to control robotics Duan al. (2016) and games Mnih et al. (2013); Silver et al. Recently, it has been widely used to optimize problems such as task scheduling Mao et al.
(2016)
2019
resource management Mao al. (2016) and optimization Li & Malik (2016); Marcus et al. Here, we formulate the search for efficient synchronization policies as an RL problem and derive RL.-based policies (RLP) which can be applied to training a specific underlying model (and varia- tions thereof): Once such a policy has been learned, it can be reapplied to future training instances. This is particularly relevant in situations where the same model (or type of model) is trained regu- larly different training data sets on the same cluster of machines (or similar cluster in the case of e.g. dynamically allocated cloud computing resources).
2019
Overall, to the best of our knowledge, our proposed RLP provides for the first time synchronization policies which are adaptive, automatic and avoid both the straggling and staleness problems.
3 METHOD
We descrihe the technical details of our proposed method in this section. First of all, we present unified framework to represent the synchronization problem of distributed SGD in the PS-setting in Section 3.1 Based on this, we formalize the synchronization policy design problem as an RL
problem in Section 3.2
Section 3.3 discusses how to train the RL-based policy.
3.1 A UNIFIED FRAMEWORK
We introduce unified framework generalizing all existing policies (BSP, ASP and SSP) and pro- viding us convenience and flexibility to design new synchronization policies. The pseudo-code of the framework is shown in Algorithm 1.
In a PS environment, let S be the parameter server holding the global parameter w and W (W1, W2.... W be a set of workers. For each worker W,, we regard its computation process as a series of steps. In each step, the state of worker W; is either active or idle. An active worker pulls the latest parameter w from S. performs the SGD computation on a mini-batch of data and pushes the gradient Vw back to S. Then, W; is scheduled by $ to run or wait. On the server side, S iteratively receives gradients sent by workers and coordinates their execution progress until the con- verge condition is met. S maintains the global parameter w and a set I of all idle workers. Initially, we set = 0.. To avoid ambiguity. we refer to the period of S receiving two consecutive gradients as iteration. In each iteration, S receives a gradient Vw submitted by worker W, and updates the global parameter w according to the SGD update rule. After that, W, is added into the idle set I, and S selects subset of active workers A according to the synchronization policy. All workers
3

---

ICLR 2020
Algorithm Unified Synchronization Policy Framework Input: a server S and a set of workers W = Wi, W Wx
Output: trained global parameter w and total time cost
I: I < 0. gets@
2: initialize the global parameter
3: set all workers in W to be active
4: while stopping condition of SGD
met do
5: receives gradient V submitted by worker W,
6: update parameter w to w' by the SGD rule
7: I I {W,)
8: select A C according the synchronization policy
9: record time cost At and set <- + At
10:
end while
11:
return w and
in A are allowed to run while all other workers in / - A keep idle. Let At denote the iteration time cost. The global parameter w and total time Σ At is returned after convergence (i.e. when certain stopping criterion is met).
By specifying a different set A of active workers in each iteration, our framework can describe vari- ous synchronization policies. For example:
BSP: we set A if [ < l and A otherwise, so all workers execute together only when all of them finish a computation step;
ASP: we set A [, so that all workers always continue running after submitting the gradi- ents:
SSP: let p. denote the number of computation steps performed by worker W,. Given a thresh-
old for the submitted worker W., if 0 < ps min, Pj we set A {W,}: otherwise if P. -min Pj = s, we set A = 6; otherwise when Ps minj Pj, we set A = I, so all workers outpace the slowest one by more than steps.
In each iteration, there exist at most 2/W] different choices. Finding a better policy using some combinatorial search algorithms or hand-crafted heuristic rules seems intractable due to the large search size. Thus, we aim at designing method to automatically learn synchronization policy.
3.2 FORMULATION AS AN RL PROBLEM
Based on the unified framework, we can formulate the synchronization policy design problem as an RL problem. To lay foundation, we briefly review some preliminaries of RL. The general setting of RL shown in Fig. [ where an agent continuously interacts with an environment. In each step n, the agent observes some state S,, of the environment and is asked to take action O,, Following this action, the environment emits a reward r. and transitions to state S,+1. 1. The state transitions and rewards are both stochastic and satisfies the Markov property, i.e. the state transitions and rewards depend only on the state S,, and action a,, The goal of the agent is to learn a sequence of actions chosen by observing the states maximize the expected cumulative reward E.,, rn].
Notice that, the agent picks actions based on a policy n, which is a probability distribution over state-action pairs: *(S, a) - [0, 1]. In most practical problems, it is impossible store the policy in tabular form but more common to represent it as function
parametrized by 6. Recently, deep neural networks (DNN) have been widely used to represent ma many RL problems Mnih et al. (2013); Li & Malik (2066); Van Hasselt et al. (2016); Marcus et al. )2019) Mao et al. (2019). Following this trend, in our work, we also adopt a DNN to represent the RL. policy. Feeding it with state vector S, it outputs value x(S, a) for all possible actions a.
We find that the synchronization policy design problem resembles prototypical RL problem, as it aims at learning how to choose an active worker set (action) in each iteration based on the SGD execution progress (state) to optimize the total time cost (reward). Fig. illustrates how to set the three key components in the RL formulation. The details are claborated follows:
State: we choose features in each SGD iteration to characterize its execution progress. For gener- alization purposes, the state feature vector should be able to represent clusters with different number of workers. To this end, we regard the execution of all workers as a black-bo box and just encode the information they submitted in the server side. For each iteration n, we just record the feature tuple
4

---

ICLR 2020
observe
policy sots. 0)
Aent
Environment
state S
(WG)
server
W
workers
W
rwp-layer DNN
WW
take action
reward v.
Figure 1: Illustration of formulating the synchronization policy problem RL problem.
T, (n, L(w), L(w") L(w), fn), where L(w) and L(w") are the loss values before and after one iteration. En records the number of pushed gradients by other workers during the execution step of this submitted worker, which reflects the level of staleness of the submitted gradients. Notice that, Tn is irrelevant to the number of workers in the cluster. Meanwhile, we do not use the gradient in- formation in our features since highly dependent on the underlying ML models and would thus impede generalization. The state vector S, (T,, T - In+H++1) ++1) keeps track of the features of the Il most recent iterations. At the very beginning, we pad the historical information all by zero. Action: the largest action space contains at most 2 w possible actions, which is relevant the worker number and impossible to train for large clusters. In our setting, we choose a very small but still powerful action space. Let the action On E {9, {Wa], J. W}. That means we have three valid actions in each iteration: @ keeps all idle worker to be idle; {W.} allows the submitted worker itself to run; and (W) allows all idle workers to run. Setting such a small size action space enables faster training of the RL policy. As we have shown earlier in Section 3.1 this action space is enough to represent and switch between the existing policies BSP, ASP and SSP with different threshold s. Moreover, it can also represent more complex policies.
Reward: we directly set the reward signal in each iteration as rn = At. Hence, maximizing the cumulative reward corresponds to finding policies minimizing the total time cost.
3.3 TRAINING RL POLICY
It has been shown in Gu al. (2016) that the off off-policy RL algorithms such as Q-learning can be more sample efficient than their policy gradient counterparts. This is largely due to the fact that policy gradient methods require on-policy samples for the new policy obtained after each update of the policy parameters. Therefore, we adopt the standard deep Q-learning method Mnih et al.)
2013): Van Hasselt et al. (2016) to perform an end-to-end training of our RL policy. Our policy network is two-layer neural networks with 64 and 32 units in each hidden layer, respectively. Leaky rectified activation units are used the two hidden layers. We present the detailed training process in Algorithm 2.
The RL policy training process is embedded into the SGD algorithm. Each training episode corre- sponds training the underlying MI models once by SGD. To stabilize the learned policy, we apply an evaluation policy network Q and a targeted policy network Q' with parameters and 6*, respec- tively. The parameters of Q are copied to Q* every e iterations. We maintain an experience replay pool 'D with size N to store transitions (S", a, r, S),, where S" is the previous state of S. In each iter- ation, we sample mini-batch of transitions from D. For each transition (S:-1, 04-1,
S,), let y: Q( 4;-1; 0) and i = Ti-1 + max, Q(S, 0) denote the estimated and targeted cumu- lative reward of the current state with discount factor ? (0, 1). We apply a square loss 11 9: to train the parameters 0 of Q. After that, we choose the action in an e-greedy manner in order to decide the execution status of workers. Finally, the policy network Q" returned for inference.
In our synchronization policy design problem, we observe that interval of possible training times of SGD under different policies has large overlaps due to the stochastic nature of SGD. To help the policy network to distinguish the difference of synchronization policies earlier and speed up the converge, we integrate the policy network with some prior knowledge on existing synchronization policies. Specifically,
apply pre-training process to train the policy network in advance. We execute SGD with existing policies (BSP, ASP and SSP with different thresholds s) and record the information state S and action a for each iteration. After SGD finishes, we obtain the truly
5

---

ICLR 2020
Algorithm Training RL Policy
1: initialize the experience replay pool D with size N for episode <- 1 to do
2:
3: while the stopping condition of SGD is not met do vector S and reward Att
4: obtain the current store the transition (S' a, S) into the replay pool sample mini-batch B of transitions from D
6:
7:
train the parameters 0 of Q by the squared loss of ys and d. on B copy the parameters of 8 of Q" every iterations random action with probability
8:
9:
choose action
arg max. Q* (S, (*) with probabilityl apply action ar set the execution status of workers
10:
11:
S" S
end while
13:
end for
14:
return policy network Q" with parameters 0"
cumulative reward -t, i.. the training time left until the end, for each iteration. Then, we apply the squared loss the difference of (S,, a: 0) and -f over all iterations to train the parameters of the policy network. We observed that using pre-training policy exploration time may be saved and faster convergence achieved.
4 EXPERIMENTS
We implement RLP in a simulated cluster/PS environment. This allows for convenient creation of various training instances. We now report our evaluation results in this section.
Instance Generation Method. We train the RLP by training DNN model consisting of several fully connected layers on simulated cluster of 10 workers. In the following we refer to this model as the "underlying" model (as opposed the RL policy model).
For each training instance, we need to configure both the cluster and the underlying model. To simulate stragglers, we have a probability p 0.3 to activate the sleep() function in each worker in the cluster. The sleeping time obeys a Gaussian distribution. For the underlying DNN model, we randomly choose a number E {0, 1, 3. as the number hidden layers. When = O, DNN degenerates to multi-class logistic regression model. Each hidden layer contains 256 units with rectified activation functions. We use cross-entropy as the loss function. In each instance, we randomly sample 50% data from the MNIST dataset and run the standard SGD for training.
Training and Testing Methods. The hyper-parameters for RLP are set as follows: historical size H 10, replay pool size N 50, mini-bated size [B] = 32, copy rate = 5, discount factor y = 0.8, exploration probability we set its batch size to 16 and learning rate to 0.01. SGD terminates once we attain 88% validation accuracy. Before training, we apply the existing policies BSP, ASP, and SSP with 2,5, 8 to pre-train the policy network on 100 instances, respectively. Then, we train the RLP policy until convergence with about 1,000 instances (episodes). When testing we run each instance 30 times with different random seeds under the same policy and report the average time cost.
0.1 and learning rate be 0.01. For the underlying DNN model,
Performance Comparison vs. Existing Synchronization Policies. First of all, we examine the performance of RLP by comparing its execution time with respect to BSP, ASP, and SSP. The result of SSP refers the result for the best threshold s. We report the average results tested on 100 instances in Fig. 2(a) On the whole, our proposed RLP runs 1.56, 1.86 and 1.44 times faster than BSP, ASP and SSP, respectively. This verifies that our RL-based method can find better synchro- nization policies in the case of our particular choice of underlying model.
To present more details, we also compare the average testing results on different models and different cluster environments. Fig. 2(b) reports the results for the underlying model with 0 and 3 hidden layers. RLP improves the running time by 48% and 43% w.r.t the best existing policy, respectively. the speedup ratio of RLP w.r.t. different models. This
There exists significant difference most likely due to very similar loss curves for these model variations. Therefore, the relative difference of different policies tends to be similar. Fig. 2(c) reports the results clusters with

---

ICLR 2020
201
25
Redattv Rodatiee Troliing Tims
Rokative Robatte Traliiog
20
27
(RIP=1)
15
(Rl
i.
86c
00c
0-
aa
D8N6
INN3
lstregler malgie-rigghe
(a)
(b)
(c)
Figure 2: Performance comparison of RLP w.r.t existing policies. (a) Average results all testing
cases, (b) Results on DNN models with 0 and 3 hidden layers. (c) Results on clusters with different
number of stragglers.
Radatbo Trabiing Tlne
Rulative Trobking Time
Trabiee
[id =1)
12:
08
Burdlass
2wordas
(a)
(b)
(c)
Figure 3: Evaluation on the generalization ability of RLP unseen cases. (a) Results on clusters
with different number of workers. (b) Results on new DNN model and new dataset. (c) Results on combination of clusters with different number of workers, new DNN model and new dataset.
different number of stragglers. RLP improves the time cost by 2%. 22% and 51% w.r.t the best existing policy when having 0, and multiple stragglers, respectively. The speedup ratio is more obvious when having more stragglers. This is due to when there exists stragglers, all existing policies perform in a similar way, so there leaves no room for RLP to improve anymore. By our observation, RLP performs in the stragglers, the straggling and staleness problems of all workers become more complex, so there exists more space for RLP explore new better policies. These detailed evaluation shows that RLP is both efficient and adaptive to speed up distributed SGD in different circumstances.
way the optimal ASP this time. When there exist more
Generalization Ability of RLP. In this set of experiments, we evaluate the generalization ability of RLP by applying it to process unseen instances with different configuration of clusters and models. First, we consider the generalization to clusters with different number of workers. We randomly choose 2,5 5 and 8 workers in the cluster and apply RLP to train the underlying DNN model having one hidden layer. The results are illustrated in Fig. ((a) We observe that RLP policy trained on cluster having 10 workers also performs much better than the existing policies on cluster with different number workers. We argue that this is due to the fact that in our RL formulation, both the state and action representation are irrelevant to the number of workers. Therefore, easy to transfer to a different cluster environment. Second, we examine the generalization ability model and data and show result in Fig. 3(b) For previously unseen DNN model with 6 hidden layers and a new CIFAR10 dataset, RLP improves the training time by 33% and 22% w.r.t. the best existing policy, respectively. This shows that our trained RLP can also generalize to train similar underlying ML models with unseen training datasets. This could be due the fact that we record only the information of the loss value in the state representation of RLP. Thus, training models with similar loss curve may also speed up by our RLP policy experience. Third, we combined these testing cases together evaluate RLP on new DNN model and dataset with a different number of workers. Fig- 3(c) shows that RLP also achieves good performance in this setting. improves the running time by 69%, 19% and 20% w.r.t. the best existing policy on 2, 5 and 8 workers, respectively, In summary, this set of experimental results shows that design of the states and actions for RLP should allow it to generalize to different settings.
RLP policy
of RLP w.r.t.
Detailed Insights into RLP. Finally, look more int the details of the learned RLP policy, To this end, in Fig. [4, we visualize the execution process of all workers in a typical case with multiple strag- glers, where the color and gray blocks represent that the worker is running and idle. respectively. Tab. summarizes statistics for each policy. Based on this, advantages and disadvantages of each policy become visible and obtain the following insights:

---

ICLR 2020
BSP
(b) ASP
(c) SSP (s=
(d) RLP
Figure 4: Visualization of the training process of SGD under different synchronization policies.
Table I: Statistics for different policies as observed in our experiments.
%
%
Policy
W;
W
W;
W;
Wio
iterations
idle time
inference time
28
28
54.24%
ASP
52
47
100
93
87
31
21
70
67
569
0
SSP
26
26
26
26
26
25
19
26
26
26
252
43.5%
8.1%
1.4%
Step
Time
9.17
10.103
4.37
4.47
15.39
22.189
4.97
5.07
Both BSP and SSP force all workers do almost the same number of computation steps but spend around half of the time in an idle state. In fact, SSP degenerates to BSP in the later steps. Therefore, although BSP performs less iterations than ASP, its total training time is longer.
For ASP, there exists no synchronization barrier for workers to wait, the number of steps done by workers are highly correlated to their step time (the correlation coefficient is around -0.9 ) However, due to the increased staleness level, ASP performs a larger number of iterations (2.4 times than BSP).
For RLP, we have two observations. First, RLP only synchronizes workers less times than neces- sary. We may observe very short idle times (only 8%) for RLP. Thus, RLP appears to never block the fastest workers. Similar ASP, the number of steps done by workers is also highly correlated to their step time, Second, the synchronization barriers chosen by RLP are all worthwhile. We find two typical cases:
the very beginning when the training loss decreases the fastest, synchronization can help to avoid staleness of parameters afterwards; 2) when some workers terminate their compu- tation steps in a similar time, worth to spend a small amount time synchronize in order to reduce staleness. Therefore, RLP exhibits lower level of staleness and needs much less iterations (only around 58%) than ASP, Moreover, the extra time spent on inference in RLP to generate actions only takes near 1%, so applying RLP adds very little extra cost to the SGD training process.
CONCLUSION AND FUTURE WORK
We have presented an RL-based framework used learn synchronization policies for PS-based training with distributed SGD. Based on the results of our experiments we argue the following points:
There exist synchronization policies for PS-based training beyond classic BSP, ASP and SSP which lead to shorter training time and improved resource utilization. It should worthwhile research endeavour to explore this topic for different types of models, in particular computationally demanding models.
Such synchronization policies may even provide better results in terms of used computational re- sources when used with slightly modified models, different cluster environments and training data.
The additional overhead incurred by having perform extra inference step for the RL policy network in the parameter server may still be less than the overall gain in efficiency from the policy when compared to BSP, ASP and SSP.
8

---

ICLR 2020
For future work we plan to deploy RLP in real-world scenario range of different model classes and gain further insights into its practicability. Moreover, we hope to further formalize our approach in order to gain theoretical insights into the existence of optimal synchronization policies.
We further believe that reinforcement learning might be a valuable research tool for the distributed systems community supporting the exploration and discovery of new policies for control problems which are typically encountered in this field.
REFERENCES
Léon Bottou. Stochastic gradient descent tricks. In Neural networks: Tricks of the trade, pp. 421- 436. Springer, 2012.
Jianmin Chen, Xinghao Pan, Rajat Monga, Samy Bengio, and Rafal Jozefowicz. Revisiting dis- tributed synchronous sgd. arXiv preprint: 1604.0091,, 2016.
Henggang Cui, Hao Zhang, Gregory R Ganger, Phillip Gibbons, and Eric P Xing. Geeps: Scalable deep learning distributed gpus with gpu-specialized parameter server. In Proceedings the 11th European Conference on Computer Systems. pp. 4. ACM, 2016.
Wei Dai, Abhimanu Kumar, Jinliang Wei, Qirong Ho, Garth Gibson, and Eric P Xing. High- performance distributed ml at scale through parameter server consistency models. In Twenty- Ninth AAAI Conference on Artificial Intelligence, 2015.
Wei Dai, Yi Zhou, Nanqing Dong, Hao Zhang, and Eric P. Xing. Toward understanding the im- pact of staleness in distributed machine learning. In 7th International Conference on Learning Representations, 2019.
Jacob Devlin, Ming-Wer Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint: 1810.04805, 2018.
Yan Duan, Xi Chen, Rein Houthooft, John Schulman, and Pieter Abbeel. Benchmarking deep reinforcement learning for continuous control. In International Conference on Machine Learning, pp. 1338, 2016.
Sanghamitra Dutta, Gauri Joshi, Soumyadip Ghosh, Parijat Dube, and Priya Nagpurkar. Slow and stale gradients can win the race: error-runtime trade-offs in distributed sgd. In International Conference Artificial Intelligence and Statistics, pp. 803 812, 2018.
Wenfei Fan, Ping Lu, Xiaojian Luo, Jingbo Xu, Qiang Yin, Wenyuan Yu, and Ruiqi Xu. Adap- tive asynchronous parallelization of graph algorithms. In Proceedings of the 2018 International Conference Management of Data, pp. 1141-1156. ACM, 2018.
Shixiang Gu, Timothy Lillicrap. Zoubin Ghahramani, Richard E Turner, and Sergey Levine. Q-prop: Sample-efficient policy gradient with an off-policy critic. arXiv preprint: 161 1.02247, 2016.
Stefan Hadjis, Ce Zhang, Ioannis Mitliagkas, Dan Iter, and Christopher Ré. Omnivore: An optimizer for dulti-device deep learning on cpus and gpus. arXiv preprint: 16t...44877704 2016.
Aaron Harlap, Henggang Cui, Wei Dai, Jinliang Wei, Gregory R Ganger, Phillip B Gibbons, Garth A Gibson, and Eric Xing. Addressing the straggler problem for iterative convergent parallel ml. In Proceedings of the Seventh ACM Symposium Cloud Computing, pp. 98-111. ACM, 2016.
Qirong Ho. James Cipar, Henggang Cui, Seunghak Lee, Jin Kyu Kim, Phillip B Gibbons, Garth A Gibson, Greg Ganger, and Eric P Xing. More effective distributed ml via a stale synchronous parallel parameter server. In Advances in neural information processing systems, pp. 1223-1231, 2013.
Jiawei Jiang, Bin Cui, Ce Zhang, and Lele Yu. Heterogeneity distributed parameter servers. In Proceedings of the 2017 ACM International Conference Management of Data, pp. 463-478. ACM, 2017.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 11..14422.9990 69800 2014.

---

ICLR 2020
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convo- lutional neural networks. In Advances in neural information processing systems, pp. 1097-1105, 2012.
Ke Li and Jitendra Malik. Learning optimize. arXiv preprini arXiv: 2016.
Hongzi Mao, Mohammad Alizadeh, Ishai Menache, and Srikanth Kandula. Resource management with deep reinforcement learning. In Proceedings of the 15th ACM Workshop on Hot Topics in Networks, pp- 50-56. ACM, 2016.
Hongzi Mao, Malte Schwarzkopf, Shaileshh Bojja Venkatakrishnan, Zili Meng, and Mohammad Alizadeh. Learning scheduling algorithms for data processing clusters. In Proceedings of the ACM Special Interest Group on Data Communication, pp. 270-288. ACM, 2019.
Ryan Marcus, Parimarjan Negi. Hongzi Mao, Chi Zhang, Mohammad Alizadeh, Tim Kraska, Olga Papaemmanouil, and Nesime Tatbul. Neo: A learned query optimizer. arXiv preprint arXiv:1904.03711, 2019.
Toannis Mitliagkas, Ce Zhang, Stefan Hadjis, and Christopher Ré. Asynchrony begets momentum, with an application to deep learning. In 2016 54th Annual Allerton Conference on Communica- tion, Control, and Computing (Allerton), pp. 997-1004. IEEE, 2016.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv 13eprint:1312566/00 2013.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mastering the game go with deep neural networks and tree search. nature, 529(7587):484, 2016.
Leslie Valiant, A bridging model for parallel computation. Communications of the ACM, 33(8): 103-111, 1990.
Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q- learning. In Thirtieth AAAI conference on artificial intelligence, 2016.
Xing Zhao, Aijun An, Junfeng Liu, and Bao Xin Chen. Dynamic stale synchronous parallel dis- tributed training for deep learning. arXiv preprint: 1908. 1848, 2019.
Shuxin Zheng, Qi Meng, Taifeng Wang, Wei Chen, Nenghai Yu, Zhi-Ming Ma, and Tie- Yan Liu. Asynchronous stochastic gradient descent with delay compensation. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pp. 4120-4129. JMLR. org, 2017.
Martin Zinkevich, John Langford, and Alex Smola. Slow learners are fast. In Advances in neural information processing systems, pp. 2331-2339, 2009.
10


Reviews:
[{"review_id": "SkedtDontS", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "This paper studies how to improve the synchronization policy for parameter server based distributed SGD algorithm. Existing work focus on either synchronous or asynchronous policy, which often results in straggler or staleness problems. Recent research proposes different ways to make the policy design fully adaptive and autonomous. This paper proposes a reinforcement learning (RL) approach and shows promising results to reduce the total training time on various scenario. \n\nA major challenge is to design the state and action spaces in the reinforcement learning setting. It requires the design can be generalized to different scenario, while ensuing efficient policy learning. Compared to existing policies such as BSP, ASP and SSP, RL has an advantage to adapt to non-stationary situations over the training process. Compared to other existing approaches, RL could be fully data-driven. \n\nThe paper formalizes an RL problem by minimizing the total training time to reach a given validation accuracy. To minimize the number of actions, only 3 actions coming from BSP, ASP and SSP are used. The state space is designed to capture similar loss curves, and to be independent of the number of workers (as much as possible). A policy network is used make decisions after trained with exiting methods. \n\nNumerical results validate that the RL policy improves the training time compared to BSP, ASP and SSP. Different number of works and models, dataset are also tested to show the RL policy is generalizable to unseen scenario. Although all the results are simulated in a controlled environment, Figure 4 gives a very interesting illustration showing the advantage of using the RL policy. I still have detailed comments (see below), but I find the paper well written, and the author(s) has obtained promising results.\n\nDetailed comments:\n-\tThe validation accuracy 88% on MINST seems pretty low to me to stop the algorithm, in particular when training multiple layer neural networks. What would happen if the accuracy is increased, can the RL approach still find a good policy? What about the validation accuracy on CIFAR-10?\n-\tI still have some concern of the computation time obtain the RL state per step. In particular, the time cost to compute the loss L on different weights w. How do you address this issue? Is L computed on the validation set? What is its size? This parameter seems to me highly sensitive when the policy is used for different dataset, in particular the dataset vary. It would be better to have more discussions in the paper or appendix. \n-\tWhat is the final test accuracy on the trained models using different policies? This allows us to see whether the approach has not over-fitted to the training/validation set. \n\nSome typo:\nPage 2 line 2 AP -> ASP\nPage 5, last line 4: converge -> convergence\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}, {"review_id": "Byl_kmjk9S", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "Background disclaimer: I work in RL research for quite an amount of time, but I do not know much about the domain of distributed systems. For this reason, I may not know the details of technical terms, and I might not be the best person to review this work (when compared with the literature in this field). Nevertheless, below I try to give my evaluation based on reading the paper. \n\n====================\nIn this work, the authors applied value-based reinforcement learning to learn an optimal policy for global parameter tuning in the parameter server (PS) that trains machine learning models in a distributed way using  stochastic gradient descent. Example parameters include SGD hyper-parameters (such as learning rate) and system-level parameters. Immediate cost is to minimize the training time of the SGD algorithm, and i believe the states are the server/worker parameters. From the RL perspective, the algorithm used here is a standard DQN with discrete actions (choices of parameters). But in general I am puzzled why the action space is discrete instead of continuous, if the actions are the hyper-parameters. State transition wise, I am not sure if the states follow an action-dependent MDP transition, and therefore at this point I am not sure if DQN is the best algorithm for this applications (versus bandits/combinatorial bandits).  While it is impressive to see that RL beats many of the SOTA baselines for parameter tuning, I also find that instead of using data in the real system to do RL training, the paper proposes generating \"simulation\" data by training a separate DNN. I wonder how the performance would differ if the RL policy is trained on the batch real data.\n\n", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "BJlg5IHCFS", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "This paper proposes to use deep RL to learn a policy for communication in the parameter-server setup of distributed training. From the perspective, the problem formulation is a nice contribution. \n\nWhile it is a reasonable idea and the initial results are promising, the lack of an evaluation on a real cluster, or for training more computationally-demanding models, is limiting. I fully appreciate the need to perform experiments in a controlled environment, such as the ones reported in the paper. These are useful to validate the idea and explore its potential limitations. However, to truly validate such an idea completely it is also necessary to implement it and run it \"in the wild\" on an actual distributed system. From my experience, although performing such experiments is certainly more involved and challenging, there can also be significant differences in the outcomes when one goes to such an implementation. Normally these are due to discrepancies between the assumed/simulated model, and real system behavior.\n\nIs it clear that deep RL is needed for this application, as opposed to more traditional RL approaches (either tabular, with suitably quantized actions, or a simpler form of function approximation? And to ask in the other direction, did you consider using a more complex policy architecture, e.g., involving an LSTM or other recurrent unit?\n\n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors consider a parameter-server setup where the learner acts a server communicating updated weights to workers and receiving gradient updates from them. A major question then relates in the synchronisation of the gradient updates, for which couple of *fixed* heuristics exists that trade-off accuracy of updates (BSP) for speed (ASP) or even combine the two allowing workers to be at most k steps out-of-sync. Instead, the authors propose to learn a synchronisation policy using RL. The authors perform results on a simulated and real environment. Overall, the RL-based method seems to provide some improvement over the fixed protocols, however the margin between the fixed and the RL get smaller in the real clusters. This is actually the main concern raised by the reviewers as well (especially R2) -- the paper in its initial submission did not include the real cluster results, rather these were added at the rebuttal. I find this to be an interesting real-world application of RL and I think it provides an alternative environment for testing RL algorithms beyond simulated environments.   As such, I’m recommending acceptance. However, I do ask the authors to be upfront with the real cluster results and move them in the main paper.

Decision: Accept (Poster)