Coordinating Distributed Example Orders for Provably Accelerated Training
Authors: A. Feder Cooper
Wentao Guo
Khiem Pham
Tiancheng Yuan
Charlie F. Ruan
Yucheng Lu
Christopher De Sa
Venue: NeurIPS 2023 poster (2023)
Keywords: permuted example ordering
distributed training
scalable training
herding
TL;DR: Provably faster convergence in distributing training via coordinated example ordering

Abstract: 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.

Coordinating Distributed Example Orders for Provably Accelerated Training
A. Feder Cooper
Wentao Gue
Khiem Pham
Tiancheng Yuan
Charlie Ruan
Yucheng Lu
Christopher De
Cornell University {afc78, vg247, dkp45, ty373, cfr54, y12967, cmd353)@cornell. edu
Abstract
Recent research on online Gradient Balancing (GraB) has revealed that there exist permutation-base 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 order examples - achieving a provably faster convergence rate than RR. However, GraB is limited by design: while it demonstrates impressive ability to scale-up training 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 distributed settings. With negligible overhead, CD-GraB exhibits a linear speedup in convergence rate over centralized GraB and outperforms distributed RR variety of benchmark tasks.
Introduction
Random reshuffling, which samples training-date 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 line of research to investigate if there exist provably better permutation-bassee orderings that afford greater scalability permuted-order SGD to the herding problem (16), and proposes the herding- online Gradient Balancing algorithm (GraB), which converges provably faster than random reshuffling, and does so with little memory 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-ban based example ordering algorithm.
training [22] 24 31). Notably, Lu et al. (24) connects
These results very exciting. suggesting that GraB should unseat random reshuffling the example ordering
for SGD; however, they only hold with respect to a single machine. GraB optimal in settings with centralized data, but does not naturally translate problems of modern-MI. scale, which demand that training workloads be distributed across multiple parallel workers that each only have access subset of the training data. This drawback raises important question:
Can we simaltaneously 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 Gradiant Balance algorithm (CD-GraB), which leverages insights
37th Conference on Neural Information Processing Systems (NeurIPS 2023).
"Equal contribution

---

