1 Introduction↩︎

Stochastic gradient descent (SGD) and its communication-efficient variants are the driving iterative optimizers behind much of modern machine learning. When the underlying problem is sparse and distributed across many processors, two forms of irregularity shape the runtime behavior of distributed SGD. First, inter-processor communication at every iteration [1], [2] requires careful algorithm design to ensure that the communication does not dominate local sparse BLAS computation. Second, the sparsity structure of the input dataset may produce work imbalance between processors which can interact non-trivially with the communication and local sparse computation costs. We show that real sparse data (e.g., rcv1, news20, and url) from the LIBSVM repository [3] exhibit heavy-tailed nonzero-per-row and nonzero-per-column distributions that produce work imbalances that inflate per-iteration runtime and interact non-trivially with the choice of data partitioner. Furthermore, models such as the classical \(\alpha\)-\(\beta\)-\(\gamma\) cost model (i.e., Hockney’s model [4]) are not well-suited for problems that contain data-dependent work imbalance.

Prior work has approached the communication cost issue from two distinct algorithmic directions. Communication-avoiding (\(s\)-step) SGD [5] reorganizes the inner loop using recurrence unrolling to group \(s\) consecutive gradient steps into a single Allreduce. This approach reduces the communication cost at the expense of additional bandwidth and computation but does not change the convergence rate. Communication-efficient methods such as FedAvg [6] perform \(\tau\) independent local updates between Allreduces. This approach reduces the per-iteration latency and bandwidth costs at the expense of convergence, where the model drifts from the global optimum. Both approaches are limited to 1D processor grid distributions which only allow a fixed data partitioning scheme. Specifically, \(s\)-step SGD partitions by columns of \(\boldsymbol{A}\), while FedAvg partitions by rows. Our approach is based on the observation that the resulting partitionings are compatible because they live on orthogonal axes of a 2D processor mesh (\(p = p_r \times p_c\)), and we exploit this to derive a family of 2D-parallel SGD algorithms (HybridSGD) that can navigate a continuous performance trade-off between the two 1D baselines as extremes in each processor mesh dimension.

Since the optimal mesh, data partitioner, and \(s\)-step recurrence unrolling length all depend jointly on the dataset’s nonzero skew and on the target machine’s hardware parameters, we also integrate hardware and data irregularity-aware performance modeling into HybridSGD. The closed-form performance model that we develop is coupled with empirical measurements of the hardware parameters \(\alpha\), \(\beta\), and \(\gamma\). The empirical measurements highlight an inflection point in MPI Allreduce bandwidth at the per-node rank boundary. This insight yields a parameter-free topology rule for the HybridSGD mesh dimensions that accurately predict the winning \((p_r, p_c)\) combination on a NERSC Cray EX system. Furthermore, we show that our load-imbalance aware model accurately predicts the empirical ranking of partitioners on every dataset and partitioner combination tested.

The contributions of this work are:

1. HybridSGD, a 2D-parallel SGD framework that generalizes \(s\)-step SGD [5] and FedAvg [6] by parameterizing the processor grid as \(p = p_r \times p_c\), where each 1D baseline is an extreme grid dimension. We present a C++/MPI implementation built on Intel MKL sparse BLAS with three selectable column partitioners (rows, nonzero-greedy, cyclic).

2. A closed-form analytic cost model for HybridSGD derived from Hockney’s \(\alpha\)-\(\beta\)-\(\gamma\) model, with closed-form optima for \(s\) (recurrence unrolling length) and \(b\) (batch size), and a topology-respecting rule for the mesh split \((p_r, p_c)\) that selects the row team to one node so the frequent row Allreduce stays within a node. The rule requires only the per-node rank count and per-core cache capacity to accurately predict the optimum on every dataset tested. This rule fixes the mesh dimension, after which the model ranks the remaining \((s, b, \tau, \text{partitioner})\) configurations and identifies the operating regime; we use it as a selection tool rather than an absolute-runtime predictor.

3. An irregularity-aware partitioning study that formulates partitioner selection as a two-objective constrained problem that minimizes the nonzero imbalance (\(\kappa\)) subject to a per-rank cache footprint. We show that a greedy partitioner which aims to balance nonzeros may lead to data distributions that overload a subset of processors (i.e., lead to cache spill), especially on column-skewed data. We introduce a cyclic partitioner that satisfies both objectives in expectation and validate the model using real (url, news20, rcv1) and skew-controlled synthetic data.

4. Strong scaling speedups on the LIBSVM benchmark data measured on a NERSC Cray EX system. HybridSGD attains \(\mathbf{53\times}\) speedup on url and \(\mathbf{14.6\times}\) on news20, matches FedAvg on rcv1, and is outperformed by FedAvg on the dense epsilon dataset. The qualitative crossover point is predicted by our cost model and confirmed empirically.

2 Background↩︎

This section briefly surveys prior and recent work that aim to improve communication efficiency of distributed optimization in theory and practice. These techniques can broadly be categorized based on whether the techniques maintain the convergence rates and behavior of classical optimization methods (i.e., communication-avoiding methods) or whether they weaken convergence rates and guarantees for better performance (i.e., communication-efficient methods).

2.1 Related work↩︎

\(s\)-step methods were originally developed to reduce the frequency of inter-process synchronization in Krylov methods for linear systems and spectral problems [7]–[11], with subsequent advances in matrix-powers and tall-skinny QR kernels [12], stability analysis and error correction [13]–[16], and parallel implementations such as \(s\)-step BiCGSTAB [17] achieving \(2.5\times\) speedups on scientific applications. The technique was later generalized to nonlinear, convex optimization [5], [18]–[23], yielding \(s\)-step variants of coordinate descent, SGD, and subsampled Newton with speedups up to \(5\times\) in parallel cluster and cloud environments and no change to convergence behavior. A complementary line of work trades sequential consistency or solution accuracy for parallel performance. Asynchronous shared-memory SGD and CD [24]–[26] eliminate blocking synchronizations under sparse update assumptions. Distributed-memory approaches reduce communication via hierarchical low-rank decomposition for kernel ridge regression [27], K-means redistribution for kernel SVM [28], and local optimization with deferred averaging [6], [28]–[30], which is the basis for federated learning. In this work we integrate the \(s\)-step and federated learning approaches into a single 2D-parallel algorithm with continuous trade-offs between the two extremes. A separate line of work reduces AllReduce traffic via gradient compression by leveraging quantization [31] and sparsification [32] to shrink each \(n\)-word gradient message by \(16\)–\(1000\times\) at the cost of biased or noisy aggregation.

