Towards fast distributed training with generic and efficient communication scheduling

Abstract

With massive computing resources and sizeable high-quality datasets, machine learning (ML) has attained state-of-the-art performance for many practical applications. To speed up the ML training, data parallelism across multiple workers is widely adopted, where each worker has a full copy of the model parameters, trains on a subset of the input dataset, and synchronizes computed gradients with others. In distributed training, popular ML frameworks employ either the Parameter Server or All-Reduce architecture to facilitate the gradient exchange. Nevertheless, due to frequent gradient synchronization, the communication cost often stands out, which may take up to 90% of the training time. Communication scheduling has thus been adopted to accelerate the training procedure. In this thesis, to minimize the communication cost and training time, we focus on designing practical and efficient schedulers for both Parameter Server and All-Reduce.

In Parameter Server, the global parameters are stored in logically centralized servers, via which workers can exchange gradients synchronously or asynchronously. To reduce the communication cost for model synchronization, prior work overlaps communication with computation for deep neural networks. Yet, these methods require not only layer-by-layer model structures but also tailor-made efforts in runtime profiling and hyperparameter tuning. To make the overlapping optimization more efficient, we propose a new Parameter Server framework. Particularly, our solution decouples the dependencies between push and pull operations, and allows workers to eagerly pull the global parameters. As a result, both push and pull operations can be easily paralleled with computations. Moreover, the overlapping strategy offers a new way to address the straggler problem, where stale updates greatly retard the training procedure. This is because, with adequate information available to workers, they can explicitly adjust their local learning rates to limit the adverse impact of stale updates on the global parameters. Experiments show that our prototype could save up to 54% less time for each iteration and up to 37% fewer iterations for model convergence, compared with those widely-used schemes.

On the other hand, in All-Reduce, there are no centralized servers, and workers only communicate with their neighbors. To mitigate the synchronization cost, efficient scheduling of all-reduce operations is essential for minimizing training time. However, existing schedulers either are agnostic to ML workload characteristics, or require runtime profiling of workloads and fine-tuning of system hyperparameters, rendering them less efficient or not generic. To this end, we propose a new All-Reduce scheduler that is both efficient and generic. Specifically, we decompose the layer-to-layer dependencies between the communication of one iteration and the computation of the next, thereby allowing all-reduce operations to be paralleled with computations; and employ a simple approximation for the parameters to guarantee the training quality. Besides, we demonstrate that other techniques like tensor fusion can be easily incorporated into the proposed scheduler. Experiments show that our prototype can accelerate training on representative vision and language tasks by up to 29.2%, compared with popular All-Reduce schedulers.