SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Coordinating Distributed Example Orders for Provably Accelerated Training
A. Feder Cooper
Wentao Guo
Khiem Pham*
Tiancheng Yuan
Charlie F. Ruan
Yucheng Lu
{afc78, wg247, dkp45, ty373, cdr54, y12967, cmd353}@cornell. edu
Recent research on online Gradient Balancing (GraB) has revealed that there exist permutation-based example orderings for SGD that are guaranteed to outperform random reshuffling (RR). Whereas RR arbitrarily permutes training examples, GraB leverages stale gradients from prior epochs to order examples — achieving a provably faster convergence rate than RR. However, GraB is limited by design: while it demonstrates an impressive ability to scale-up training on centralized data, it does not naturally extend to modern distributed ML workloads. We therefore propose Coordinated Distributed GraB (CD-GraB), which uses insights from prior work on kernel thinning to translate the benefits of provably faster permutation-based example ordering to distributed settings. With negligible overhead, CD-GraB exhibits a linear speedup in convergence rate over centralized GraB and outperforms distributed RR on a variety of benchmark tasks.
Random reshuffling, which samples training-data examples without replacement, has become the de facto example-ordering method in modern deep-learning libraries [34], given that it tends to accelerate optimizer convergence in practice. However, some recent theoretical work has identified cases in which random reshuffling can lead to data orderings that have a poor effect on convergence [8][36][49]. This has encouraged a line of research to investigate if there exist provably better permutation-based orderings that afford greater scalability in training [22][24][31]. Notably, Lu et al. [24] connects permuted-order SGD to the herding problem [16], and proposes the herding-based online Gradient Balancing algorithm (GraB), which converges provably faster than random reshuffling, and does so with little memory or computational overhead. In fact, in follow-on work, Cha et al. [6] proves that GraB is optimal: in theory, GraB is the fastest possible permutation-based example ordering algorithm.
These results are very exciting, suggesting that GraB should unseat random reshuffling as the example ordering method-of-choice for SGD; however, they only hold with respect to a single machine. GraB is optimal in settings with centralized data, but does not naturally translate to problems of modern-ML scale, which demand that training workloads be distributed across multiple parallel workers that each only have access to a subset of the training data. This drawback raises an important question:
Can we simultaneously achieve the scalability benefits of distributed training and provably faster permutation-based example ordering for SGD — both in theory and in practice?
In this work, we show that it is indeed possible to attain these twin objectives. To do so, we suggest the online Coordinated Distributed Gradient Balance algorithm (CD-GraB), which leverages insights
from kernel thinning to elevate the herding framework of centralized GraB (GraB) to the parallel setting. Felicitously, as a side effect, this choice of formulation brings about positive practical performance benefits (that can also improve the empirical behavior of centralized GraB). Using the exact same assumptions as the original GraB paper, we show analytically that coordinating example orders across parallel workers leads a linear speedup in convergence rate. For  $T$  epochs and  $m$  parallel workers, each with access to  $n$  examples, CD-GraB's convergence rate is  $\tilde{O}((mnT)^{-2/3})$  on smooth, non-convex objectives and  $\tilde{O}((mnT)^{-2})$  under the Polyak-Lojasiewicz (P.L.) condition $^2$
We run a series of experiments to verify these improvements in practice, implementing CD-GraB on a single node that distributes computation across multiple GPUs. We also run an ablation study in order to disentangle the benefits of parallelism from the positive side effects of using kernel thinning to formulate the CD-GraB algorithm. Similar to how centralized GraB demonstrates improved generalization over centralized random reshuffling (RR), we observe that CD-GraB exhibits improved generalization over distributed random reshuffling (D-RR). Altogether, the success of our work suggests a new distributed training paradigm to explore in future work, which we call the Order Server (Section 6). In summary, we:
- Propose the online Coordinated Distributed Gradient Balancing (CD-GraB) algorithm, which enables provably accelerated training using SGD in the parallel setting (Section3);
- Prove that the convergence rate for CD-GraB exhibits a linear speedup over GraB, using the exact same assumptions as the original GraB paper (Section 4);
- Produce extensive empirical validation of CD-GraB's improved scalability on a variety of tasks in deep learning and on large-scale logistic regression (Section5).
2 Preliminaries and Related Work
In this section, we discuss the preliminaries and prior scholarship on permutation-based example ordering, with particular attention paid to the centralized online Gradient Balancing Algorithm (GraB) [24]. This lays the groundwork for how our coordinated, distributed GraB algorithm (Section3) imparts the efficiency guarantees of GraB to the parallelized regime (Section4).
Ordering data examples during training. Training a model can be formulated as minimizing a differentiable loss function  $f: \mathbb{R}^d \to \mathbb{R}$  over  $N$  data examples. The goal of this minimization is to obtain the target model weights  $\boldsymbol{w}^* = \operatorname{argmin}_{\boldsymbol{w}} f(\boldsymbol{w})$ , where  $f(\boldsymbol{w}) = \frac{1}{N} \sum_{j=1}^{N} f(\boldsymbol{w}; j)$ , for which  $f(\boldsymbol{w}; j)$  denotes the loss incurred on the  $j$ -th example. A typical training process iteratively updates the model parameters  $\boldsymbol{w}$  by scanning over the  $N$  data examples repeatedly, with  $t$ -th scan (or epoch) following
$$
\boldsymbol {w} _ {t} ^ {j + 1} = \boldsymbol {w} _ {t} ^ {j} - \alpha \nabla f \left(\boldsymbol {w} _ {t} ^ {j}; \pi_ {t} (j)\right), \forall j \in [ N ], \tag {1}
$$
where  $\alpha$  denotes the learning rate, and  $\pi_t: [N] \to [N]$  denotes a permutation ordering adopted in the  $t$ -th epoch from which the examples are chosen to compute gradients,  $\boldsymbol{w}_t^1$  denotes the initial model weights for the  $t$ -th epoch, and  $\boldsymbol{w}_t^j$  denotes the model weights after  $j-1$  gradient updates in the  $t$ -th epoch.
The choice of ordering  $\pi$  can have a significant effect on optimizer performance. Two popular methods, which can demonstrate convergence speedups in practice, are 1) random reshuffling (RR) [46], for which the permutations are random and differ over epochs, and 2) Shuffle Once (SO) [413], for which a random permutation is computed once and remains fixed for all epochs. Recht and Ré [37] conducted the first theoretical investigation of RR, while subsequent works like Yun et al. [49] and De Sa [8] have given counterexamples in which RR leads to orderings that have a poor effect on convergence. Altogether, many studies indicate that RR and SO only provide efficiency benefits under certain conditions [14, 15, 30].
These limitations of RR and SO have motivated research to identify permutations that outperform random ones. Rajput et al. [36] introduces an RR variant that achieves improved convergence for quadratics by reversing the ordering every other epoch. Other non-RR-based methods pick efficient orderings based on correlations between adjacently selected examples. In a recent line of work, Lu et al.
In this paper, we use  $\tilde{O}$  by convention to hide logarithmic factors in the problem parameters.
3While without-replacement orderings are most common in large-scale learning [5], ordering strategies need not be permutations, e.g., with-replacement sampling [23,32,39] or curriculum learning [12,27,41].
4 Note that we write  $\boxed{1}$  in terms of per-example-  $j$  gradients.
[22] proves that faster convergence is possible for SGD when the averages of consecutive stochastic gradients converge faster to the full gradient. Based on this result, in follow-on work Lu et al. [24] proposes the centralized online Gradient Balancing algorithm (GraB), which outperforms RR, and upon which we base this work.
2.1 GraB: Optimal, online, permutation-based example ordering for centralized ML
GraB is a permutation-based example-ordering algorithm that identifies provably better-than-random orderings in centralized, single-node settings for SGD. GraB finds such orderings by leveraging information in stale stochastic gradients from previous epochs to guide ordering in the next epoch. More formally, for smooth, non-convex objectives, Lu et al. [24] proves that any permutation  $\pi^{*}$  that guarantees
$$
\left. \max  _ {k \in [ N ]} \left\| \sum_ {j = 1} ^ {k} \nabla f (\boldsymbol {w}; \pi^ {*} (j)) - \nabla f (\boldsymbol {w}) \right\| _ {\infty} = \tilde {O} (1) \quad (\nabla f (\boldsymbol {w}) \text {i s t h e a v e r a g e g r a i d e n t}), \right. \tag {2}
$$
will yield a convergence rate of  $\tilde{O}((NT)^{-2/3})$  (for epochs  $T$ ) for SGD, which is superior to the  $O(N^{-1/3}T^{-2/3})$  convergence rate of random reshuffling [30].
GraB's connection to herding and balancing. To find such a permutation  $\pi^{*}$ , Lu et al. [24] connect [2] to the herding problem and vector balancing [16, 45]. Understanding why GraB does not naturally extend to the distributed setting — and our main contributions (Sections 3 and 4) — requires some additional details on the fundamentals of herding:
Given  $N$  vectors  $\boxed{5}\{\pmb {z}_j\}_{j = 1}^N (\pmb {z}_j\in \mathbb{R}^d)$ $\| \pmb {z}_j\| _2\leq 1$ $(\forall j)$ , herding identifies a permutation  $\pi^{*}$  such that
$$
\left. \max  _ {k \in [ N ]} \left\| \sum_ {j = 1} ^ {k} \left(\boldsymbol {z} _ {\pi^ {*} (j)} - \bar {\boldsymbol {z}}\right) \right\| _ {\infty} = \tilde {O} (1), \quad \text {w h e r e} \bar {z} = \frac {1}{N} \sum_ {j = 1} ^ {N} \boldsymbol {z} _ {j}. \right. \tag {3}
$$
It is clear that 3 generalizes 2, which is a specific case of herding in an optimization setting.
Harvey and Samadi solve (3) with a method called balancing (16). Balancing uses a signed version of the herding problem to optimize any given permutation  $\pi$  to reduce the bound in (3). That is, balancing formulates the signed herding problem
$$
\left. \max  _ {k \in [ N ]} \left\| \sum_ {j = 1} ^ {k} s _ {\pi (j)} \left(\boldsymbol {z} _ {\pi (j)} - \bar {\boldsymbol {z}}\right) \right\| _ {\infty}, \quad \text {w h e r e} \{s _ {j} \} _ {j = 1} ^ {N} \in \{+ 1, - 1 \}. \right. \tag {4}
$$
Given a group of such signs  $\{s_j\}_{j=1}^N$  and an arbitrary permutation  $\pi$ , Harvey and Samadi prove that Algorithm produces a new permutation  $\pi'$  such that
$$
\max _ {k \in [ N ]} \left\| \sum_ {j = 1} ^ {k} \big (\pmb {z} _ {\pi^ {\prime} (j)} - \bar {\pmb {z}} \big) \right\| _ {\infty} \leq \frac {1}{2} \max _ {k \in [ N ]} \left\| \sum_ {j = 1} ^ {k} s _ {\pi (j)} \big (\pmb {z} _ {\pi (j)} - \bar {\pmb {z}} \big) \right\| _ {\infty} + \frac {1}{2} \max _ {k \in [ N ]} \left\| \sum_ {j = 1} ^ {k} \big (\pmb {z} _ {\pi (j)} - \bar {\pmb {z}} \big) \right\| _ {\infty}.
$$
This says that, with new permutation  $\pi^{\prime}$ , the objective of (3) now approaches the bound of (4). Importantly, recent advances show that it is quite cheap to find a group of signs, such that (4) is on the order of  $\tilde{O}(1)$  (e.g., Alweiss et al. [2], in Algorithm 2). We are therefore able to call Algorithm 1 repeatedly, which will eventually obtain the  $\pi^{*}$  that solves the  $\tilde{O}(1)$  herding objective in (3).
Algorithm 1 Reordering Vectors based on Balanced Signs [Harvey and Samadi [16]]
input: a group of signs  $\{s_j\}_{j = 1}^N$  , initial order  $\pi$
initialize: two order-sensitive lists  $L_{\mathrm{pos}} \leftarrow [, L_{\mathrm{neg}} \leftarrow []$
for  $j = 1\dots N$  do
$L_{\mathrm{pos}}$  .append  $(\pi (j))$  if  $s_j$  is  $+1$  else  $L_{\mathrm{neg}}$  .append  $(\pi (j))$
return: new order  $\pi^{\prime}\coloneqq \mathrm{concat}(L_{\mathrm{pos}},\mathrm{reverse}(L_{\mathrm{neg}}))$
GraB's application of herding to gradient balancing. Lu et al. [24] applies this framework of herding and balancing to develop GraB, i.e., to minimize [2]. The main challenge for the success of this approach is to find the right gradients  $z_{j}$  in the optimization context of [2]. Notably, the herding and balancing framework requires the vector mean  $\bar{z}$  in advance. To satisfy this requirement, GraB "centers" the gradient vectors using a stale mean. That is, GraB runs the herding algorithm on vectors that are defined as
$$
\boldsymbol {z} _ {j} = \nabla f \left(\boldsymbol {w} _ {t} ^ {j}; \pi_ {t} (j)\right) - \frac {1}{N} \sum_ {p = 1} ^ {N} \nabla f \left(\boldsymbol {w} _ {t - 1} ^ {p}; \pi_ {t - 1} (p)\right), \tag {5}
$$
where  $\pmb{w}_t^p$  denotes the model weights after  $p - 1$  updates in the  $t$ -th epoch, and  $\pi_t$  denotes the permutation adopted in the  $t$ -th epoch. Lu et al. [24] proves that this definition of  $z_j$  preserves the
5 Herding does not have an optimization context. Here,  $N$  does not refer to the number of data examples used in training [1]; rather,  $N \in \mathbb{Z}^{+}$  describes the size of a set of arbitrary vectors. We slightly abuse notation because we execute the herding subroutine on exactly  $N$  gradients (Section 3), which happen to equal the number of  $N$  examples.
benefits of balancing with negligible noise or overhead. The only overhead comes from storing the running average of the gradients in epoch  $t - 1$  to "center" the gradients in the subsequent epoch  $t$ .
With this approach, Lu et al. [24] proves that GraB demonstrates more efficient convergence than RR for SGD. Better still, Cha et al. [6] demonstrates that GraB is in fact the optimal permutation-based ordering method for SGD: it is not possible to produce a permutation-based ordering in the centralized setting that achieves a faster convergence rate for SGD.
Despite GraB's clear benefits over RR, it assumes local access to all examples. This assumption does not hold for popular, modern, parallel settings (e.g., parameter server [20]), in which workers only have access to subsets of examples. No present work has attempted to investigate GraB's applicability to this setting. While some work has studied distributed RR (D-RR) [18,26,38,48], it remains an open question if GraB's efficiency benefits for SGD can be conferred to the modern-scale, distributed-ML setup.
3 CD-GraB: A Provably Efficient Ordering Algorithm for Distributed Training
Our main contribution is to elevate GraB to the parallel regime, so that distributed training can enjoy the efficiency benefits of provably better example ordering. Based on the preliminaries, we can now explain why this is not a straightforward task: While GraB achieves the optimal convergence rate for SGD on centralized data, it does not naturally translate to a distributed setting (Section 3.1). Our key insights for resolving these problems are to reformulate the herding framework in Lu et al. [24] to work in parallel, and to leverage insights from kernel thinning [3,10,11] to derive the online PairBalance algorithm, which solves this parallelized herding objective (Section 3.2). Lastly, we present the fullstack CD-GraB algorithm that makes our solution work in practice (Section 3.3). The server implements online PairBalance, which coordinates gradient information from the distributed workers in training epoch  $t$  in order to determine a provably efficient example order for the next epoch  $t + 1$  (Section 4).
3.1 Issues with GraB in the distributed setting
To clarify the issues with distributing GraB, we first need to define the distributed training setup more precisely. We consider the standard data-parallel regime with  $m$  parallel workers, where each worker keeps a copy of the model weights  $\boldsymbol{w} \in \mathbb{R}^d$  and maintains  $n = N / m$  local examples [6] As in many data-parallel training applications such as geo-distributed model training [47], we assume the data examples cannot be shared or moved across workers. More formally, this setup can be expressed as
$$
\min  _ {\boldsymbol {w} \in \mathbb {R} ^ {d}} \left[ f (\boldsymbol {w}) = \frac {1}{m} \sum_ {i = 1} ^ {m} f ^ {i} (\boldsymbol {w}) \right] \quad \text {w i t h} \quad f ^ {i} (\boldsymbol {w}) = \frac {1}{n} \sum_ {j = 1} ^ {n} f ^ {i} (\boldsymbol {w}; j), \tag {6}
$$
where  $f^i(\boldsymbol{w}; j) : \mathbb{R}^d \to \mathbb{R}, j \in [n]$ , denotes the loss incurred on the  $j$ -th example on the  $i$ -th worker for model weights  $\boldsymbol{w}$ . We can now consider running (1) using this setup, for which each worker scans over their  $n$  local-data examples using (potentially) different permutations. We denote  $\pi_{t,i} : [n] \to [n]$  as the permutation-based ordering adopted on the  $i$ -th worker in the  $t$ -th training epoch. Adjusting (1) to accommodate the setup in (6), the update to the model can be summarized as
$$
\boldsymbol {w} _ {t} ^ {j + 1} = \boldsymbol {w} _ {t} ^ {j} - \frac {\alpha}{m} \sum_ {i = 1} ^ {m} \nabla f ^ {i} \left(\boldsymbol {w} _ {t} ^ {j}; \pi_ {t, i} (j)\right), \forall j \in [ n ]. \tag {7}
$$
That is, in epoch  $t$ , each worker  $i$  selects their respective, local  $j$ -th example according to  $\{\pi_{t,i}\}_{i=1}^{n}$  in order to compute stochastic gradients (Appendix).
Following this setup, Algorithm 1 no longer guarantees the  $\tilde{O}(1)$  bound to the herding problem (3), a bound that is valid only when all data examples can be permuted freely [16]. This constraint is fine for centralized GraB, but, in distributed training, workers only have access to a subset of examples. Distributed training requires that worker-specific permutations only involve the examples in their respective local subsets. Further, recall that GraB uses stale means to center gradients (5) in order to solve the herding objective. This, too, causes problems in distributed training. In practice, it is typical to employ larger learning rates  $\alpha$  for greater scalability [40]; larger  $\alpha$  increases the discrepancy between averaged gradients in adjacent epochs, which, in turn, would make GraB's use of stale means unreliable.
Without loss of generality, we assume the  $N$  examples are divided evenly among the  $m$  workers and  $n$  is even.
One such popular paradigm is federated learning [28] e.g.]. Federated learning typically involves highly imbalanced loads, heterogeneous data, partial user participation, and additional privacy-preserving mechanisms. These characteristics are orthogonal to what we consider here for example order. If we were to allow for such data organization, we would need to assume non-global communication per iteration or additional constraints on how global communication occurs. For CD-GraB, we focus on the regime of using parallelism to accelerate training.
3.2 Our efficient solution: parallel herding and pair balancing
To address the limitations presented in the prior section, which preclude the direct application of GraB to distributed training, we will need to 1) reformulate the herding problem to fit the parallel setting, and 2) redesign how to do gradient balancing, such that it both solves our new herding formulation and allows for reliability with higher learning rates. We now present our solution to both these problems; we introduce the parallel herding problem and the online PaIRBalance subroutine that solves it.
Parallel Herding. To extend herding to the parallel setting, consider the following setup: There are  $m$  workers, which each have local access to  $n$  vectors. Let  $\mathbf{z}_{i,j} \in \mathbb{R}^d$  denote the vector indexed by  $j$  on the  $i$ -th worker. Assuming  $\| \mathbf{z}_{i,j} \|_2 \leq 1$  ( $\forall i \in [m], \forall j \in [n]$ ), the goal of parallel herding is to find  $m$  permutations,  $\pi_1, \pi_2, \dots, \pi_m$  where  $\pi_i : [n] \to [n]$  ( $\forall i \in [m]$ ), so as to minimize:
$$
\left. \max  _ {k \in [ n ]} \left\| \sum_ {j = 1} ^ {k} \sum_ {i = 1} ^ {m} \left(\boldsymbol {z} _ {i, \pi_ {i} (j)} - \bar {\boldsymbol {z}}\right) \right\| _ {\infty}, \quad \text {w i t h} \quad \bar {\boldsymbol {z}} = \frac {1}{m n} \sum_ {i = 1} ^ {m} \sum_ {j = 1} ^ {n} \boldsymbol {z} _ {i, j}. \right. \tag {8}
$$
When directly comparing (8) with (3), it is clear that parallel herding differs in two notable ways from the original herding problem. First, each permutation  $\pi_i : [n] \to [n] (\forall i \in [m])$  only decides the ordering of the  $n$  vectors that are associated with worker  $i$ . Second, the prefix sum taken in the objective norm is accumulated over all the workers (the inner sum from  $i = 1 \dots m$ ). This formulation naturally captures the setting in a distributed environment: workers need to decide permutations collaboratively, and the worker-specific vectors are processed simultaneously rather than sequentially.
Given that this formulation fits the distributed setting, we next need to show that parallel herding does in fact address the limitations posed by centralized GraB: that it is possible recover the original  $\tilde{O}(1)$  herding bound, and that we can solve the issue of unreliable stale gradients (Section 3.1). The solution that we present in the remainder of this section is a new vector balancing subroutine: online PairBalance. To give an intuition, as its name suggests, online PairBalance leverages insights from kernel thinning to balance vector differences over vector pairs. This also eliminates the need to perform vector centering, and thus solves the stale mean problem.
Using kernel thinning to solve parallel herding. We call our solution to the parallel herding objective (8) pair balancing, which we derive from key insights in kernel thinning [3] [10] [11]. In particular, Dwivedi and Mackey show that it is possible to solve the herding objective in  $\tilde{O}(1)$  by only examining differences on pairs of examples [10]. They derive an algorithm that generalizes Alweiss et al. [2] subroutine in Algorithm [2], which solves herding in  $\tilde{O}(1)$  (Section [2]), and does so by operating only on vector-pair differences [8]. This comes with a very useful property: eliminating the requirement of knowing the maximum vector norm ahead of time and centering the vectors (i.e., making all the vectors sum to zero) in order to solve the herding problem. This is the key to solving the parallel herding objective (8) in  $\tilde{O}(1)$ , and elevating the benefits of GraB to a distributed setting.
Following Dwivedi and Mackey [10], we will balance over paired vectors, and will do so in an online fashion (Section 3.3). This eliminates GraB's requirement of using a stale mean to center gradient vectors (Section 2.1), but still minimizes the parallel herding objective to  $\tilde{O}(1)$ . We defer proving this result to Section 4 and first describe our concrete algorithm. Online PairBalance applies Algorithm 1 on the "flattened" and "paired" sequence of all of the workers' paired-difference gradients, i.e.,
$$
\boldsymbol {y} _ {n (k - 1) + i} = \boldsymbol {z} _ {i, 2 k - 1} - \boldsymbol {z} _ {i, 2 k}, \quad \forall k \in [ \frac {n}{2} ], \quad i = 1 \dots m.
$$
That is, we fit these ordered-paired differences  $\{\pmb{y}_i\}_{i=1}^{mn/2}$  into the herding and balancing framework (Algorithm 1): if sign  $s$  is associated with  $\pmb{y}_{n(k-1)+i}$ , then  $\pmb{z}_{i,2k-1}$  and  $\pmb{z}_{i,2k}$  receive  $s$  and  $-s$ , respectively.
3.3 The full-stack CD-GraB algorithm
Having solved the parallel herding problem with pair balancing, we now demonstrate how to bring everything together in an optimization context to coordinate distributed gradient balancing for distributed training. That is, we can now introduce our full-stack CD-GraB algorithm, which trains models in a distributed setting (Section 3.1) while efficiently ordering the examples by using PairBalance (Section 3.2, Algorithm 2) in an online manner.