HybridSGD takes a structurally different approach by partitioning the weight vector \(\boldsymbol{x}\) across \(p_c\) column-team ranks. This reduces the message size of weights synchronization to \(n/p_c\) words losslessly. The row-team is then utilized for the \(s\)-step recurrence unrolled computations which carries \(O(s^2b^2/p_r^2)\) words and is independent of \(n\). We analyze the computation, latency, and bandwidth costs of HybridSGD in 5. The HybridSGD approach can integrate both asynchronous communication and compression techniques as these methodologies are orthogonal to the HybridSGD approach. However, in this work, we aim to fully characterize the HybridSGD framework before combining orthogonal strategies that require rigorous convergence analysis to show effectiveness. While alternatives to SGD exist for logistic regression (coordinate descent, IRLS, quasi-Newton, higher-order methods), we focus on SGD as it is the workhorse optimization method for many machine learning problems. We therefore focus on SGD and leave a broader methods generalization to future work.

2.2 Notation↩︎

Our datasets are sparse matrix-vector pairs \((\boldsymbol{A}, \boldsymbol{y})\) with \(\boldsymbol{A} \in \mathbb{R}^{m \times n}\) (\(m\) samples, \(n\) features) and binary labels \(\boldsymbol{y} \in \{+1, -1\}^m\). Bold uppercase letters denote matrices, bold lowercase vectors, and nonbold lowercase scalars. All Greek letters are tunable scalar quantities. Subscripts are used either to index iterations of an algorithm or to index entries of a vector/matrix. We distinguish the latter by typeface. For example, \(y_i\) is the \(i\)-th entry of \(\boldsymbol{y}\), \(\boldsymbol{a}_{i,:}\) is the \(i\)-th row of \(\boldsymbol{A}\), and \(a_{i,j}\) is the \((i,j)\) entry. The function \(\mathop{\mathrm{diag}}(\boldsymbol{y}) \colon \mathbb{R}^m \mapsto \mathbb{R}^{m \times m}\) maps \(\boldsymbol{y}\) to a diagonal matrix with \(d_{i,i} = y_i\). Bracketed superscripts identify per-processor quantities, e.g.\(\boldsymbol{A}^{[i]}\) is the local block of \(\boldsymbol{A}\) on processor \(i\).

3 Optimization Problem↩︎

Given a matrix \(\boldsymbol{A} \in \mathbb{R}^{m \times n}\) and a label vector \(\boldsymbol{y} \in \{+1, -1\}^m\), we seek \(\boldsymbol{x} \in \mathbb{R}^n\) which solves the unregularized logistic regression problem \[\label{eq:logreg} f(\boldsymbol{A}, \boldsymbol{y}, \boldsymbol{x}) := \min_{\boldsymbol{x} \in \mathbb{R}^{n}}\frac{1}{m} \sum_{i = 1}^{m} \log\bigl(1 + \exp(-y_i \cdot \boldsymbol{a}_{i,:}\boldsymbol{x})\bigr).\tag{1}\] Since 1 does not have a closed-form solution, we use the gradient descent formulation, \(\boldsymbol{x}_k = \boldsymbol{x}_{k-1} - \eta \boldsymbol{g}_k\), with a fixed step-size, \(\eta\), where \[\begin{align} \boldsymbol{u}_k &= \frac{\vec{\boldsymbol{1}}}{\vec{\boldsymbol{1}} + \exp\bigl(\mathop{\mathrm{diag}}(\boldsymbol{y}) \cdot \boldsymbol{A} \cdot \boldsymbol{x}_{k-1}\bigr)}\tag{2}\\ \boldsymbol{g}_k &= -\frac{1}{m}\bigl(\mathop{\mathrm{diag}}(\boldsymbol{y}) \cdot \boldsymbol{A}\bigr)^\intercal \boldsymbol{u}_k.\tag{3} \end{align}\] We pre-compute \(\mathop{\mathrm{diag}}(\boldsymbol{y}) \cdot \boldsymbol{A}\) once by scaling rows of \(\boldsymbol{A}\) by the corresponding labels. Each iteration therefore consists of one SpMV to form \(\boldsymbol{u}_k\), a nonlinear sigmoid step on an \(m\)-dimensional vector, and a transposed SpMV to form \(\boldsymbol{g}_k\). Mini-batch SGD (1) sub-samples \(b\) rows of \(\boldsymbol{A}\) via \(\boldsymbol{S}_k \in \mathbb{R}^{b \times m}\) built from \(b\) rows of the identity, reducing per-iteration work by \(b/m\) while converging faster per unit computation in practice [33]. We design parallel variants of SGD throughout.

4 Parallel Algorithms Design↩︎

The main computation and communication bottlenecks in 1 are the SpMV with the subsampled matrix \(\boldsymbol{S}_k \cdot \mathop{\mathrm{diag}}(\boldsymbol{y}) \cdot \boldsymbol{A}\) at [alg:sgd-uk,alg:sgd-gk]. In a parallel setting, forming \(\boldsymbol{u}_k\) and/or \(\boldsymbol{g}_k\) requires communication. 2 shows the three layouts. Under 1D-row partitioning, [alg:sgd-uk] requires an Allreduce on an \(n\)-dimensional vector during the transposed-SpMV. Under 1D-column partitioning, [alg:sgd-uk] requires an Allreduce on the \(b\)-dimensional vector, \(\boldsymbol{u}_k\). Under 2D partitioning (\(p = p_r \times p_c\)) both SpMV computations require an Allreduce but are smaller (\(b/p_r\) and \(n/p_c\) words respectively). These partitionings expose a rich performance trade-off space which we analyze in 5.

4.1 Communication-efficient SGD↩︎

We focus on two communication-efficient 1D variants: federated SGD with Averaging (FedAvg) [6] which combines 1D-row partitioned SGD with deferred communication and \(s\)-step SGD [5] which combines 1D-column partitioned SGD with recurrence unrolling.

3 partitions \((\boldsymbol{A}, \boldsymbol{y})\) into 1D-row layout across \(p\) processors and performs \(\tau\) local sequential SGD iterations on each processor with an Allreduce of the local solutions \(\tilde{\boldsymbol{x}}_k^{[i]}\) with averaging. The local-SGD analysis [6] (summarized in the FedAvg row of 1) shows that FedAvg trades convergence for performance as \(p\) and \(\tau\) grow. \(\tau = 1\) degenerates to synchronous mini-batch SGD on an effective global batch of \(p b\), while \(p = 1\) reduces to sequential SGD.

