A Model Parallel Proximal Stochastic Gradient Algorithm for Partially Asynchronous Systems

TL;DR

This paper introduces AsyB-ProxSGD, an asynchronous model parallel stochastic gradient algorithm designed for large-scale machine learning models, achieving convergence and near-linear speedup in distributed settings.

Contribution

It generalizes ProxSGD to asynchronous, model parallel environments and provides theoretical convergence guarantees along with practical implementation results.

Findings

Achieves $O(1/ oot{K})$ convergence rate for nonconvex problems.

Demonstrates near-linear speedup with increasing workers.

Validates effectiveness on real-world large-scale datasets.

Abstract

Large models are prevalent in modern machine learning scenarios, including deep learning, recommender systems, etc., which can have millions or even billions of parameters. Parallel algorithms have become an essential solution technique to many large-scale machine learning jobs. In this paper, we propose a model parallel proximal stochastic gradient algorithm, AsyB-ProxSGD, to deal with large models using model parallel blockwise updates while in the meantime handling a large amount of training data using proximal stochastic gradient descent (ProxSGD). In our algorithm, worker nodes communicate with the parameter servers asynchronously, and each worker performs proximal stochastic gradient for only one block of model parameters during each iteration. Our proposed algorithm generalizes ProxSGD to the asynchronous and model parallel setting. We prove that AsyB-ProxSGD achieves a convergence rate of $O(1/\sqrt{K})$ to stationary points for nonconvex problems under \emph{constant} minibatch sizes, where $K$ is the total number of block updates. This rate matches the best-known rates of convergence for a wide range of gradient-like algorithms. Furthermore, we show that when the number of workers is bounded by $O(K^{1/4})$, we can expect AsyB-ProxSGD to achieve linear speedup as the number of workers increases. We implement the proposed algorithm on MXNet and demonstrate its convergence behavior and near-linear speedup on a real-world dataset involving both a large model size and large amounts of data.