Communication-Efficient, 2D Parallel Stochastic Gradient Descent for Distributed-Memory Optimization

TL;DR

This paper introduces HybridSGD, a 2D parallel stochastic gradient descent method that reduces communication costs and improves scalability in distributed-memory systems, outperforming existing algorithms in convergence and speed.

Contribution

It generalizes prior 1D SGD methods to a 2D approach, providing a theoretical framework and empirical evidence of improved performance and scalability.

Findings

HybridSGD converges faster than FedAvg at similar scales.

Achieves up to 5.3x speedup over s-step SGD.

Achieves up to 121x speedup over FedAvg.

Abstract

Distributed-memory implementations of numerical optimization algorithm, such as stochastic gradient descent (SGD), require interprocessor communication at every iteration of the algorithm. On modern distributed-memory clusters where communication is more expensive than computation, the scalability and performance of these algorithms are limited by communication cost. This work generalizes prior work on 1D $s$-step SGD and 1D Federated SGD with Averaging (FedAvg) to yield a 2D parallel SGD method (HybridSGD) which attains a continuous performance trade off between the two baseline algorithms. We present theoretical analysis which show the convergence, computation, communication, and memory trade offs between $s$-step SGD, FedAvg, 2D parallel SGD, and other parallel SGD variants. We implement all algorithms in C++ and MPI and evaluate their performance on a Cray EX supercomputing system. Our empirical results show that HybridSGD achieves better convergence than FedAvg at similar processor scales while attaining speedups of $5.3\times$ over $s$-step SGD and speedups up to $121\times$ over FedAvg when used to solve binary classification tasks using the convex, logistic regression model on datasets obtained from the LIBSVM repository.