4 defers communication for \(s\) iterations without affecting convergence behavior [5] at the cost of additional computation and bandwidth because message sizes grow proportional to \(sb\). [alg:sstep-sgd-gram] forms a Gram matrix \(\boldsymbol{G}\) whose blocks correct \(\boldsymbol{u}_{sk+j}\) in [alg:sstep-sgd-correction] and [alg:sstep-sgd-correction-end] for the deferred solution updates. 1D-column partitioning of \(\boldsymbol{A}\) yields the cheapest communication for 4. Any other layout introduces an additional 1D or 2D matrix multiply to form \(\boldsymbol{G}\).

4.1.0.1 HybridSGD Design

We combine the two by arranging \(p = p_r \times p_c\) processors into a 2D grid where row teams perform FedAvg on \(n/p_c\) fractions of \(\boldsymbol{x}\) and column teams perform \(s\)-step SGD on \(b/p_r\) independent batches of rows of \(\boldsymbol{A}\). Algorithmically, HybridSGD is obtained by 2D-partitioning \(\boldsymbol{A}\) in 3 and replacing the SGD call with a call to 4. We require \(s \leq \tau\) because the local solution vectors \(\tilde{\boldsymbol{x}}_k^{[i]}\) are averaged only every \(\tau\) iterations.

5 Algorithms Analysis↩︎

[tab:comp-costs,tab:comm-costs] summarize the theoretical costs of the parallel SGD variants studied. We assume \(\boldsymbol{A} \in \mathbb{R}^{m \times n}\) has \(\bar{z}\) nonzeros per row uniformly distributed across \(p\) processors, yielding \(m\bar{z}/p\) nonzeros per rank and \(b\bar{z}\) total nonzeros per mini-batch of size \(b\). \(\boldsymbol{A}\) is stored in Compressed Sparse Row (CSR) layout. Sub-sampling of rows is performed cyclically via \(i = (i + b) \bmod m\) which permits the row-index array to be reconstructed cheaply, and we pad \(\boldsymbol{A}\) so \(m \equiv 0 \pmod{s_{\text{max}}\cdot b}\).

5.1 Computation, Convergence, and Storage↩︎

We model algorithm cost as \(T = T_{\text{comp}} + T_{\text{comm}}\) with \(T_{\text{comp}} = \gamma F\), where \(F\) is the flop count and \(\gamma\) is seconds per floating-point operation. Convergence rates for SGD, \(s\)-step SGD, and FedAvg are well-known, and we cite them in 1 and use them below as needed. \(s\)-step SGD is an algebraic reformulation of 1 [5] and converges identically up to floating-point error. FedAvg [6] attains rate \(1/(\tilde{K}bp)\) provided \(\tau = O(\sqrt{\tilde{K}/(bp)})\). Exceeding this bound at fixed \(\tau\) as \(p\) grows degrades convergence through approximation error. This effect is pronounced on non-I.I.D. datasets. We note that this convergence degradation further motivates HybridSGD since it allows for tunable \(p_r\), where the column dimension \(p_c\) can absorb additional parallelism without reducing the FedAvg convergence rate or reducing \(\tau\).

The FedAvg and HybridSGD rows of 1 follow from 1 by specializing the 2D-SGD bound. 1D/2D parallel SGD and FedAvg have comparable compute and storage, while \(s\)-step SGD and HybridSGD inflate both. FedAvg beats parallel SGD when \(\tilde{K} < K/\tau\), which follows from its linear-in-\(p\) convergence speedup at bounded \(\tau\). \(b\) is treated as a global batch size for cross-algorithm comparison. In practice each algorithm tunes its own \(b\), which we exploit in 7.

5.2 Communication↩︎

We use Hockney’s (\(\alpha\)-\(\beta\)) model, \(T_{\text{comm}} = \alpha L + \beta W\), where \(L\) is messages, \(W\) is words moved, and \((\alpha, \beta)\) are hardware parameters. We assume all algorithms use MPI Allreduce with the reduce-scatter + all-gather bound \(L = 2\log p\), \(W = d\) [1], [2], and that computation is performed redundantly on all ranks after reduction. 2 summarizes the resulting per-algorithm bounds. Per-iteration bounds follow by dropping the \(K, \hat{K}, \tilde{K}\) factors. We prove 2D SGD below and specialize to 1D variants. The \(s\)-step SGD bound is from [5].

The FedAvg and HybridSGD rows of 2 follow from 2 by the same specialization. \(s\)-step SGD trades compute, bandwidth, and latency as a function of \(s\). FedAvg trades convergence against performance as a function of \(p\) and \(\tau\). 1 shows that HybridSGD’s 2D grid interpolates between the two extremes, with compute and bandwidth relative to \(s\)-step SGD shrinking by \(p_r^2\) at the cost of a convergence rate that is faster than FedAvg whenever \(p_c > 1\). In particular, \(p_r\) selects a convergence curve and \(p_c\) scales parallelism along it without further convergence cost. The additional bandwidth from communicating across \(p_c\) is mitigated by tuning \(s\), \(b\), \(p_r\), and \(p_c\), while \(n\) is fixed by \(\boldsymbol{A}\).

6 \(\alpha\)-\(\beta\)-\(\gamma\) Cost Model and 2D-Mesh Optimum↩︎

This section derives a closed-form runtime model \(T(p_r, p_c, s, b, \tau, \alpha,\beta,\gamma)\) for HybridSGD that unifies the bounds of 5. We use the model primarily as a ranking and selection tool: its role is to order candidate \((p_r, p_c, s, b, \tau, \text{partitioner})\) configurations and identify the operating regime, not to predict wall-clock time to high accuracy. As we show in 6.5, the model’s absolute-runtime error can reach \(2\) to \(10\times\), but its ranking of configurations is reliable, and ranking is the only property the selection rules below depend on. Minimizing \(T\) over the possible mesh dimensions \(p = p_r \times p_c\) yields a continuous family of optima that interpolates between pure 1D \(s\)-step SGD (\(p_r{=}1\), \(p_c{=}p\)) and pure FedAvg (\(p_r{=}p\), \(p_c{=}1\)).

6.1 Machine model and per-solver costs↩︎

We use Hockney’s two-term model, where one Allreduce over \(q\) ranks carrying a payload of \(W\) words costs \(T_{\text{comm}}(q, W) = 2\lceil\log_2 q\rceil\alpha + W\beta\), with the \(2\log_2 q\) factor following from bandwidth-optimal reduce-scatter + all-gather [1], [2]. Computation is \(T_{\text{comp}} = \gamma F\). Machine parameters \((\alpha, \beta, \gamma)\) are empirically measured on the NERSC Cray EX system used in the experiments in 7.1, and \(w\) is the word size in bytes. 3 reports per-sample \((\alpha, \beta, \gamma)\) costs for all six solvers, amortized over each solver’s communication period.

