PASO: Step Parallel Stochastic Optimization

Distributed-memory stochastic gradient descent (SGD) for sparse workloads is performance-limited by two forms of irregularity, namely inter-process communication and nonuniform nonzero distributions. Inter-processor communication cost grows irregularly based on the SGD variant employed to solve the problem. Heavy-tailed nonzero distributions typical of real sparse data drive load imbalance that further inflates per-iteration runtime. This paper develops HybridSGD, a 2D-parallel SGD method that generalizes existing 1D s-step SGD and 1D Federated SGD with Averaging (FedAvg) into a continuous family by employing a 2D processor mesh, and shows that the optimal split is datasetand machine-dependent. We derive a closed-form α-β -γ cost model whose optimum interpolates between s-step SGD and Federated SGD, refine it with cache-aware compute, rank-aware bandwidth, and load-imbalance terms calibrated on a NERSC Cray EX system, and use it to select the algorithmic and partitioning parameters predictively. We further introduce a cache-friendly data partitioner that simultaneously bounds nonzero imbalance and per-rank cache footprint, replacing static, balanced partitioners that we show are suboptimal for skewed data. Our irregularity-aware cost model predicts the winning partitioner on every measured dataset and partitioner combination. Finally, on the LIBSVM benchmark suite, HybridSGD achieves 53× and 14.6× time-to-target-loss speedups on the sparse, high-dimensional url and news20 datasets over FedAvg, matches FedAvg on rcv1, and is outperformed by FedAvg on the dense epsilon dataset. This crossover is predicted by the cost model and shows that HybridSGD is most beneficial in sparse, high-dimensional, communication-and skew-limited regimes, while simpler FedAvg can remain preferable when dense local computation dominates. Our software is available at