from kernel thinning to elevate the herding framework of centralized GraB (GraB) to the parallel setting. Felicitously, as 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 same assumptions as the original GraB paper, we show analytically that coordinating example orders across parallel workers leads linear speedup in convergence rate. For T epochs and m. parallel workers, each with access to examples, CD-GraB's convergence rate is ( - on smooth, non-convex objectives and ((mm]) ²) under the
(P.L.) condition
We run a series of experiments 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 future work, which we call the Order Server (Section 60. In summary, we:
Propose the online Coordinated Distributed Gradient Balancing (CD-GraB) algorithm, which enables provably accelcrated training using SGD in the parallel setting (Section
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[]:
Produce extensive empirical validation of CD-GraB's 's improved scalability on variety of tasks in deep learning and on large-scale logistic regression 3)
2 Preliminaries and Related Work
In this section, we discuss the preliminaries and prior scholarship on permutation-basee 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 (Section imparts the efficiency guarantees of GraB to the parallelized regime (Section 4.
Ordering data examples during training. Training model can be formulated as minimizing differentiable loss function f: + lK over N data examples. The goal of this minimization is to obtain the target model weights w= = argminw (w), where f(w) = N Σ Ew:j(,; for which (wr;j) denotes the loss incurred on the j-th example. A typical training process iteratively updates the model parameters by scanning over the N data examples repeatedly, with 1-th scan (or epoch) following
w +1 = w
V (I VjE
(1)
where denotes the learning rate, and ,.. [N] + [N] denotes permutation ordering Fladppted in the i-th epoch from which the examples are chosen to compute gradients, wi denotes the initial model weights for the 1-th epoch, and w denotes the model weights after j - gradient updates in the t-th epoch
The choice of ordering can have significant effect 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) 44] [3), for which a random permutation is computed once and remains fixed for all epochs. Recht and Ré 377 conducted the first theoretical investigation of RR, while subsequent works like Yun et al. 49] and De Sa [ 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 [4 [5] 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-basee methods pick efficient orderings based on correlations between adjacently selected examples. In a recent line of work, Lu et al.
PI this paper. we use Ô by convention to hide logarithmic factors in the problem parameters.
'Wiile sethout-replacemeet not he permutations, e.g., with-replacemmet sampling (23 32 ky or curriculum learning 12 (1)
orderings are most common in large-scale scale learning [5). ordering strategies need
Note that we write in terms of per-example-i gradients.

---

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. 224 proposes the centralized online Gradient Balancing algorithm (GraB), which outperforms RR, and upon which we base this work.
22]
2.1 GraB: Optimal, online, permutation-baser example ordering for centralized ML
GraB is permutation- example-ordering algorithm that identifies provably better-than-random orderings centralized, single-node settings for SGD. GraB finds such orderings by leveraging infor- mation in stale stochastic gradients from previous epochs to guide ordering in the next epoch. More for- mally, for smooth, non non-convex objectives, Lu et al. [24] proves that any permutation π that guarantees
[ Σド= Vf(ww (1) (f))v((w) =0(1) (V is the average gradient),
(2)
will yield a convergence rate of Ö((NT) 2/33) (for epochs T) for SGD, which is superior to the (/VV1//T-2/3) convergence rate of random reshuffling 33)
GraB's connection to herding and balancing. To find such a permutation n², et al. 24 connect
2 to the herding problem and vector bulancing [16 45). Understanding why GraB does naturally
extend to the distributed setting - and our main contributions (Sections and requires some additional details on the fundamentals of herding:
Given N vectors
{ (2, E ERR). || - (Vj), herding identifies permutation क such that Σ (z) 2) =Ō where 2 = Σ 22111
maX
is clear that generalizes which specific of herding in optimization setting.
Harvey and Samadi 33 with method called balancing [66.. Balancing uses a signed version of
the herding problem to optimize any given permutation to reduce the bound in (3). That is, balancing formulates the signed herding problem
max Σ (m()) (= (zx()) =)
where { (+1,-1).
Given a group of such signs {s and an arbitrary permutation n, Harvey and Samadi prove that Algorithm produces a new permutation such that
maxx RE(N]
E)=(z= (z=10) =)
max RE[N]
S0n x) max kE[N]
Ej= (Z+()) -=) =)
This says that, with permutation n', 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 Õ(1) (e.g.. Alweiss et al. [2]. in Algorithm 2). We are therefore able to call Algorithm I repeat- edly, which will eventually obtain the π that solves the (1) herding objective in 3
Algorithm 1 Reordering Vectors based Balanced Signs [Harvey and Samadi [16]
input: group of signs initialize: two arder-sensitive lists Lyes < []. Lung []- for forj=1.... do
و initial order i
append (T (J)) if s, is +1 else Lueg appond( (j)).
end for
return: new order x' concat Lownaa( Lnez))
GraB's 's's application of herding to gradient balancing. et 2 applies this framework of herding and balancing to develop GraB, i.e., to minimize (2) The main challenge for the success of this approach find the right gradients zj in the optimization context of (2). Notably, the herding and balaneing framework requires the vector mean in advance. To satisfy this requirement, GraB "centers" the gra- dient vectors using a stale mean. That is, GraB runs the herding algorithm on vectors that are defined as
(w)/ ())) + E V (
(5)
where / denotes the model weights after p - updates in the t-th epoch, and me denotes the permutation adopted in the i-th epoch. Lu et al. 24 proves that this definition of 2, preserves the
"Herding does not have an optimization context. Here, does not refer to the number of data examples used in training [ rather, N E 2t describes the size of a set of arbitrary vectors. We slightly abuse notation because we ex- ecute the herding subroutine on exactly N gradients (Section which happen 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 to "center" the gradients in the subsequent epoch t.
With this approach, Lu et al. for SGD. Better still, Cha al. (6] demonstrates that GraB is in fact the optimal permutation-based ordering method for SGD: it is not possible to produce a permutation-bes -aased ordering in the centralized setting that achieves faster convergence rate for SGD.
[24] proves that GraB demonstrates more efficient convergence than RR
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 's applicability to this setting. While some work has studied distributed RR (D-RR) [8 26 38 48) it remains an open ques- tion if GraB's s efficiency benefits for SGD can be conferred to the modern-scale, distributed-Mi 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 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 to derive the online Pair algorithm, which solves this parallelized herding objective (Section Lastly, we present the full- stack CD-GraB algorithm that makes our solution work in practice (Section 3.3). The server implemetts online PairBalance, which coordinates gradient information from the distributed workers in training epoch in order to determine provably efficient example order for the next epoch [++ 1 (Section
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 copy of the model weights w E Rd and maintains n = N/m //m local examples 6 As in many data-parallel training applications such 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 [fw = 1, [w)] with ('(w)= [(
(6)
where (w::)) [m]. denotes the loss incurred on the j-th example on the i-th worker for model weights w. We can now consider running (I) using this setup, for which each worker scans over their n local-data examples using (potentially) different permutations. We denote : - [r] as the permutation-basse
ordering adopted on the i-th worker in the 1-th training epoch. Adjusting (T) to accommodate the setup in (6). the update to the model can be summarized as
w m a E. (.) Vj [m]]
(7)
m
That is, in epoch 1, each worker selects their respective, local j-th example according to ( in order compute stochastic gradients (Appendix).
Following this setup, A no longer guarantees the Ó(1) bound to the herding problem
a bound that is valid only when all data examples can be permuted [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 center gradients (5) in order to solve the herding objective. This, too, causes problems in distributed training. In practice, it typical to employ larger learning rates for greater scalability [40]; larger a increases the discrepancy between averaged gradients in adjacent epochs, which, in turn, would make GraB's 's use of stale means unreliable.
"Without loss of generality, we assume the N examples are divided evenly among the m workers and v is even. "One such popular paradigm federated learning (28 e.g.]. Federated learning typically involves highly imbalanced loads, heterogeneous data, partial user participation, and additional privacy- preseerii 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-globa communication per iteration or additional constraints on how global communication occurs. For CD- GraG, we focus the regime of using parallelism 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 do gradient balancing, such that 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
Parallel Herding, To extend herding to the parallel setting, consider the following setup: There are m workers, which each have local access to vectors. Let 21,5 E R. denote the vector indexed by j on the i-th worker. Assuming | ll2 S 1 ( mjj.jje [rz]), the goal of parallel herding is to find m permutations, Ti TTT2,., where n; [m] + [n]] (vie [m]]. so as to minimize:
max
Ej=E ()) 2)
with = 1 E -E=
(8)
When directly comparing the original herding problem. First, each permutation n; [n] - [na] (Vi E (ml) 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...m.) This formulation naturally captures the setting in a distributed environment: workers need to decide permutations collaboratively and the worker-specific vectors are processed simultancously rather than sequentially.
1 with (3), it is clear that parallel herding differs in two notable ways from
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 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, 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 the parallel herding objective (8) pair balancing, which we derive from key insights in kernel thinning [ 10 M.] In particular,
Dwivedi and Mackey show that it is possible to solve the herding objective in O(1) by only examining differences on pairs of examples [TO]. They derive an algorithm that generalizes Alweiss al. (2]
subroutine in Algorithm 2 which solves herding Õ(1) (Section 2 and does so by operating only on vector-pair differences This comes with 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 Ö(1) and elevating the benefits of GraB to a distributed setting.
Following Dwivedi and Mackey 100) we will balance over paired vectors, and will do so in an online fashion (Section 3.3) This climinates GraB's 's requirement of using stale mean to center gradient vectors (Section 2.1 but still minimizes the parallel herding objective to (1). We defer proving this result Section and first describe our concrete algorithm. Online PairBalance applies Algorithm on "llatteed and "paired" sequence of all of the workers' paired-difference gradients, i.e.,
yn(k-1)++ Zi,2k Zi,2k,
( 1=
mn/2 Ts we fit these aidered-paire differences (3/ }m/ into the herding and balancing framework (Al- gorithm TD: if sign is associated with (n then z,, 2 and 21,2% receive 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 I while efficiently ordering the examples by using Pair (Section 3.2, Algorithm in an online manner.
Dwivedi and Mackey minimize the maximum mean discrepancy (MMD) between a selected coreset and an empirical distribution. They develop a new self- balancing Hilbert walk on differences of pairs of examples to select exactly half of the dataset points, and solve coreset selection by iteratively halving the input vector sequence into balanced coresets then selecting and refining candidate coreset to minimize MMD with the input sequence.
5

---

Algorithm/ 2 PairBalance
The inputs, outputs and subroutine for this algorithm are seder-sensitive
gs
ga
Call Pair Balance
6. Call
input: current running sum r, paired vectors 21, #2
Algorithm
Data Permutations
g
92
gs
94
compute: 8.7 +- Randomizeedallnce((--e
returm: 8 (sign for 21
Send Gradients server
(sign for Z2).
r (updated running sum)
Watoer
Worker
Workert
Worker 2
Workett
Wotoor 2
X
X
gi
gi
X
Adapted from Alweiss al.
define subroutine: RandomizedBallnceen...
X
X2
92
g3
X
compute: p+ (-(r,e)
9
g
X
compute: s - +1 with probability p:
8 s6-1 with probability 1-p update: r p &C
X
X
g
9
Examples
2. Compute gradients
Reorder examples with
return:
ordered with
Fe-12
Figure I: Left: The PairBalance algorithm, which the server runs online. Right: CD-GraB running on one server (top) and two workers (bottom). The workers do not share data examples.
We describe CD-GraB at two levels of abstraction: a high-level illustration (Figure steps 1-7) and detailed pair of worker-server algorithm statements (Figure Since the workers only have access to subset of the training data, in parallel they compute local, per-example stochastic gradients and send them to the server. The server simultaneously calls PairBalance online, which coordinates information from all the workers' gradients (i.e., using adjacent example-specific gradients) to determine the next epoch's 's worker-specific permutations. In more detail:
In epoch f, (Figure ! step 1) the two workers have permutations TT,1 and T2,2, respectively. Each worker computes per-example gradients 9 (2; Algorithm 3 4), and sends them to the server (3; Algorithm 5).. The server we implement functions as parameter server [20]: It computes the average of The workers' per-example gradients (Algorithm 46), and sends it back to all workers (Algorithmm47) so that they can update their local models (Algorithm 3 6-7)) Simultaneously, server receives gradients (Algorithm 4 5)) it calls PairBalance (Algorithmp) 2 on adjacent vectors (4; Algorithm PairBalance produces signs to supply to the reordering algorithm (Algorithm(1) [[]. which, using the current worker permutations Tr produces the new per-worker permutations for the next epoch (5; Algorithm 4 14). In Figure I these correspond to T1+1.1 and 1 which the server then sends back to the respective workers (6; Algorithmm 15). Lastly, before the start of the next epoch, the workers reorder their examples according to the new permutations (7: Algorithm 99..
the
Convergence Analysis
We next demonstrate formally that our CD-GraB algorithm (Section 3.3) of centralized GraB (Section to the distributed setting. In brief, our main theoretical results show that CD-GraB enjoys linear speedup in convergence rate under two sets of conditions: smoothness (Theoremp) and the
confers the efficiency benefits
(P...) condition (Theorem 3 Both results guarantee that CD-GraB is faster than distributed random reshuffling (D-RR). Our proofs rely on Corollary 7 from (subroutine in Algorithm from Alweiss et al. (2)) guarantees a (1) bound to the signed herding objective (4
Dwivedi and Mackey [10), which shows that, with high probability, Randomized
To begin, we restate this result cohere with framework, for which the vectors zj are gradients in an optimization context:
Theorem 1 (Corollary 7, Dwivedi and Mackey [10) Consider any vectors { (z, E (Rd) with
|| < supplied as input to the RandomizedBalancee subroutine Algorithm 2 Then for any d 5>0
E with probability at least 1-8, 8, RandomizedBalance outputs a sequence of signs E (-1,1)
that satisfy maxke[N]
Σ 1
<, where À = 2log log( ANN =
"Corollary 7 from Dwivedi and Mackey
10
improves the result of Theorem 1.1 from Alweiss el. 2

---

Algorithm 3 CD-GraB Workers
Algorithm 4 CD-GraB Parameter Server
require: m workers. n:= ex. worker input: initial wi. epochs T, learning rate a
require: m workers, n=== N ex. per worker
input: epochs T
{=
1: receive: initial permutations for eurehhc I...7 do
1: send: initial permutations {m
for epoch t:= :etttt:..... 1...I do
D> Run in parallel for workers 1...m
3: initialize: running sum ht 0: empty list
3: for example j == do
4:
for example j 1...n 1..n. do
compute: gj + V ( (w.
j-th stochastic grad. 9]
send: g]
receive: { from all workers
6: compute: avg. gradient: g; 1 -i=1g]
receive:
9j
avg. j-th stochastic grad. g)
send: gj to all the workers
8:
for worker i= ::..m.mmd
7:
update: w/" +1 ww -ogs
9:
if j mod 2=0:
10:
h,s h Pairralancceh.ggaaaaaa S.append 8 )) S append (8))
end for
11:
12:
end for end for
13:
D Call Alg. 1 i I...m on ma, S
14:
compute: next permutations (T ) send: {Ir++1a) Mm to each worker
9:
receive: next permutation
15:
10:
update: wi. w
16:
end for
11:
end for
12:
return: WT+1 wh+1
Figure 2: CD-GraB worker and server (here, a parameter server (20) algorithms.
To integrate this result with our parallel setting, we need some additional assumptions that are standard in the literature distributed optimization that the variance of the per-example gradients each worker is uniformly bounded (Assumption [). and that the variance between worker-specific gradients similarly bounded (Assumption More precisely, following the distributed setup (7), denote the global loss gradient to be f (w), each i-th worker's ' local loss gradient to be V ()) (ViE [m]),
and each i-th worker's per-example loss gradients be V f (w:j) (Vj E [n]). We assume:
Assumption 1 (Bounded Gradient Variance). For all E [m] there exists constant o >> 0 such that for all j E [n] and all w E Rd, holds that (w:
Assumption 2 (Bounded Data Heterogeneity). There exists a constant > such that Vi E [m]] [V) V f 11?
Lastly, we include one additional assumption from the original GraB paper 24 assume cross
norm L2 (which can be easily adapted to Lg-smoothness by setting 22 to be vdLz).
Assumption 3 (Smoothness). There exists constant L2 such that any w. Rd any i E [m].
and any j E [n], it holds that V (w:j) Vf ||
o2 aolu-lles
Given these assumptions, we can prove convergence guarantee for CD-GraB:
Theorem 2. Suppose that Assumptions 2 and hold. For any 8> 0, if we set learning rate to be
4F
ai
16I2, (2n (2m+A/m) A/m) (s+o)) 18L3 m2 o
where Fi = ((w))) innf f(w) and À comes from Theorem 1 -88,
Then, with probability at least
T
(V 9(F1 (( (2 (s o2 +64F L2, (2+ A/mm))
(mnT)2/33
T
( (mn7)2/3
1
We can also prove accelerated rate for CD-GraB if we additionally assume the P.L., condition:
7

---

Assumption 4 (P.L. Condition). We say the loss function f fulfills the P.L. condition if there exists Jt > such for any w E R. 2 ( veR f())..
Theorem 3. Suppose that Assumptions C3 to be
120 and hold. For any d 0, we set constants W and
C (Fi+o²) )p²
and W W-Wo(T
224L3
000 (2+
where À tomesfrom Theorem
F1 is from Theorem and Wo the Lambert- function. If we set
learning rate a
2W Tap
and if the number of epochs T satisfies
T2 10+ 22 (m0
then, with probability least 1 1-T8, it holds that
(Fi (2
11212, (=to)__(
= =m(mtry))
Fr+15 (mnT)²
+
C3
where Fr+1- = f(w(+())))))))))))))))11))))())- inf f
We prove Theorems and in the Appendix. Together, they show that CD-GraB exhibits linear speedup in the number of workers m over GraB [24]ss convergence rates (Ó((nT)-2/4 and O(n'()) -)), respectively) under both smoothness and the P.L. condition. Further, CD-GraB's convergence rate of ((mmT) -?) " is faster than many previous rates such as the high probability bound of (7m)))17-22 for D-RR in Yun et. 48
CD-GraB in Practice: Distributed and Simulation Experiments
We next verify CD-GraB's accelerated convergence variety of empirical tasks For case of comparison, we follow the experimental plan from the original GraB paper! 11 large-scale logistic regression experiments. We also run ablation study to isolate the effects of different improvements in CD-GraB. We do this because online PairBalance exhibits performance benefits that are separate from parallelism namely, removing the need for gradient centering with a stale mean and allowing for higher learning rates (Section 3.2 14
[1 and add some additional
Evaluating CD-GraB's convergence speedup. We use the following three tasks for evaluating distributed training efficiency: logistic regression on a large-scale mortgage application (New York 2017 subset, 244, 107 examples with 18 features) [7] (Figure 3ap. Long Short- Term Memory (LSTM) [17 on the WikiTexx-2 dataset (29] (FigureBb) and autoregressive Multi-Layer Perceptron (MLP) on the M4 Weekly dataset (25) (Figure Sc). We measure the loss incurred on the entire training set (Full Train Loss) and task-appropriate test metrics during evaluation, with respect to both the number of epochs and wall-clock time. Regarding test metrics, we measure test accuracy for the mortgage application, perplexity for WikiText-2, and SMAPE for M4. Additional details regarding the datasets, models, and test metrics can be found in the Appendix.
For all three tasks, we use a single 128 GiB memory machine with 4 NVIDIA GeForce RTX 2080 Ti GPUs. For the mortgage application and WikiText-2 (FiguresBaaandd 3b). we launch 4 workers (processes), where each worker runs on one GPU. For the M4 task, we launch m = 32 workers, where each of the 4 GPUs hosts process workers. We use NCCL as the distributed communication backend (33] for the mortgage application and WikiTexxt2 tasks, and GLOO II] as the distributed communication backend for the M4 task.
As shown in F we compare CD-GraB's convergence to the standard beastribuued-trainiin example- ordering method: random reshufiling (D-RR). From all subfigures in Figure observe that CD-GraB outperforms the D-RR baseline significantly and consistently; CD-GraB cxhibits better training loss and test metrics, measured against both the number of epochs and wall-clock time. We also note that the
10 For centralized GraB, the total number of examples N = n and m =1
11 "Thee exclusively focus the P.L.. case, so we compare CD-GraB them under the same condition.
12 Our GitHub repository is https::/githuub.co. www..ooooooo/ooooGGiiiiG
13 Following et al. 24 for our LSTM experiment on WikiText-2, we set embedding dimension 32.
We note that we can improve perplexity if we set the dimension higher.
GraB can also implement online PairBalance, in place of Balancc [22] (Appendix).
8

---

0.339
GraB
0.339
Test Accuracy
330
038
Full Train
www
ww
0.337
0.337
82.2-
82.2
12:1
0.335
0.335
62.0-
82.0
Epochs
Seconds
Epochs
Seconds
(a) Logistic regression on morigage application (NY 2017 subset)
7
5.7
CD GraB
5.7
Test Perplexity
300
300
D- RR
Full Train
5.5-
260
260
5.4-
5.4
240-
240
20
40
150
300
450
750
20
150
300
450
600
750
Epochs
Seconds
Epochs
Seconds
LS) 17
on WikiTexx-2
29
0:315
0:115
9.2
M
0.310
CD-GraB
D-RR
0.110
Test SMAPE
9.0
9.0
Full Train
0:305
0.125
6.8
88-
0.100
0.100
86
0.205
0,005
8.4
84
1000
2000
1000 Seconds
4000
5000
1000 2000
3000
4000
Scon
Epochs
Epochs
Seconds
(c) Autoregressive MLP on M4 Weekly
25
Figure 3: Convergence of CD-GraB in comparison to D-RR. For each experiment, we show train loss over epochs and time (left of each subfigure) and test performance over epochs and time (right of each subfigure). We run at least 3 random seeds, and plot the mean + STD.
results for CD-GraB are much smoother than for D-RR. This is likely due to the variance of stochastic gradients during training, which CD-GraB reduces as side-effect (so, too, does GraB, in comparison to RR). For smoother D-RR results, we can reduce the learning rate (Appendix). CD-GraB allows for the use of larger learning rate, which accelerates training while preserving the final model's performance.
Ablation simulation study: the importance of coordination at large scale. CD-GraB has several de- sign benefits over the original centralized GraB algorithm [24] coordinating parallel workers' specific permutations using PairBalance on the server (Algorithm) 2) and removing the dependency on a stale mean (Section 2.Ip. which enables the ability to using larger learning rates reliably (Section 3.2). Clearly, not all of these benefits come directly from distributing training. For example, being able to use larger learning rates, a side effect of our solution to develop CD-GraB, not our main contribution. There- fore, we run a simulation ablation study to disentangle the relative importance of each of CD-GraB's B efficiency benefits over GraB. To do so, we compare the convergence of CD-GraB to two additional base- lines in the distributed setting, beyond D-RR: (1) ID-GraB (Bal), where each independent worker runs GraB locally using RandomizedBalance (subroutine in Algorithm to perform gradient vector bal- ancing; (2) ID-GraB (PairBal), where each independent worker runs GraB locally using Pair
Figure 4 summarizes the results, with convergence curves for m E 14,,,64) workers training LeNet on CIFAR-10, 10. We choose this task and architecture to cohere with the experiments done in the original GraB paper. For these experiments, the workers, which refers to the number of stochastic examples used for an overall optimization step; each worker thus has subset of this minibatch an equivalently-sinn make two main observations. First, when scaling up training with more workers, CD-GraB converges increasingly faster than the no-coordination ordering methods ID-GraB (Bal) and ID-GraB (PairBal). This result aligns with our theory and intuition that, when the number of workers increases, the parallel herding bound (8) will increase linearly if there no coordination. Second, as we scale up to larger m, the convergence curves of ID-GraB (Bal) and ID-GraB (PairBal) gradually approach
denote B be the aggregated minibatch across all
subset of B examples We
15 For example, we have workers with an aggregated minibatch size of 32, each worker would compute their respective local gradients with 8 examples, and then all-reduce these gradients to obtain the aggregated minibatch gradient for all 32 examples for the optimization step. We discard N mod B examples at random ensure examples per worker.

---

4
Workers
8 Workers
16 Workers
32 Workers
64 Workers
CD Gr3B
1.0
D Ki
Full Train Loss
ID-GreB (bwe)
0.9
ID-GraBi (PairBa))
0.8
0.7
68
Test Accuracy
66-
64-
62
50
100
50
100
50
100
50
100
50
100
Epochs
Epochs
Epochs
Epochs
Epochs
Figure Convergence for CD-GraB, D-RR, ID-GraB (Bal), and ID-GraB (PairBal) training LeNet on CIFAR-10, with m (4,8,16,32) 64) workers. For each, the aggregated minibatch size per update is 64.
the curve for D-RR: at larger scales, herding- beasedg-based example ordering will be no better than randomly permuting the dataset. Both observations give strong evidence that coordination (i.e., running online PairBalance the server to coordinate per-worker permutations) is critical for accelerating training.
We note that of these experiments use SGD, since both the theoretical results of the original GraB paper and our results for CD>GGraB here are for SGD. In the Appendix, we additionally include results for training GPT-2 on WikiText- 103,, for which we use A as the optimizer. We find that GraB with AdamW works in practice; however, our theory results do not directly apply to these experiments. We additionally include results on memory usage in the Appendix, which show that CD-GraB results in negligible overhead in practice.
Conclusion and Future Work: Toward an Order Server Architecture
We clevate the benefits of provably faster, permutation-basee example ordering to the contemporary ML distributed-training setting. We focus on reformulating the online Gradient Balancing algorithm (GraB) [24] because, even though it is the provably optimal permutation baaed example-orderine method it is limited by design to centralized settings (Section[3.1] To overcome these limitations, we redesign GraB's 's herding and balancing framework to account for parallel workers: A parallel herding objective, which we solve with an online PairBalance subroutine, based on key insights from kernel thinning [1] Pair operates on ordered pairs of vectors to do balancing, which enables our full-stack, low-overhead, Coordinated and Distributed online CD-GraB algorithm. We give full specification of our online CD-GraB algorithm (Section 3.3), provide convergence rate guarantees regarding its speedups on both !) smooth non-convex and 2) P.L. objectives (Section 4. and verify these speedups in practice single-node distributed tasks and a simulated ablation study (Section
Both our theory and experiments demonstrate that training epochs (Appendix). This is another reason that we do not emphasize experiments involving fine- tuning pre-trained models like GPT-2, as fine-tuning can be achieved just a couple of epochs. As noted above, it is also more common to train such models using optimizers from the Adam family. In future work, we intend to extend the theory on GraB and CD-GraB to such optimizers, which would make the results on optimal, commutationnbaaee example ordering more useful for base-model pre-training.
really shines when there are multiple
Pre-training from scratch would demonstrate the tremendous power of CD-GraB to scale to very large models; however, we did not have the training budget to perform such experiments for the present work. Further, to truly exercise the benefits of CD-GraB tigate moving beyond the single-node setup that we present. Notably, to train larger models, our results suggest novel distributed training architecture. The ordering operation performed by the server (Algo- rithm[4) is not very latency sensitive; the server has the duration of the entire epoch to compute the new permutations for the next, /+1 epoch. Given this relaxed latency requirement, and the success of our al- gorithmic results, it would be an exciting direction for future MI--ssttem research to invest in building an Order Server architecture. Such architecture, which could be composed with traditional parameter servers, would afford the scalability benefits of CD-GraB to host of massive-scal scale ML applications.
such large-scale settings, future work should inves-
10

---

Acknowledgments
A. Feder Cooper is supported by Christopher De Sa's NSF CAREER grant. Yucheng Lu is supported by Meta Ph.D. Fellowship. We also acknowledge gift from SambaNova Systems.
References
[1] Collective communications library with various primitives for multi-maching machine training, 2023.
com http://ww...ooeecceoooocccaaaara
[2] Ryan Alweiss, Yang P Liu, and Mehtaab Sawhney. Discrepancy minimization via a self self-balancing walk. In Proceedings of the 53rd Anmual ACM SIGACT Symposium on Theory of Computing. pages 14-20, 2021.
[3] Alessandro Barp, Carl-Johann Simon-Gabriel, Mark Girolami, and Lester Mackey. Targeted separation and convergence with kernel discrepancies. arXiv preprint arXiv:2209. /2835 2022.
[4] Dimitri P. Bertsekas. Incremental Gradient, Subgradient, and Proximal Methods for Convex Optimization: A Survey, In Optimization for Machine Learning. The MIT Press, 2011.
[5] Léon Bottou. Stochastic gradient descent tricks. In Neural networks: Tricks of the trade, pages 421 4-6. Springer, 2012.
[6] Jaeyoung Cha, Jaewook Lee, and Chulhee Yun. Tighter Lower Bounds for Shuffling SGD: Random Permutations and Beyond, In Proceedings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org, 2023.
[7] A. Feder Cooper, Katherine Lee, Solon Barocas, Christopher De Sa, Siddhartha Sen, and Baobao Zhang. Is My Prediction Arbitrary? Mcasuring Self Colf-Consittnncy Fair Classification. arXiv preprint arXiv:2301. /562, 2023.
[8] Christopher De Sa. Random Reshuffling is Not Always Better, In Advances in Neural Information Processing Systems, 2020.
[9] Jacob Devlin, Ming-Wes Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pages 4171 4186, 2019.
[10] Raaz Dwivedi and Lester Mackey. Kernel thinning, arXiv preprint arXiv:2105.05844 2021.
[11] Raaz Dwivedi and Lester Mackey. Generalized Kernel Thinning. In Tenth International Conference Learning Representattinss 2022.
[12] Alex Graves, Marc G Bellemare, Jacob Menick, Remi Munos, and Koray Kavukcuoglu. Automated curriculum learning for neural networks. In international conference on machine learning, pages 1311-1320. PMLR, 2017.
[13] Mert Gürbüzbalaban, Asuman E. Ozdaglar, and Pablo A. Parrilo. Convergence Rate of Incremental Gradient and Incremental Newton Methods. SIAM Journal on Optimization, 29 (4):2542-2565, 2019.
[14] Mert Gürbüzbalaban, Asu Ozdaglar, and Pablo A Parrilo. Why random reshufiling beats stochastic gradient descent. Mathematical Programming, 186(1):49 84, 2021.
[15] Jeff HaoChen and Sra. Random Shuftling Beats SGD after Finite Epochs. In Proceedings of the International Conference on Machine Learning, volume 97, pages 2624-2633, 2019.
[16] Nick Harvey and Samira Samadi. dear-Optimal Herding. In Proceedings of 27th Conference on Learning Theory, volume 35, pages 1 1165-11 1182,1 2014.
[17] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, (8):1735-1780, 1997.
[18] Kun Huang, Xiao Li, Andre Milzarek, Shi Pu. and Junwen Qiu. Distributed Random Reshuffling over Networks. arXiv preprint 1rX11221 2021.
[19] Hakan Inan, Khashayar Khosravi, and Richard Socher. Tying word vectors and word classifiers: 2016.
A loss framework for language modeling. arXiv preprint
11

---

[20] Mu David G. Andersen, Jun Woo Park, Alexander J. Smola, Amr Ahmed, Vanja Josifovski, James Long, Eugene Shekita, and Bor- Yiig Su. Scaling Distributed Machine Learning with the Parameter Server. In Proceedings of the 1thh USENIX Conference Operating Systems Design and Implementation, OSDI' 14, page 583-598, USA, 2014. USENIX Association. ISBN 978 1931971164.
[21] Ily Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017.
[22] Yucheng Lu, Si Yi Meng, and Christopher De Sa. A General Analysis of Example-Selcctee Selection for Stochastic Gradient Descent. In International Conference on Learning Representations, 2021.
[23] Yucheng Lu, Youngsuk Park, Lifan Chen, Yuyang Wang. Christopher De Sa, and Dean Foster. Variance Reduced Training with Stratified Sampling for Forecasting Models. In Proceedings of the International Conference on Machine Learning, pages 7 7145-7155. PMLR, 2021.
[24] Yucheng Lu, Wentao Guo, and Christopher De Sa. GraB: Finding Provably Better Data Permutations than Random Reshuflling. In Alice H. Oh, Alckh Agarwal, Danielle Belgrave, and Kyunghyun Cho, Advances in Neural Information Processing Systems, 2022. URL qKHTp
[25] Spyros Makridakis, Evangelos Spiliotis, and Vassilios Assimakopoulos. The m4 competition: 100,000 time series and 61 forecasting methods. International Journal of Forecasting, 36(1): 54-74, 2020. ISSN 0169-2070. doi: http:://oi.org//0..000 http:///ww....ooooooooooaaaaoaa. URL https://wwit
M4
Competition.
[26] Grigory Malinovsky, Konstantin Mishchenko, and Peter Richtárik. Server-Side Stepsizes and Sampling Without Replacement Provably Help in Federated Optimization. arXiv preprint arXiv:2201.11066, 2022.
[27] Tambet Matiisen, Avital Oliver, Taco Cohen, and John Schulman. Teacher staleet--ttenttte curriculum learning. IEEE transactions on neural networks and learning systems, 31(9):3732-377 2019.
[28] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. communication-effcccennn learning of deep networks from decentralized data. In Artificial intelligence and statistics, pages 1273-1282. PMLR, 2017.
[29] Stephen Merity, Nitish Shirish Keskar, and Richard Socher. Regularizing and optimizing Istm language models. In International Conference on Learning Representations, 2018.
[30] Konstantin Mishchenko. Ahmed Khaled, and Peter Richtárik. Random Reshufiling: Simple Analysis with Vast Improvements. In Advances in Neural Information Processing Systems, 2020.
[31] Amirkeivan Mohtashami, Sebastian Stich, and Martin Jaggi. Characterizing & Finding Good Data Orderings for Fast Convergence of Sequential Gradient Methods. arXiv preprint urXiv:2202.018323 2022.
[32] Deanna Needell, Rachel Ward, and Nathan Srebro. Stochastic Gradient Descent, Weighted Sampling, and the Randomized Kaezmarz algorithm. In Advances in Neural Information Processing Systems, pages 1017-1025,2014
[33] NVIDIA. NVIDIA Collective Communication Library, 2023. URL https: //developer.. wwwpperr.viiiiee wwwiia.iiiaoc.m... com/nccl
[34]
Py PyTorch Contributors. DataLoader API, 2023.
data. html
[35] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
[36] Shashank Rajput, Kangwook Lee, and Dimitris Papailiopoulos. Permutation- -Based SGD: Is Random Optimal? In International Conference on Learning Representations, 2022.
[37] Benjamin Recht and Christopher Ré. Toward Noncommutative fitthetticgeommttr Mean Inequality: Conjectures, Case-studies, and Consequences. In Conference on Learning Theory. volume 23, pages 11.1-11.24, 2012.
[38] Abdurakhmon Sadiev, Grigory Malinovsky, Eduard Gorbunov, Igor Sokolov, Ahmed Khaled, Konstantin Burlachenko, and Peter Richtárik. Federated Optimization Algorithms with Random Reshuffling and Gradient Compression. arXiv preprint arXiv:2206.07000 2022.
12

---

[39]
Mark Schmidt, Nicolas Le Roux, and Francis R. Bach. Minimizing finite sums with the stochastic average gradient. Mathematical Progranmingg 162(1-2):83-112 2017.
[40] L. Pieter-Jan Kindermans, Chris Ying, and Quoc Le. Don't Decay the Learning Rate, Increase the Batch Size. In International Conference on Learning Representations, 2018 88.
[41] Petru Soviany, Radu Tudor lonescu, Paolo Rota, and Nicu Sebe. Curriculum learning: A survey. International Journal of Computer Vision, pages 1-40, 2022.
[42] Evangelos Spiliotis, Andreas Kouloumos, Vassilios Assimakopoulos, and Spyros Makridakis. Are forecasting competitions data representative of the reality? International Journal of Forecasting, 36(1):37--))) 1):3-53,, 2020.
[43] Merity Stephen, Xiong Caiming, Bradbury James, and Richard Socher. Pointer sentinel mixture models. Proceedings ICLR, 2017.
[44] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pages 353-355, Brussels, Belgium, November 2018. Association for Computational Linguistics.
doi: 10. 10.18653/v1/W18-566
URL
https://aclanthooogyy wwwphhooog...og/////. 2rg//11-5446.
[45] Max Welling. Herding dynamical weights to learn. In Proceedings of the 26th Annual International Conference on Machine Learning, pages 1121-1128, 2009.
[46] Bicheng Ying, Kun Yuan, Stefan Vlaski, and Ali H. Sayed. On the performance of random reshuffling in stochastic learning. 2017 Information Theory and Applications Workshop (ITA), pages 1-5, IEEE, 2017.
[47] Binhang Yuan, Yongjun He, Jared Davis, Tianyi Zhang, Tri Dao, Beidi Chen, Percy S Liang, Christopher Re, and Ce Zhang. Decentralized training foundation models heterogeneous environments. Advances in Neural Information Processing Systems, 35:25464-25477, 2022.
[48] Chulhee Yun, Shashank Rajput, and Suvrit Sra. Minibatch vs Local SGD with Shuffling: Tight Convergence Bounds and Beyond. In International Conference on Learning Representations, 2021.
[49] Chulhee Yun, Suvrit Sra, and Ali Jadbabaie. Open Problem: Can Single-Shon Shuffle SGD be Better than Reshuffling SGD and GD? In Conference on Learning Theory, 2021.
13


Reviews:
[{"review_id": "5aHhkxS5nC", "reviewer": "Reviewer_dDcg", "summary": "This paper proposes a method to extend the optimal data order finding method GraB to centralized distributed setting. This is achieved by using an alternative balancing method that does not require using stale averages which would cause a problem given higher learning rates in the distributed settings. Furthermore the ordering task is determined by a global coordinator leading to CD-GraB. The effectiveness is demonstrated empirically over random reshuffling as well as running GraB locally on each worker (with and without the alternative balancing method). Finally, an analysis of the convergence rate is given for non-convex smooth functions as well as functions satisfying PL. ", "questions": "1. Looking at Figure 3, it seems using GraB allows the optimization to reach a better train loss. Is this true or if the normal RR is run for a longer period it finally reaches GraB's loss? \n\n2. How does using CD-GraB affect the learning rate schedule? ", "limitations": "No specific discussion of limitations is included. I think a quantitative discussion around additional resource requirements (e.g. memory) is needed. Furthermore, as a replacement for the current optimization methods, a discussion about its ability to combine with existing solutions in distributed settings is missing. For example, I am not sure if this method can be combined with existing privacy frameworks which could also highlight a direction for future work.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The proposed alternative balancing method is novel and can be useful both for the single-node and multi-node versions of GraB. On the experiments CD-GraB is faster both in terms of epoch and wall clock time than random reshuffling. Ablation studies clearly show the importance of using the central coordinator. The paper is written clearly and is easy to follow and understand.", "weaknesses": "One of the main drawbacks is that the experiment are quite small in scale. The main comparison is using MLPs and logistic regression while the ablation is done using LeNet instead of the larger ResNet models. For the ablation, the number of epochs is also limited and as a result the obtained accuracy is no where near the state of the art. This makes it hard to judge whether the observed improvement will continue to be there as the training goes on further. Since the paper is advocating for using CD-GraB in practice, large-scale evidence is essential. I would also suggest considering the effect of different architectures including Transformers. If for some reason the method is not applicable to some settings this should be clearly discussed.\n\nIn addition, in many cases other optimizers such as Adam are used in practice. The possibility of using this method with such optimizers (particularly Adam) as well as its effect on the performance is not investigated which again makes it harder to judge whether it can be directly deployed in practice or not.  \n\nMoreover, one thing that I would assume can considerably affect the performance in the distributed setting is the data heterogeneity. However, no discussion on this has been made. It would be nice to have at least one set of experiments in the heterogeneous settings.\n\nAs a side note, since the pair balance algorithm can also be applied in the single-node settings, a good comparison to show the effect of this switch is comparing it with the original GraB in the single node setting, both in terms of speed and the best achievable loss. Though this is covered to some extent given that the performance of local-GraB and CD-GraB match with 4 workers (according to Figure 4), I think a direct comparison in the single node setting would still be useful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review, and for appreciating the writing quality of our submission. We respond directly to your comments about weaknesses below.\n\n**Scale of experiments** \n* We note in the paper/Appendix that we ran we could afford to. CD-GraB converges faster than D-RR, which the existing results show. We are academics, not an industry lab. We can run LeNet for longer, but saved our training budget for other tasks. \n\n* Based on the intended scope of our contributions (i.e., primarily theory, with empirical implications to be explored in future systems work), the fairest comparison for our experiments is to the original GraB paper (also a theory-focused paper). We used this prior work as the basis for designing our experiments. \n\n* We additionally include results for a small GPT-2 on WikiText-103 in the Appendix. We note in the conclusion why we did not focus on transformers. Models like LLMs today are finetuned typically for a really small number of epochs. The benefits of CD-GraB (i.e., getting the theory-backed benefits of achieving the herding bound) only come in the **multi-epoch** setting, and we cannot afford the pretraining costs of a LLM for a lot of epochs. \n\n* Even so, we demonstrate the benefits of CD-GraB on other ML tasks, and believe that this is an important contribution. It is worth investing the effort in future ML systems research to build and stress-test a full implementation of an order server at scale. Without the present work, we would not have been able to know that such an investment was worth it. Such a system necessarily is the realm of future work. \n\n**Adam**: The theoretical guarantees for GraB (upon which we base CD-GraB) are for SGD. So it makes sense to run our experiments using SGD. We can provide more details about this context in revisions. This is another reason also we de-emphasize our GPT-2 experiment, which uses AdamW for pretraining and finetuning.\n\n**Data heterogeneity:** Data heterogeneity is common in federated learning. We state in footnote 5 that our formulation is not designed for federated learning; it is not in scope. 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. \n\n**Single node pair-balance:** The point of our method is to coordinate and distribute computation where we have >1 worker. We can add this to our Appendix, but do not believe this is important for show-casing our main results.\n\n**D-RR v. GraB:** GraB often performs better than RR in practice (see also the original GraB paper), and in cases like LSTM on WikiText-2 is is clear that running RR longer still cannot reach GraB's training loss, and this is demonstrated both in our paper and the original GraB paper.\n\n**Learning rate schedule**: We used standard settings for the tasks that we investigated, and for the LSTM on WikiText-2 experiment we still used learning rate scheduler. Investigating how to tune the learning rate schedule in response to our improved convergence is the realm for future theory work. \n\n**Memory:** We believe that the extra memory requirements follow directly from our algorithm, so we cut the histogram of LSTM on WikiText-2 memory for space, which is associated with the Figure 3.b. We need 1 copy of the model parameters as the accumulator in the order server, and both obtaining a pair gradient element-wise difference and calculating inner product computation during the dataset sorting step is O(1). The accumulator update is in-place. Notice that in experiments we always place the order server in the GPU, but we could further alleviate the memory concern by moving it to the CPU as getting next-epoch permutation is not blocking any of the training stages in GPU.\n\n| Memory Overhead (MiB) for LSTM on WikiText-2 | Optimizer | Forward & Backward | Communication | Dataset Sorter\n| :---        | :----: |   :---: | :---: | :---: |\n| CD-GraB  | 4.13   | 36.29   | 16.51 |  4.36 |\n| D-RR        | 4.13   | 36.29   |  0.00 |  0.03 |\n\n**Overall comment about review and project scope:** Given the length of conference submission, we don’t have much more room for experiments than the quantity that we already present. Our main contribution, which fills the majority of our paper and Appendix, is theoretical. To include 3 additional sets of experiments (as requested) would require us to downplay our main contribution. We aren’t claiming that our work is a general-purpose solution for all ML (e.g., privacy-preserving frameworks). We situate our discussion in relation to parameter-server architectures and the distributed training setup. How our work could be combined with privacy techniques is the subject of another paper.\n\n**Request to please consider revising score**: \nThe review acknowledges that we’ve made a sound and novel theory-backed contribution that we support with experiments. We understand that our experiments are not the largest scale that they could be, but we ran what we could afford to and believe that the experiments we ran are appropriate for the contributions we are making. They are in line with and go beyond what GraB did. \nThe emphasis on empirics in this review does not address the value of our theoretical contributions. Given the length of a conference paper, to make all of the additions that this review requests (which, as discussed above, we don’t believe are entirely appropriate for the present work) would require us to significantly cut down on our discussion of our main contribution. This feedback is helpful for us to consider when scoping out future work. \nNevertheless, given the novelty and promising nature of our results, our theory contributions, and how our work suggests that it is worth investing future effort in building a large scale system (which we could not have known was worth doing before the present work), we believe such a low score (3/rejection) is quite harsh. \n\nThank you for your consideration."}, {"review_id": "X5dOVlQOB1", "reviewer": "Reviewer_AQrA", "summary": "The paper introduces Coordinated Distributed GraB, a new method that leverages stale gradients from prior epochs to order examples and achieve a provably faster convergence rate than random reshuffling. CD-GraB demonstrates a linear improvement in convergence rate compared to centralized GraB and showcases superior performance over alternative methods across various benchmark tasks. The contributions of this paper encompass a thorough theoretical analysis of CD-GraB's convergence rate as well as empirical evaluations conducted on three benchmark tasks.", "questions": "The theoretical analysis in this paper assumes an ideal scenario where the number of data examples is evenly divisible by the number of workers. However, it is crucial to consider the implications when the numbers cannot be perfectly divided or when workers have different numbers of update iterations. How does CD-GraB address these challenges? Is there any theoretical analysis available to support its handling of such problems?", "limitations": "No.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.The paper presents an innovative method called Coordinated Distributed GraB, which builds upon previous work on GraB and kernel thinning. This novel approach capitalizes on the utilization of outdated gradients from prior epochs to prioritize examples, resulting in a demonstrably faster convergence rate compared to random reshuffling.\n\n2.This paper offers a comprehensive theoretical analysis of CD-GraB's convergence rate.  The analysis reveals a linear acceleration in the convergence rate of CD-GraB.", "weaknesses": "1. This method is designed specifically for the paradigm involving a parameter server and workers. However, it is important to note that as the number of workers increases, there will be a linear increase in overhead, potentially affecting performance during large-scale training sessions.\n\n2. In comparison to the widely used ring allreduce method for gradient aggregation, the communication of gradients between the parameter server and workers in this method is less efficient. Consequently, the improvement in wall-clock time may not be convincing. Furthermore, the unavailability of the code makes it difficult to reproduce the experiments.\n\n3. The convergence analysis presented in the paper may not be easily applicable to adaptive optimization methods such as AdamW, which is a popular choice for training deep neural networks.\n\n4. It is worth noting that the proposed method is based on stochastic gradients, while it is common practice to train models using minibatch stochastic gradients.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and questions, \n\n> 1. This method is designed specifically for the paradigm involving a parameter server and workers. However, it is important to note that as the number of workers increases, there will be a linear increase in overhead, potentially affecting performance during large-scale training sessions.\n\n**Overhead**: We do not observe a meaningful impact for overhead in our distributed experiments or in our simulations. It is of course possible that communication (of gradients) could cause overhead, but we do not observe this to be meaningful in practice for our experiments. As noted in our response to reviewer dDcg, we could not afford to perform larger scale experiments, which we do not believe should be grounds for rejection. Memory overhead is minor for CD-GraB/ PairBalance, which we discuss in the paper. \n\n> 2. In comparison to the widely used ring allreduce method for gradient aggregation, the communication of gradients between the parameter server and workers in this method is less efficient. Consequently, the improvement in wall-clock time may not be convincing. Furthermore, the unavailability of the code makes it difficult to reproduce the experiments.\n\n**Parameter server**: Parameter servers are also a widely-used distributed training paradigm (see cited works in the paper). Our empirical results demonstrate improvements in wall-clock time, as this review notes. We are not sure why this would not be convincing evidence of the success of our method, and we could provide a histogram of forward / backward / communication / dataset permutation computation time per iteration if needed. \n\n> 3. The convergence analysis presented in the paper may not be easily applicable to adaptive optimization methods such as AdamW, which is a popular choice for training deep neural networks.\n\n**Theory and AdamW**: The theory applies to SGD, as the theory for the original GraB algorithm also applies to SGD. Theory for other optimization methods is deferred to future work. We believe that this contribution is itself useful, laying the groundwork for future improvements. This is not a drawback; there is just more research to do in this area in the future. We nevertheless test CD-GraB using AdamW when training a small GPT-2 in the Appendix (which we observe performs well, though does not directly inherit the theory we contribute).  \n\n> It is worth noting that the proposed method is based on stochastic gradients, while it is common practice to train models using minibatch stochastic gradients.\n\n**Minibatches**:  Our theory applies to stochastic gradients, just as the original GraB paper does (see their Appendix). As stated in top-level comments A and B, our main contribution is theoretical. Similar to the original GraB paper (see Appendix), we could do additional engineering work to make our approach work for minibatching. This is not something that wouldn’t work for our method, but we don’t believe that this work is necessary for showcasing our main results. Instead, this is work we will include for future systems work that stress-tests all of the systems and implementation efficiency considerations in modern distributed training. Here, we emphasize showing that our method works for per-example gradients in order to map directly onto our problem setup.  Also notice that *when the number of workers is greater than 1, the gradients computed by the central parameter server are also minibatched*. We also consider the setting where the *microbatch per worker is greater than 1* in our logistic regression on mortgage application experiment, where we use 4 workers and a microbatch of 4 examples per worker (so the minibatch per iteration on parameter server aggregation $B$ is 4 * 4 = 16).\n\n> The theoretical analysis in this paper assumes an ideal scenario where the number of data examples is evenly divisible by the number of workers. However, it is crucial to consider the implications when the numbers cannot be perfectly divided or when workers have different numbers of update iterations. How does CD-GraB address these challenges? Is there any theoretical analysis available to support its handling of such problems?\n\n**Ideal scenario**: We discuss this in the paper and in the Appendix. We divide the examples M mod N on the available servers, and discard (at random) the remaining up to N - 1 examples. For large datasets (and even something like N = 64 in the worst case), this should not matter (indeed, we do not observe this to matter empirically). Indeed even in a non-industry scale dataset like CIFAR-10 that has only 50,000 examples, discarding 63 examples in the worst case is still far less than 50,000 and we did not observe a visible performance degradation. "}, {"review_id": "vrU1k6V6CO", "reviewer": "Reviewer_ZrZn", "summary": "The authors propose a distributed version of Grab. The authors both theoretically and experimentally show that there exists a pattern which can improve training convergence of such a method in a distributed setup. \n\nThe paper is very well and tightly written with clear limitations. The authors verified their claims both experimentally and theoretically. \n", "questions": "Compared to D-RR graphs CD-Graph seems to very smooth. Do the authors have an explanation for this ? ", "limitations": "The authors have discussed it well. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "\n* The authors have clearly motivated the problem and discussed it in detail how their work compare to other prior works\n\n* The experiments are well executed and clearly demonstrate the benefit of CD-GRAB. \n\n* The authors discussed the proof in detail in the appendix and organized it well.\n \n* All limitations are discussed. \n", "weaknesses": "* As the authors have stated that there is additional memory overhead keeping older gradients. As memory becomes more scarce due to larger models,do authors think there is a way of approximating the previous epoch gradients ?\n\n* The loss for baselines seems a bit high, can authors comment on that, for ex- the PyTorch LSTM example on wikitext-2 gives a loss of around 100, as shown here - https://arxiv.org/abs/1905.13727 why the perplexity score is so far off ?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and questions. Please see our answers below.\n\n> As the authors have stated that there is additional memory overhead keeping older gradients. As memory becomes more scarce due to larger models, do authors think there is a way of approximating the previous epoch gradients?\n\n**Memory overhead, gradient approximation**: We have included quantitative information about the low memory overhead in our response to reviewer dDcg in ML tasks we are afford to run. Overall, the overhead is small for CD-GraB, but we are actively investigating other ways to approximate gradients for example ordering (which will be the subject of future submissions). Our hope is that our next (in-progress) work on this will also no longer require computing per-example gradients. We can include additional details on limitations and future work in our conclusion (which notes other limitations) in revisions.\n\n> The loss for baselines seems a bit high, can authors comment on that, for ex- the PyTorch LSTM example on wikitext-2 gives a loss of around 100, as shown here - https://arxiv.org/abs/1905.13727 why the perplexity score is so far off?\n\n**LSTM loss on wikitext-2**:\nWe follow the original GraB paper (except that we tied input embedding matrix weight with the output word classifier (the weight-typing strategy) as indicated in the appendix) to use the embedding dimension as 32. If we use a larger embedding dimension (200 or higher), we could also reach test perplexity below 100. \n\n**Smoothness of CD-GraB results compared to D-RR**: \nThe non-smoothness of D-RR is likely due to the variance of stochastic gradients during training, and CD-GraB would reduce the variance of stochastic gradients during training. If we want to make the D-RR training loss curve smoother, we could use a tiny learning rate. We provide an ablation study in Appendix D.3 on the impact from the learning rate and we could observe that D-RR becomes smoother when we decrease the learning rate. CD-GraB, on the other hand, allows the use of a large learning rate to accelerate the training while preserving the final model’s performance. This is a practical advantage of using CD-GraB: stabilizing training with large learning rate would need significantly less optimization steps to reach the same performance.\n"}, {"review_id": "yTZUPHxssj", "reviewer": "Reviewer_xnQq", "summary": "Propose an online method, CD-GraB for shuffling data samples during distributed model training (parameter-server setting) in a fashion that accelerates convergence in both theory and practice.\n\nTheoretical assumptions are reasonable and quite standard in the literature. Empirically, this work mostly builds on the GraB method, which was described for centralized settings, and extended to the distributed setting in this work.\n\nThe difficulty in extending the algorithm to the distributed setting is that\n* the considered problem setting only allows workers to produce permutations of their local data samples\n* large step sizes results in relatively large disparities across gradients at different workers", "questions": "Please see questions described in-line in weaknesses section.", "limitations": "Limitations of the work are not adequately addressed.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "**Quality**\nThe work appears to be of a high quality. The proposed method and proofs are non-trivial, and the problem addressed by this work, building on the GraB method, requires careful thought.\n\n**Clarity**\nDiagrams are polished, notation is reasonable. The proposed method is non-trivial, and the authors have done a reasonable job at explaining the procedure, with additional illuminating diagrams in the appendix.", "weaknesses": "**Significance**\nWhile the overall contribution of the work is impressive, it remains the case that, a) the method is an extension of GraB, and b) not obvious to me that the proposed problem setting is very practical for the considered use cases. It seems to me that much of the difficulty of the analyzed problem is self-imposed: The difficulty in extending the algorithm to the distributed setting is that\n* the considered problem setting only allows workers to produce permutations of their local data samples… but often in large-scale DDP model training, data can be read by any worker, and is not stored directly on the computing node. This would enable one to simply extend GraB to such settings.\n* while large step sizes may indeed results in large disparities across gradients at different workers… why not AllReduce the gradients, run a gossip step, or simply average the gradients in the parameter server as is typically employed in DDP? In fact, the proposed algorithm already has an aggregation step at the parameter server.\n\nI am also concerned about the overhead of running the PairBalance subroutine in large-scale settings with many workers, especially as the size of model parameters balloons.\n\nOn the empirical side, I appreciate the breadth of the exploration, but the improvements training LeNet on CIFAR10 seem quite odd. Why should performance of D-RR change  with the number of workers, when training with the same batch size?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for observing the challenging nature of our work and the value of our contributions (both in terms of novel results, and in clarifying them / communicating them effectively). We respond below.\n\n> Significance While the overall contribution of the work is impressive, it remains the case that, a) the method is an extension of GraB, and b) not obvious to me that the proposed problem setting is very practical for the considered use cases. It seems to me that much of the difficulty of the analyzed problem is self-imposed: The difficulty in extending the algorithm to the distributed setting is that the considered problem setting only allows workers to produce permutations of their local data samples… but often in large-scale DDP model training, data can be read by any worker, and is not stored directly on the computing node. This would enable one to simply extend GraB to such settings.\n\n**Problem Setting**: Such an extension of GraB to the distributed setting is an important contribution. GraB is the optimal permutation-based ordering method (see citations within the paper), but does not extend to distributed training and thus is fundamentally limited. \n\n> while large step sizes may indeed results in large disparities across gradients at different workers… why not AllReduce the gradients, run a gossip step, or simply average the gradients in the parameter server as is typically employed in DDP? In fact, the proposed algorithm already has an aggregation step at the parameter server.\n\n**Other DDP approaches** We understand and acknowledge that there exist other solutions to the disparities across gradients in distributed learning, but our solution is orthogonal to theirs. We only consider the problem that how we can find the data permutation for each clients so that their disparities are reduced, and the gossip step approach is not directly related to our solution. \n\n> I am also concerned about the overhead of running the PairBalance subroutine in large-scale settings with many workers, especially as the size of model parameters balloons.\n\n**Memory**: We do not observe overhead issues with the experiments that we could afford to run. We include quantitative results in our response to Reviewer dDcg. As we discuss in our response to reviewer ZrZn, we are actively investigating other ways to approximate gradients as a part of follow-on work to achieve even greater scalability. However, this is out of the scope of the present paper. We can include additional details on limitations and future work in our conclusion (which notes other limitations) in revisions.\n\n> On the empirical side, I appreciate the breadth of the exploration, but the improvements training LeNet on CIFAR10 seem quite odd. Why should performance of D-RR change with the number of workers, when training with the same batch size?\n\nWe agree that theoretically D-RR will not be impacted when we use the same batch size in the parameter server side, but in practice, the changes in Figure 4 for D-RR is likely due to the fact that as we change the number of workers, different workers would have a different partition of the dataset. \n\nWe also notice that **the difference across the number of workers for D-RR is even less than the standard deviation (the shaded area) when we have the same number of workers** in Figure 4 where we only change the initialization and data permutation of D-RR as we choose multiple seeds. So the fact that the performance of D-RR changes with the number of workers should be of less concern.\n"}, {"review_id": "Wi6K6ZwAts", "reviewer": "Reviewer_T8gT", "summary": "This paper is devoted to the development of the Gradient Balance for the centralized distributed setting.\nThe proof is highly non-trivial because the extension to distributed ML is not natural. To this end,  the authors use kernel thinning to translate the benefits of provably faster permutation-based example ordering to distributed settings.", "questions": "What is the main technical difficulty of convergence analysis for distributed GraB.\n\nMinor: line 32: \"Gradiant Balance algorithm\"  should be \"Gradient Balance algorithm\" ", "limitations": "Yes.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper proposes the online Coordinated Distributed Gradient Balancing (CD-GraB) algorithm enabling provably accelerated training in the parallel setting. The new algorithm leverages insights from kernel thinning to elevate the herding framework of centralized GraB (GraB) to the parallel setting.\n\n2. They prove that the convergence rate for CD-GraB exhibits a linear speedup over GraB, which is non-trivial.\n\n3. Empirical results are presented to demonstrate the efficiency of the proposed algorithm.", "weaknesses": "1. The presentation may need to be improved. For example, if we turn to  Table 1 （Algorithm 1), the algorithm is not self-explanatory because some notation are not defined.\n\n2. The paper uses a different smoothness assumption. The authors seem to want to hide the degradation from the dimension $d$.\n\n3. The technical difficulty of convergence analysis can be presented as roadmap for the reader.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We respond below.\n\n> 1. The presentation may need to be improved. For example, if we turn to Table 1 (Algorithm 1), the algorithm is not self-explanatory because some notation are not defined.\n\n**Algorithm 1**: Thank you for this feedback. Can you please indicate where the presentation is unclear? Which notation did we miss? We will address this in the following discussion period and the subsequent revisions. \n\n> 2. The paper uses a different smoothness assumption. The authors seem to want to hide the degradation from the dimension $d$\n\nSmoothness assumptions follow from the prior work (the original GraB paper) in this space. There is nothing hidden here. \n\n> 3. The technical difficulty of convergence analysis can be presented as roadmap for the reader.\n\n**Proof roadmap**: Thank you for your feedback about this. We provide such a roadmap below, which we will include in revisions of the Appendix. \n"}]

Author rebuttal:


# Top-level response to all reviewers

Thank you for your comments. We appreciate the detailed feedback, and think that there are ideas in the reviews that would be interesting to pursue in future work (but unfortunately don’t fall in scope for the present project). We clarify below.

## A. Prior work, our work, & future work 

We clarify how our work fits with prior work, what we intend to (and do) solve, and what this enables for future work.

### i. Prior
Provably better permutation-based example orderings **for SGD** is a relatively new area. We consider there to be 3 papers (starting in 2022) that directly precede ours: 
1. “A General Analysis of Example-Selection for Stochastic Gradient Descent” provides a framework for reasoning about example order and **SGD**

2. “GraB” finds provably better orders **for SGD** in the **single node** case

3. “Tighter Lower Bounds for Shuffling SGD Random Permutations and Beyond” proves that GraB is optimal (for SGD, single node)

**None of this work applies to the modern setting of distributed training.** It assumes local access to all examples, which isn't the case in popular distributed settings (e.g., parameter server with workers with access to subsets of examples). 

### ii. Our work
Our main contribution is to prove essential theory that lifts this promising prior work to the distributed setting. **Our main contribution is theoretical.** We make GraB work in the distributed setting (i.e., better permutation-based example ordering for **SGD**) and guarantee a linear speed-up in the number of distributed workers. We solved the problem we set out to solve (and verified our solution with experiments **for SGD**). Based on our reading of the reviews, the novelty and role of our theory is not in dispute (and this is our main contribution).

### iii. Future work 
Our theory has direct empirical implications for future work. We'd like to build a distributed ML system, which has an order server, parameter server(s), and workers that efficiently orchestrate distributed example selection. **But to know that it was worth building such a system, we first had to first make the contributions in this paper.** We had to:
 
1. Prove that it was possible to make GraB work in the distributed setting (which we do); and,
2. Verify empirically that we're able to see improvements in convergence time (which we also do). 

These two aims more than fill a conference-length paper (we have a large Appendix). Contributions beyond this are out of scope, e.g.: 

* **Systems considerations**: Future work should address implementation optimizations, the engineering particulars of minibatching and CD-GraB (also see the original GraB paper’s Appendix for a sketch of how such an implementation would work), full system-level profiling of our method in comparison to other balancing strategies, detailed HPO, etc. **Such analysis would be significantly more useful for such a fully-implemented ML systems paper, which is not what our present paper is about.**

* **Other optimizers**:  It'd be interesting to study example ordering for Adam and other optimizers (both single node and distributed). These are contributions for other papers, not the present one. E.g., Single-node guarantees for Adam would be a different way to extend prior work, but that isn't the direction we took; we solved something else. This is also why our experiments in the main paper involve **SGD**, not other optimizers; see Appendix for additional experiments.

## B. We ran the experiments we could afford to

Since our main contribution is theoretical (see A.ii), **our empirical contributions are secondary**. We'll adjust the paper presentation to make this clearer (this change is cosmetic, not substantive). As academics, we had a limited training budget for this project, which we exhausted for a project with these theoretical aims. Further, larger scale stress-testing of distributed example selection requires future research investment to build a working distributed ML system, which we intend to budget for and do. **But first**, we needed to make sure our approach was theoretically sound --- to justify making such an investment in systems research in this direction (see A.iii). **Our experiments achieve this aim: They show that our theory works in practice; they validate that it'd be worth investing significant resources in building such a system and testing it at a much larger scale.**

## C. Our empirical work should use GraB as the basis for comparison
For the reasons above (theoretical aims, relationship to prior work), **our experiments should be evaluated in comparison to what the original GraB paper did**. We made our experiments consistent with those the original authors ran, **and we did more** (including larger scale, see our Appendix). We believe that this is the fair comparison for what was sufficient to stress-test our theoretical contributions. 

## D. Memory footprint results

Memory requirements follow directly from our algorithm: We need 1 copy of the model parameters as the accumulator in the order server; both obtaining a pair gradient element-wise difference (could be done in-place) and calculating inner product computation during the sort step is O(1). The accumulator update is in-place so there is no memory overhead. We cut figures for space, but will include in revisions, e.g., for LSTM on Wikitext2 (Figure 3b), see the attached PDF. The memory bottleneck is more relevant when it comes to manifesting and communicating the per-example gradients on the order-server side. However, notice this is still a minor concern because **clients only need to know the next epoch data permutation at the end of the epoch**. We could place the order server in CPU and ***offload the gradients in CPU RAM**, which usually has much larger memory than GPU (we didn't do this for simplicity of our implementation, but could, as dataset sorting is not blocking any of the forward / backward / optimization stages).

Meta-review: This paper proposes the Coordinated Distributed Gradient Balancing method (CD-GraB) to accelerate training in the distributed setting. The theoretical analysis is provided for general smooth nonconvex functions and functions with the PL condition. Experimental results on logistic regression, LSTM, and autoregressive MLP show that the proposed method improves over the distributed random reshuffling (D-RR). Reviewers unanimously recommended accepting this paper and all concerns were well-addressed in the rebuttal. Therefore I recommend accepting this paper.
Decision: Accept (poster)