6.2 Closed-form runtime model↩︎

Each rank holds \(m/p_r\) local rows and \(n/p_c\) local columns of \(\boldsymbol{A}\). One full round is \(\tau\) consecutive 1D \(s\)-step SGD bundles on each row communicator (\(sb\tau/p_r\) row-team samples) followed by one column Allreduce. With the \(p_r\) row teams running in parallel, the per-epoch round count is \(m/(sb\tau)\). The per-rank compute contribution to the per-epoch wall is \((m/p)\cdot(6\bar{z} + 2sb)\gamma\) — the total floating-point work \(m(6\bar{z} + 2sb)\) is shared across \(p\) ranks — while each per-rank communication contribution is \(m\) times the corresponding per-sample Allreduce wall, since Allreduce runtime is borne by every participating rank and is not divided across them. Summing yields

with \(p = p_r p_c\). Compute grows linearly with \(sb\), dominated by the \(O(s^2b^2)\) correction loop at large \(sb\). Latency decreases as \(sb\tau\) grows. Gram bandwidth tracks the \(\binom{s}{2}b^2\) entries of the block Gram message per row team. Here \(b\) is the per-row-team batch size (identical to the 1D \(s\)-step SGD row), so the Gram payload is independent of \(p_r\). Sync bandwidth decreases as \(sb\tau p_c\) grows since only \(n/p_c\) words are Allreduced per column sync.

6.2.0.1 Baselines as limits

At \(p_r{=}1, p_c{=}p, \tau \to \infty\) the column Allreduce vanishes and [eq:model] reduces to \(\frac{m}{p}(6\bar{z}+2sb)\gamma + m\bigl[\frac{2\alpha\log p}{sb} + \frac{(s-1)bw\beta}{2}\bigr]\), the pure 1D \(s\)-step SGD cost. At \(p_r{=}p, p_c{=}1, s{=}1\) the row Allreduce vanishes and it reduces to \(\frac{m}{p}(6\bar{z}+2b)\gamma + m\bigl[\frac{2\alpha\log p}{b\tau} + \frac{nw\beta}{b\tau}\bigr]\), the pure FedAvg cost. MB-SGD follows from \(\tau{=}1\) in the FedAvg limit. These limits confirm that [eq:model] subsumes the solver hierarchy of 3 and that HybridSGD is a strict generalization.

6.3 Optimal parameters↩︎

Write \(\widetilde{L} = \tau\log p_c + \log p_r\). Both the latency and sync-BW terms decrease strictly in \(\tau\), so the communication-optimal \(\tau\) is unbounded. In practice, convergence drift from local weight staleness sets a dataset- and step-size-dependent upper bound [6], [34]. Collecting terms in \(s\) at fixed \(b, \tau, p_r, p_c\) yields a convex \(A_s s + B_s/s + C_s\) with \(A_s = 2\gamma b/p + b w\beta/2\) and \(B_s = 2\alpha\widetilde{L}/(b\tau) + nw\beta/(b\tau p_c)\), minimized at \[\label{eq:sstar} s^* = \sqrt{B_s / A_s} = \sqrt{\tfrac{2\alpha\widetilde{L}/(b\tau) + nw\beta/(b\tau p_c)} {(2\gamma/p + w\beta/2)\,b}}.\tag{4}\] The analogous derivation in \(b\) gives \[\label{eq:bstar} b^* = \sqrt{\tfrac{2\alpha\widetilde{L}/\tau + nw\beta/(\tau p_c)} {(2\gamma s/p + (s{-}1)w\beta/2)\,s}}.\tag{5}\] One step of fixed-point iteration on 4 and 5 yields the joint \((s^*, b^*)\). Equating the Gram-BW and sync-BW terms gives the bandwidth balance \((s{-}1)sb^2\tau p_c \approx 2n\). With \((p_r, p_c)\) pinned by the topology rule (6 ), the balance controls only the choice of \(s, b, \tau\). Above the balance the Gram message dominates and \(s\) or \(b\) should shrink. Below the balance the weight sync dominates and \(\tau\) should grow.

6.3.0.1 Optimal mesh split

With \(p_c = p/p_r\) and \((s, b, \tau)\) fixed, \(\partial T/\partial p_r = 0\) produces a transcendental fixed-point equation in \(p_r\) whose coefficients depend on \(\beta_{\text{row}}(p_c)\) (the row-Allreduce bandwidth over \(p_c\) ranks) and \(\beta_{\text{col}}(p_r)\) (the column-Allreduce bandwidth over \(p_r\) ranks) through the rank-aware Hockney bandwidth. The two coupled bandwidths in turn depend on the mesh split itself, so \(p_r^{*}\) has no closed form in general. A simple structural observation removes the need to solve the fixed point. The calibrated \(\beta(q)\) on Perlmutter CPU (7) is approximately a step function in the per-Allreduce rank count \(q\). Intra-node values (\(q \leq R\), where \(R\) is the per-node rank count) sit in \([5{\times}10^{-11},\,2.7{\times}10^{-9}]\) s/B, while inter-node values (\(q > R\)) jump to \([2.7{\times}10^{-9},\,6.6{\times}10^{-9}]\) s/B, an order-of-magnitude discontinuity at \(q = R\). Holding the row team to one node (\(p_c \leq R\)) therefore keeps the frequent row Allreduce on shared-memory transport, while crossing the node boundary (\(p_c > R\)) both inflates the Gram-BW term by the \(\sim 5{\times}\) \(\beta\) step and crosses into the slow inter-node regime. Sliding \(p_c\) upward along the constraint \(p_r p_c = p\) reduces the sync-BW term (\(\propto n/p_c\)) monotonically inside the intra-node piece \(p_c \leq R\). The kink at \(p_c = R\) then sets the optimum: \[\label{eq:meshrule} p_c^{*} = \max\Bigl(\bigl\lceil n w / L_{\text{cap}} \bigr\rceil,\; \min(R,\,p)\Bigr), \qquad p_r^{*} = p / p_c^{*},\tag{6}\] where the cache term \(\lceil n w / L_{\text{cap}} \rceil\) raises \(p_c^{*}\) above \(R\) only when the per-rank weight slab \(n w / p_c\) would spill the chosen cache level \(L_{\text{cap}}\) at \(p_c = R\). 6 requires only the two machine constants \((R, L_{\text{cap}})\) and the dataset’s \(n w\), with no \(\alpha\)-\(\beta\)-\(\gamma\) calibration. On Perlmutter CPU \(R = 64\) and \(L_{\text{cap}} = 1\) MB (L2 per core on AMD EPYC 7763). The cache term is non-binding on every LIBSVM dataset we measure (\(n w \leq R \cdot L_{\text{cap}} = 64\) MB), so the rule reduces to \(p_c^{*} = \min(R, p)\) in this benchmark suite. 4 confirms that 6 predicts the empirical winner on news20, rcv1, and uniform-density synthetic data (\(m{=}2^{21}, n{=}3.15\text{M}, \rho{=}0.004, \kappa{=}1\)) used as a controlled per-mesh sweep at \(p{=}128\) to isolate the topology effect from column skew. On url the rule’s prediction is the immediate-neighbor mesh of the empirical winner and ties within \(9\%\) on per-iteration runtime. The cost model [eq:model] is then used at the selected mesh to rank candidate \(s, b, \tau\) settings and locate the operating regime (5); we use it to order configurations rather than to predict absolute wall-clock time, for which it is accurate only up to a constant per-call overhead (6.5).

