Learning Efficient Parameter Server Synchronization Policies for Distributed SGD

1 INTRODUCTION . In recent years , Stochastic gradient descent ( SGD ) Bottou ( 2012 ) and its variants Kingma & Ba ( 2014 ) ; Chen et al . ( 2016 ) , have been adopted as the main work horse for training machine learning ( ML ) models . To be able to train large models , which are both computationally demanding or require very large training datasets , SGD is often parallelized across several machines , with the well-known parameter-server ( PS ) framework being one of the most widely adopted distribution strategies . In 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 . Challenges of Distributed SGD in PS . In the PS setting , a central task is to design a synchronization 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 at the parameter server . Thus , it directly determines the overall training time . However , finding a good synchronization policy is difficult , as this will at 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 synchronous 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 a heterogeneous cluster , it is common that some workers , also referred to as 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 a negative impact on the overall convergence speed Chen et 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 ) . 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 a fixed threshold s , such that the fastest worker is allowed to outpace the slowest one by at most s steps ( “ bounded staleness ” ) . However , an optimal value of s is hard to specify without enough prior knowledge . Moreover , the optimal value of smay change during the entirety of the training process . Other synchronization policies ( see in Section 2 ) may either degenerate to the above mentioned three policies or may need extensive manual tuning of hyper-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 in the following , whenever we speak of “ overall training time ” or “ time until convergence ” , we specifically mean the amount of time needed until a model reaches a certain pre-defined accuracy 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 to 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 a 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 et al . ( 2013 ) ; Van Hasselt et al . ( 2016 ) . Furthermore , we design a pre-training process to speed up the convergence . Empirical results demonstrate the validity of our approach and the advantages of our learned RLP policy in 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 framework 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 ( 2014 ) in order to reduce the number of hyper-parameters and allow for a 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 demanding models , such as deep Convolutional Neural Networks Krizhevsky et al . ( 2012 ) or transformerbased 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 our 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 directly applied for distributed SGD in a PS-setting . As the performance of BSP is heavily diminished by the straggling problem , other variants , such as the K-sync BSP and K-batch-sync BSP have been proposed in Dutta et al . ( 2018 ) . They alleviate the straggling problem by slightly relaxing the synchronization condition . However , as 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 ) . However , they may case additional , non-negligible communication overhead Jiang et al . ( 2017 ) . Another extreme is to run distributed SGD using the ASP policy , which avoids the straggling problem but typically forces one to consider the staleness problem . The negative effects of staleness have been shown in Zinkevich et al . ( 2009 ) ; Chen et al . ( 2016 ) ; Cui et al . ( 2016 ) ; Dutta et al . ( 2018 ) ; Dai et al . ( 2019 ) on typical ML models , where Dai et al . ( 2019 ) performed a comprehensive experimental analysis . To correct bias caused by staleness , Hadjis et al . ( 2016 ) ; Mitliagkas et al . ( 2016 ) ; Zheng et al . ( 2017 ) 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 s in SSP which bounds the staleness is hard to tune and fixed during the training process . To overcome this , Jiang et al . ( 2017 ) ; Zhao et al . ( 2019 ) proposed a dynamic SSP policy where s is tuned during the training process . Fan et al . ( 2018 ) proposed a more flexible adaptive asynchronous parallel policy to allow different values of s for each worker . Although more adaptive , they need manual tuning of hyper-parameters . As a result , until now , there appears to exist no synchronization policy that is both fully adaptive and automatic . This paper , for the best of our knowledge , marks the first instance where RL is used find synchronization policy for distributed SGD in a completely data-driven fashion . RL has been successful applied to control robotics Duan et al . ( 2016 ) and games Mnih et al . ( 2013 ) ; Silver et al . ( 2016 ) . Recently , it has been widely used to optimize problems such as task scheduling Mao et al . ( 2019 ) , resource management Mao et al . ( 2016 ) and optimization Li & Malik ( 2016 ) ; Marcus et al . ( 2019 ) . 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 variations 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 regularly on different training data sets on the same cluster of machines ( or similar cluster in the case of e.g . dynamically allocated cloud computing resources ) . 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 .

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