Within the FedAvg convergence regime where Stich’s bound \(\tau = O(\sqrt{\tilde{K}/(bp)})\) holds, direct measurement at \(\eta \in [0.005, 0.02]\) shows FedAvg and HybridSGD converge to within \(5\%\) of the same loss given enough iterations (7.5), so time-to-accuracy speedups arise from the winning mesh and partitioner reducing per-iteration runtime by \(1\) to \(2\) orders of magnitude on column-skewed data. Outside that regime (large \(p\) with \(\tau b p\) exceeding Stich’s bound for the iteration budget), FedAvg’s local-update drift inflates iterations-to-loss and HybridSGD compounds its per-iteration win with a sample-efficiency win, as observed on url at \(p{=}256\) (7.5). Load imbalance (\(\kappa\), 6.5) further biases the optimum toward cache-friendly partitioners.

6.4 Regime analysis↩︎

5 summarizes the four operating regimes exposed by [eq:model].

In the latency-bound regime \(s^*b^*\) grows as \((\alpha/\beta)^{1/2}\) and the topology rule 6 pins the mesh away from both 1D corners. The interior \(p_r\) predicted by the rule lands in \([2, 16]\) for LIBSVM datasets at moderate \(p\). In the Gram-BW-bound regime \(s\) shrinks toward \(1\) and the model reduces to FedAvg. In the sync-BW-bound regime, increasing \(p_c\) distributes the weight across more ranks and cuts each column-Allreduce payload from \(n\) to \(n/p_c\). This is the principal communication benefit of \(p_c > 1\) over pure FedAvg. The CA overhead of \(2sb\) extra FLOPs/sample is beneficial when \(\alpha\log p_c / \gamma > s^2 b^2\). On Perlmutter \(\alpha/\gamma \approx 10^6\) to \(10^8\) (7), so the inequality holds for all \(s \leq 32\), \(b \leq 64\), \(p_c \geq 2\).

6.5 Empirical refinements↩︎

The leading-order \((\alpha, \beta, \gamma)\) model of 3 captures regime direction correctly on every LIBSVM dataset, but quantitatively over/underestimates by \(2\) to \(10\times\). We present several empirical refinements to the model, which address irregularities in hardware and data.

6.5.0.1 Cache-aware compute

The \((4\bar{z} + 2n/b)\gamma\) compute term assumes the weight vector \(\boldsymbol{x} \in \mathbb{R}^n\) is fetched from DRAM at every mini-batch. In practice the \(\tau\) inner SGD steps between Allreduces share a single weight vector that stays cache-resident after the first access. We replace the per-outer-iter weight-access cost with \(T_{\text{weights}} = nw[\gamma_{\text{DRAM}} + (\tau{-}1)\gamma_{\text{cache}}(nw)]\), where \(\gamma_{\text{cache}}\) is a step function over the cache hierarchy. Empirical values from cblas_ddot on Perlmutter EPYC 7763 are \(\gamma_{L1} \approx 4{\times}10^{-12}\), \(\gamma_{L2} \approx 1.25{\times}10^{-11}\), \(\gamma_{L3} \approx 1.5{\times}10^{-11}\), and \(\gamma_{\text{DRAM}} \approx 2.6{\times}10^{-11}\) s/byte. HybridSGD’s local weights \(\boldsymbol{x} \in \mathbb{R}^{n/p_c}\) may fall in a tighter cache level than FedAvg’s full-\(n\) vector, giving HybridSGD a cache-locality advantage on large-\(n\) datasets.

6.5.0.2 Rank-aware \(\beta\) and load imbalance

\(\beta(p)\) on Perlmutter spans \(5.3{\times}10^{-11}\) (intra-node \(p{=}1\), shared-memory) to \(6.6{\times}10^{-9}\) s/B (\(p{=}16384\), inter-node Slingshot). This is a \(125\times\) swing dominated by shared-memory bandwidth contention below \(p \leq 64\) and NIC saturation above. Substituting the rank-appropriate \(\beta(p)\) improves single-node predictions \(2\) to \(5\times\). Load imbalance from heavy-tailed nnz distributions multiplies the sparse-compute term by \(\kappa = \max_p(\text{nnz})/\overline{\text{nnz}}\). Measured \(\kappa\) ranges from \(1.0\) (epsilon) through \(1.9\) (news20 at \(p{=}64\), 1D row) to \(\mathbf{482}\) (url at \(p_r{=}4,\,p_c{=}1024\), 2D row+col with column skew dominating).

6.5.0.3 Cache-aware partitioning

Given sparse matrix inputs, the simplest approach for SGD-based methods would be to simply balance nnzs to ensure equal workload in the SpMV computations. While this is sufficient on moderate-skew data (news20, where it yields a measured \(2.3\times\) HybridSGD speedup), but can be harmful on extreme column-skew data (i.e., heavy-tailed distributions). Measurement on url at \(p_c{=}64\) shows three regimes. Row partitioner gives \(n_{\text{local}}{=}n/p_c{=}50{,}499\) cols/rank (cache-friendly, fits L2) with \(\kappa{=}33.8\) (nnz-imbalanced). The nnz-balanced partitioner gives \(n_{\text{local}} \in [1, 1{,}409{,}992]\) per rank, where \(\kappa{=}1.3\) on nnz but the rank holding \(1.4\)M columns has \(11.2\) MB of weights, spilling out of L2 (\(1\) MB/core) into L3 or DRAM and degrading per-iteration runtime by \(2.4\times\). Cyclic partitioner gives \(n_{\text{local}}{=}n/p_c\) exactly with \(\kappa{=}1.9\) (near-optimal) at the cost of a column permutation in the reader. The right partitioner depends on the column-skew structure and the target cache hierarchy. It is a two-objective constrained problem, \(\min_P \kappa(P)\) s.t.\(\max_p n_{\text{local}}^p(P)\,w \leq L_{\text{cap}}\). Naive nnz-balanced partitioning satisfies neither. Cyclic partitioning satisfies both but pays in read complexity. The cost model predicts, and measurement confirms (7.3), that no single partitioner dominates across the LIBSVM suite.

6.5.0.4 Sync-skew term

Per-phase timing on url HybridSGD at \(p_r{=}4, p_c{=}64\) across the three partitioners (10) shows that the dominant cost of poor partitioning manifests as sync-skew waiting time inside the row-team Allreduce, not as compute time on the slowest rank. The \(s\)-step SGD comm timer (row Allreduce of the \(sb\) residual) measures \(477\,\mu\)s under rows partitioner (\(\kappa{=}34\)) versus \(142\,\mu\)s under cyclic partitioner (\(\kappa{=}1.9\)). The \(\sim\)​335 \(\mu\)s gap is wait-for-slowest-rank time, not MPI bandwidth or latency cost, since the payload is \(\sim\)​1 KB in both cases. We capture this with \(T_{\text{sync skew}} \approx (\kappa_{\text{local}} - 1) \cdot T_{\text{compute,avg}}\), applied to the row Allreduce. The term is zero for \(\kappa{=}1\) and grows linearly with imbalance.

6.5.0.5 Validation

We validate the model against the property we actually rely on — ranking fidelity, i.e.whether it orders configurations correctly — and treat absolute-runtime accuracy as a secondary, weaker check. The model gets the regime direction and partitioner ranking correct on all \(9\) (dataset, partitioner) cells measured (7). On url and news20 the predicted ranking is cyclic \(<\) rows \(<\) nnz (cache spill on the latter), matching observation, and on rcv1 all three are tied within \(5\%\) both predicted and measured. This ranking fidelity is the property required to use the formalism for mesh and partitioner selection on new datasets. Absolute runtime is predicted less accurately: the combined refinements match rcv1 within \(1\%\) across all three partitioners (predicted/measured ratio \(0.94\) to \(0.99\)), but the ratio falls to \(0.77\) on news20 cyclic and \(0.34\) on url cyclic, with the gap growing in \(n/p_c\). At url’s \(n/p_c{=}50{,}499\) columns per rank and \(\bar{z}/p_c{\approx}2\) nonzeros per row, the per-call overhead of MKL’s mkl_sparse_syrkd inspector (which scans the full column-index array) and the transpose SpMV scatter into the \(n/p_c\)-length gradient vector dominate actual floating-point work. Standalone microbenchmarking confirms that syrkd has a \({\sim}10\,\mu\)s floor at \(n/p_c{=}50\)K regardless of nonzero count, and the transpose SpMV adds \({\sim}9\,\mu\)s, both scaling linearly with \(n/p_c\). On rcv1 (\(n/p_c{=}2{,}952\)) these overheads drop below \(3\,\mu\)s and the flop-based \(\gamma\) term is accurate. The model does not capture this per-call overhead because it is proportional to the column dimension, not to the flop count; a piecewise \(\max(\text{flop cost},\, c \cdot n/p_c)\) compute term would close the gap but requires two additional calibration constants. Crucially, because we use the model only to rank configurations, this absolute-runtime gap does not affect the selection decisions it drives.

7 Experiments↩︎

We evaluate on the four LIBSVM datasets summarized in 6. We implement the SGD variants of 4 in C++ with MPI [35] (Cray MPICH 9.0.1) and Intel oneAPI MKL [36] (2025.3) for dense and sparse BLAS, in particular mkl_sparse_d_mv for SpMV and mkl_sparse_syrkd for the \(s\)-step Gram computation. All experiments use the NERSC Cray EX (Perlmutter) CPU partition [37] at \(64\) MPI ranks per node \(\times\) \(2\) cores per task (one rank per physical core, no SMT). We performed offline experiments to determine this setting. We also performed experiments using a mixed OpenMP+MPI hybrid parallelism model, but we did not observe improvements over a flat MPI model. So we use flat-MPI throughout our experiments. \(\boldsymbol{A}\) is stored in three-array CSR format and a new MKL sparse handle is created per batch of \(b\) rows. All experiments use FP64 to match the \(\bar{z}\)-dependent conditioning of the \(s\)-step Gram matrix, which was unstable at FP32 on news20 (\(\bar{z}{=}455\)). Memory is aligned to \(64\)-byte boundaries and experiments are timed using std::chrono.

7.1 Measured \(\alpha\), \(\beta\), and \(\gamma\)↩︎

The cost model in 6 is parameterized by \(\alpha\) (s/message), \(\beta\) (s/byte), and \(\gamma\) (s/byte). We measure these quantities empirically on Perlmutter CPU nodes using microbenchmarks (2\(\times\)AMD EPYC 7763, Slingshot-11, \(64\) ranks/node, one rank per physical core). \(\alpha\) and \(\beta\) come from MPI_Allreduce sweeps at \(N \in \{1,\dots,256\}\) nodes (inter-node) and at \(p \in \{1,\dots,64\}\) ranks single-node (intra-node), fitting \(T = 2\lceil\log_2 p\rceil\alpha + W\beta\) over \(W \in \{2^{10},\dots,2^{26}\}\)  Bytes. \(\gamma\) is measured via single-thread cblas_ddot with increasing input sizes to exercise cache effects. The cache-aware refinement (6.5) used the parametrized function \(\gamma(W)\) to capture the varying nature of performance due to caching. The rank-aware refinement selects \(\beta(p)\) based on whether the Allreduce is intra- or inter-node.

7.2 Mesh tuning and crossover↩︎

6 predicts the per-iteration optimum at mesh dimension \(p_c^{*} = \min(R, p)\), \(p_r^{*} = p/p_c^{*}\), with \(R{=}64\) on Perlmutter CPU. At \(p{=}256\) this gives \(p_r^{*}{=}4, p_c^{*}{=}64\). We validate against the full mesh sweep in 8, which enumerates all nine factorizations \(p_r p_c = 256\) with \(b{=}32\), \(s{=}4\), \(\tau{=}10\) at fixed iteration count and reports per-iteration runtime. The per-iteration minimum sits at \((p_r{=}8, p_c{=}32)\), the immediate-neighbor mesh of the rule’s prediction. The rule’s \((p_r{=}4, p_c{=}64)\) is within \(9\%\) of the empirical optimum on per-iteration runtime and is the time-to-target winner (11) because fewer averaging groups (\(p_r{=}4\) versus \(8\)) reduce local-update drift. The sharp climb in per-iteration cost once \(p_c\) falls below \(32\) is consistent with the row Allreduce crossing the per-node rank boundary at \(R{=}64\), after which the inter-node \(\beta\) step inflates the Gram-BW term.

8 reports per-iteration runtime at each dataset’s empirically best HybridSGD mesh alongside FedAvg. HybridSGD’s per-iteration advantage emerges only on url, where the FedAvg Allreduce of the full \(n{=}3.2\text{M}\) weight vector dominates per-iteration cost. On news20, rcv1, and epsilon FedAvg’s cheaper per-iteration compute (\(4\bar z\gamma\) versus HybridSGD’s \(6\bar z\gamma\)) wins. As \(\bar z\) grows (epsilon, news20) or \(n\) shrinks (rcv1), the topology rule 6 pushes \(p_c\) toward \(p\) (the 1D FedAvg corner). HybridSGD’s per-iteration disadvantage on compute-bound datasets does not preclude time-to-target wins. HybridSGD achieves \(53\times\) on url and \(14.6\times\) on news20 (11) because the smaller column-Allreduce payload (\(n/p_c\) versus \(n\)) and reduced local-update drift compound over many iterations.

7.3 Irregularity-aware partitioning↩︎

6.5 frames partitioning as a two-objective problem that minimizes \(\kappa\) subject to a per-rank weight partitioning that fits in fast memory. We implement three partitioners in our experiments (5 illustrates the resulting column-to-rank assignments on a small skewed matrix). Rows uses uniform contiguous \(n/p_c\) columns per rank, which is cache-friendly but nnz-imbalanced on skewed data. Nnz is a contiguous greedy partitioner that walks columns in order and advances to the next rank once its cumulative nnz reaches the target \(m\bar z/p\), achieving \(\kappa{\approx}1\) at the cost of storing heavy columns on a single rank. On heavy-tailed distributions, this partitioner can produce a weight partitioning that does not fit in cache, which leads to a performance imbalance. Cyclic uses round-robin column assignment, giving perfect \(n_{\text{local}}{=}n/p_c\) with near-optimal \(\kappa\) in expectation.

9 reports \(\kappa\), max-rank \(n_{\text{local}}\), and the per-iteration HybridSGD runtime for each partitioner at each dataset’s best strong-scaling configuration. The url dataset at \(p_c{=}64\) is the worst case. The nnz partitioner gives one rank \(1.4{\times}10^6\) columns (\(\sim\)​11 MB of weights), which spills out of L2 (\(1\) MB/core) and L3 (\(\sim\)​512 KB/core) into DRAM. This leads to a \(2.4\times\) per-iteration penalty over the rows partitioner on url (9). Cyclic partitioning is the consistent winner on heavy-tailed data (\(1.86\times\) over rows on url, \(3.5\times\) on news20, ties on rcv1). The ordering cyclic \(<\) rows \(<\) nnz on url contradicts the assumption that balancing nnzs alone leads to better performance. In contrast, our cost model predicts that the right partitioner solves both objectives simultaneously.

A controlled synthetic skew sweep (6) corroborates these regime claims by showing cyclic partitioner is invariant to the column-skew exponent while rows partitioner degrades smoothly and nnz partitioner stays competitive at small \(n\) but spills cache at large \(n\).

7.3.0.1 Runtime breakdown on url

10 reports a timing breakdown for url HybridSGD \(4{\times}64\) for the three partitioners to illustrate the performance variation. The \(s\)-step SGD comm timer grows from \(142\,\mu\)s under cyclic to \(477\,\mu\)s under rows (\(\kappa{=}34\)) to \(1{,}905\,\mu\)s under nnz (cache spill). The increase in the comm timer is not due to the Allreduce communication bandwidth or latency, but rather idle time waiting for the slowest rank to synchronize. The linear increase in the comm timer qualitatively matches the linear increase in \(\kappa\) from cyclic to rows to nnz partitioner.

The refined predictor (6.5) gets the partitioner ranking correct on all \(9\) (dataset, partitioner) cells we measured and is within \(1\%\) on rcv1, \(30\%\) on news20 cyclic, and \(17\) to \(60\%\) on url. The url residual is dataset-specific and partitioner-independent (\(\approx 200\,\mu\)s/iter) and stems from MKL mkl_sparse_syrkd per-call overhead that the bandwidth-bound model does not capture. 7 visualizes the predictor against measurement. The rcv1 dataset sits on the diagonal, news20 cyclic and news20 nnz lie inside the \(0.5\) to \(2\times\) band, and the outliers (all three url cells and news20 rows) reflect the same MKL sparse_syrkd overhead. Qualitatively, our predictor’s partitioner ranking matches empirical measurements on all \(9\) dataset and partitioner combinations. Based on these results, cyclic partitioning is best for column-skewed sparse data, followed by nnz partitioning for balanced sparse data.

7.4 Solver-family transition↩︎

8 sweeps \(p_r\) across all factorizations of \(p\) using the cyclic partitioner for each dataset. We trace the full continuum from 1D \(s\)-step SGD (\(p_r{=}1\)) through interior HybridSGD meshes to FedAvg (\(p_r{=}p\)). The url panel exhibits a U-shaped per-iteration profile with an empirical minimum at \(p_r{=}8\) (\(0.557\) ms/iter). The topology rule 6 predicts \(p_r{=}4, p_c{=}64\) (\(0.606\) ms/iter), the immediate-neighbor mesh and within \(9\%\) of the empirical optimum. On news20 (\(p{=}64\)) and rcv1 (\(p{=}16\)) the rule’s \(p_c{=}\min(R,p)\) saturates at \(p\), so it predicts the 1D \(s\)-step SGD corner (\(p_r{=}1\)). The empirical curves are monotone with the minimum at the same corner on both panels.

7.5 Time-to-target loss↩︎

We perform offline tuning on the learning rate and set it to \(\eta{=}0.01\), which yields convergence on all solvers. We report single-trial times with inter-trial variance \({\le}\,4\%\) (see 11) and pick each solver’s best \(p\). For HybridSGD, we also pick the best mesh and partitioner. Per-dataset target losses (11) are calibrated to the slower solver’s terminal loss within the iteration budget. This corresponds to FedAvg on url (\(0.50\)) and news20 (\(0.45\)). All solvers reach the same loss on rcv1 (\(0.40\)). Epsilon’s loss decays slowly within the iteration budget, so we chose a mid-range value (\(0.65\)) as the target. This dataset stresses the convergence rate of the solver. 9 shows the convergence behavior of HybridSGD, 1D \(s\)-step SGD, and FedAvg for each solver’s best configuration. We show the convergence behavior on the url, news20, and rcv1 datasets. HybridSGD is fastest to the target on url with speedups of \(2.8\times\) over 1D \(s\)-step SGD and \(53\times\) over FedAvg. The speedup over FedAvg arises from two factors. First, HybridSGD reduces the weights vector Allreduce bandwidth cost to \(n/p_c = 101{,}000\) words vs FedAvg’s \(n = 3{,}231{,}961\)). Second, HybridSGD converges faster since it uses \(p_r\) processors instead of \(p\) processors in the row dimension (see 1 for convergence rate difference). On rcv1, ties emerge from near-identical trajectories.

7.5.0.1 Solution quality and the convex objective

1 is convex, so every solver minimizes the same objective and, at a comparable step size and a sufficient iteration budget, reaches the same loss (9). The speedups we report are therefore time-to-equal-loss, not accuracy trade-offs. The winning solver reaches a fixed objective value sooner, it does not converge to a different solution.

10 (left) shows per-iteration speedup on url as a function of \(p\) for FedAvg, HybridSGD \(1{\times}p\), and HybridSGD \(8{\times}(p/8)\) using the cyclic partitioner. FedAvg’s per-iteration time is flat at \({\approx}1\times\) since increasing \(p\) also increases work per rank. Setting \(\tau{=}10\) keeps the AllReduce overhead below \(1\%\) of per-iteration cost, so the communication bottleneck does not appear in the FedAvg experiment. HybridSGD \(1{\times}p\) is similarly flat near \(1\times\). This bottleneck is due to url’s extreme column skew and is not a communication bottleneck from Gram Allreduce size. HybridSGD \(8{\times}(p/8)\) reaches \(5.7\times\) at \(p{=}1024\). This is achieved through the use of a smaller row team (\(p_c{=}p/8\)) which shrinks the weights vector and the Gram Allreduce size while keeping the row team intra-node through \(p{=}512\).

10 (right) shows per-iteration speedup on a synthetic uniform nnz distributed sparse matrix. We observe speedups from 1D \(s\)-step SGD when column skew is removed. HybridSGD continues to show greater per-iteration speedup than the other solvers due to the additional mesh dimension, which allows for finding a better Allreduce bandwidth balance point than the extremes of FedAvg and 1D \(s\)-step SGD.

In both experiments, we observe superlinear speedups at \(p = \{128, 256\}\) due to caching effects, where cache size grows while local problem size decreases. This superlinear speedup disappears once communication becomes a larger fraction of runtime.

8 Conclusion↩︎

HybridSGD is a 2D-parallel SGD method that generalizes 1D \(s\)-step SGD and 1D FedAvg into a continuous family indexed by the mesh split \(p = p_r \times p_c\), recovering each 1D baseline at an extreme grid dimension and exposing the interior meshes as a productive design space. The two 1D approaches were previously seen as incompatible because each fixes a single data partitioning axis. We show that the partitionings are in fact compatible because they live on orthogonal axes of the 2D mesh. This compatibility lets a single solver navigate the trade-off between the recurrence-unrolling regime of \(s\)-step SGD and the deferred-averaging regime of FedAvg by varying \((p_r, p_c)\). The framework is implemented in C++/MPI on Intel MKL sparse BLAS with three selectable column partitioners (rows, nonzero-greedy, cyclic) and a configurable mesh, exposing the full \((p_r, p_c, s, b, \tau, \text{partitioner})\) design space for evaluation.

We present a closed-form \(\alpha\)-\(\beta\)-\(\gamma\) cost model that decomposes the per-epoch wall into compute, latency, Gram-bandwidth, and sync-bandwidth contributions, with optima for the recurrence length \(s\) and mini-batch size \(b\) and a bandwidth balance condition \((s{-}1)sb^2\tau p_c \approx 2n\) that distinguishes the Gram-bandwidth and sync-bandwidth regimes. We also present a cache-aware compute parameter \(\gamma(W)\), a rank-aware bandwidth parameter \(\beta(p)\) that captures the intra-to-inter-node transition, and a nonzero-imbalance multiplier \(\kappa\) to capture load-imbalance and characterize slow ranks. We present a parameter-free topology rule, \(p_c^{*} = \max(\lceil nw/L_{\text{cap}}\rceil,\,\min(R, p))\), which sets the mesh split from the per-node rank count \(R\) and the per-core cache capacity \(L_{\text{cap}}\) alone. The rule reduces mesh selection to two machine constants and removes the transcendental fixed-point that the analytic optimum would otherwise require.

Cache-aware partitioning is a two-objective constrained problem. We minimize the nonzero imbalance \(\kappa\) subject to a per-rank cache-footprint capacity. A nonzero-balanced contiguous partitioner achieves \(\kappa \approx 1\) at the cost of overloading the rank that absorbs the heavy columns, which spills out of L2 into DRAM and degrades per-iteration runtime by up to \(2.4\times\) on column-skewed data. A cyclic partitioner satisfies both objectives in expectation, bounding \(n_{\text{local}}\) exactly to \(n/p_c\) while keeping \(\kappa\) near \(1\), at the cost of a column permutation in the reader. The cost model predicts and measurement confirms that no single partitioner dominates across the LIBSVM suite. The appropriate choice is dataset-conditional and is a first-class HybridSGD parameter.

At the strong-scaling limit on LIBSVM, HybridSGD attains \(\mathbf{53\times}\) time-to-target-loss speedup on url and \(\mathbf{14.6\times}\) on news20 over FedAvg, matches FedAvg on rcv1, and is outperformed by FedAvg on dense epsilon where the cheaper per-iteration compute dominates. The qualitative regime crossover at each dataset is predicted by the cost model and confirmed empirically. These results position HybridSGD as a regime-aware solver, most beneficial in sparse, high-dimensional, communication- and skew-limited settings, rather than a universal replacement for FedAvg. The speedups decompose into a per-iteration throughput component (driven by the smaller column-Allreduce payload of \(n/p_c\) versus the full-\(n\) payload of FedAvg) and a sample-efficiency component (driven by the smaller effective per-sync batch satisfying Stich’s bound at large \(p\), where FedAvg’s local-update drift inflates iterations-to-loss). The framework extends directly to other convex losses for which \(s\)-step methods apply, with the cost model and topology rule transferring unchanged.

Acknowledgements↩︎

Thanks to Grey Ballard for numerous, helpful discussions on this work. This work was supported by the U.S. Department of Energy, Office of Science, Advanced Scientific Computing Research (ASCR) program under Award Number DE-SC-0023296. This research used resources of the National Energy Research Scientific Computing Center (NERSC), a Department of Energy Office of Science User Facility using NERSC award ASCR-ERCAP0026261 and ASCR-ERCAP0028617.

References↩︎