Title: Coordinating Distributed Example Orders for Provably Accelerated Training
Venue: NeurIPS 2023
Accepted: yes
Keywords: permuted example ordering

distributed training

scalable training

herding
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

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.

1 Introduction

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 C oordinated D istributed Gra diant B alance algorithm (CD-GraB), which leverages insights

37th Conference on Neural Information Processing Systems (NeurIPS 2023).

∗

Christopher De Sa 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 ˜ O (( mnT ) -2 / 3 ) on smooth, non-convex objectives and ˜ O (( mnT ) -2 ) under the PolyakŁ ojasiewicz (P.L.) condition. 2

Werun 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 C oordinated D istributed Gra dient B alancing (CD-GraB) algorithm, which enables provably accelerated training using SGD in the parallel setting (Section 3);

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 (Section 5).

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 (Section 3) imparts the efficiency guarantees of GraB to the parallelized regime (Section 4).

Ordering data examples during training. Training a model can be formulated as minimizing a differentiable loss function f : R d → R over N data examples. The goal of this minimization is to obtain the target model weights w ∗ =argmin w f ( w ) , where f ( w ) = 1 N ∑ N j =1 f ( w ; j ) , for which f ( w ; j ) denotes the loss incurred on the j -th example. A typical training process iteratively updates the model parameters w by scanning over the N data examples repeatedly, with t -th scan (or epoch) following

where α denotes the learning rate, and π t : [ N ] → [ N ] denotes a permutation ordering 3 adopted in the t -th epoch from which the examples are chosen to compute gradients, w 1 t denotes the initial model weights for the t -th epoch, and w j t denotes the model weights after j -1 gradient updates in the t -th epoch. 4

The choice of ordering π 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) [4, 13], 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.

[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 π ∗ that guarantees will yield a convergence rate of ˜ 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 π ∗ , 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 5 { z j } N j =1 ( z j ∈ R d ), ‖ z j ‖ 2 ≤ 1 ( ∀ j ), herding identifies a permutation π ∗ such that

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 π to reduce the bound in (3). That is, balancing formulates the signed herding problem

∥ Given a group of such signs { s j } N j =1 and an arbitrary permutation π , Harvey and Samadi prove that Algorithm 1 produces a new permutation π ′ such that

Algorithm 1 Reordering Vectors based on Balanced Signs [Harvey and Samadi [16]]

input: a group of signs { s j } N j =1 , initial order π initialize: two order-sensitive lists L pos ← [ ] , L neg ← [ ] . for j =1 ...N do L pos . append ( π ( j )) if s j is +1 else L neg . append ( π ( j )) . end for return: new order π ′ := concat ( L pos , reverse ( L neg )) .

This says that, with new permutation π ′ , 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 ˜ 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 π ∗ that solves the ˜ O (1) herding objective in (3).

GraB's application of herding to gradient balancing. Luet al. [24] applies this framework of herding andbalancing 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 ¯ 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

where w p t denotes the model weights after p -1 updates in the t -th epoch, and π t denotes the permutation adopted in the t -th epoch. Lu et al. [24] proves that this definition of z j preserves the

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:AProvablyEfficientOrderingAlgorithmforDistributed 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 w ∈ R d and maintains n = N/m local examples. 6 As in many data-parallel training applications, 7 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 where f i ( w ; j ): R d → R , j ∈ [ n ] , denotes the loss incurred on the j -th example on the i -th worker for model weights 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 π t,i : [ n ] → [ 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

That is, in epoch t , each worker i selects their respective, local j -th example according to { π t,i } n i =1 in order to compute stochastic gradients (Appendix).

Following this setup, Algorithm 1 no longer guarantees the ˜ O (1) boundtotheherdingproblem(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 α for greater scalability [40]; larger α increases the discrepancy between averaged gradients in adjacent epochs, which, in turn, would make GraB's use of stale means unreliable.

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 z i,j ∈ R d denote the vector indexed by j on the i -th worker. Assuming ‖ z i,j ‖ 2 ≤ 1 ( ∀ i ∈ [ m ] , ∀ j ∈ [ n ]) , the goal of parallel herding is to find m permutations, π 1 , π 2 ,..., π m where π i : [ n ] → [ n ] ( ∀ i ∈ [ m ]) , so as to minimize:

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 π i : [ n ] → [ n ] ( ∀ i ∈ [ 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 ...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 ˜ 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. Wecall 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 ˜ 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 ˜ 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 ˜ 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 ˜ O (1) . Wedefer 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.,

That is, we fit these ordered-paired differences { y i } mn/ 2 i =1 into the herding and balancing framework (Algorithm 1): if sign s is associated with y n ( k -1)+ i , then z i, 2 k -1 and z i, 2 k 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.

Algorithm 2 PairBalance

Worker 1

Worker 2

1. Examples

ordered with

Worker 1

Worker 2

2. Compute

gradients

-1

+1

-1

+1

Worker 1

Worker 2

7. Reorder

examples with

4. Call Pair

Balance

6. New Data

Permutations

Server

5. Call

Algorithm 1

3. Send Gradients

to server

Wedescribe CD-GraB at two levels of abstraction: a high-level illustration (Figure 1, steps 1-7 ) and a detailed pair of worker-server algorithm statements (Figure 2). Since the workers only have access to a 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 worker-specific permutations. In more detail:

In epoch t , (Figure 1, step 1 ) the two workers have permutations π t, 1 and π t, 2 , respectively. Each worker computes per-example gradients g i j ( 2 ; Algorithm 3:4), and sends them to the server ( 3 ; Algorithm 3:5). The server we implement functions as a parameter server [20]: It computes the average of the workers' per-example gradients (Algorithm 4:6), and sends it back to all workers (Algorithm 4:7) so that they can update their local models (Algorithm 3:6-7). Simultaneously, as the server receives gradients (Algorithm 4:5), it calls PairBalance (Algorithm 2) on adjacent vectors ( 4 ; Algorithm 4:4-13). PairBalance produces signs to supply to the reordering algorithm (Algorithm 1), which, using the current worker permutations π t,i , produces the new per-worker permutations for the next epoch ( 5 ; Algorithm 4:14). In Figure 1, these correspond to π t +1 , 1 and π t +1 , 2 , which the server then sends back to the respective workers ( 6 ; Algorithm 4:15). Lastly, before the start of the next epoch, the workers reorder their examples according to the new permutations ( 7 ; Algorithm 3:9).

4 Convergence Analysis

Wenext demonstrate formally that our CD-GraB algorithm (Section 3.3) confers the efficiency benefits of centralized GraB (Section 2.1) to the distributed setting. In brief, our main theoretical results show that CD-GraBenjoys a linear speedup in convergence rate under two sets of conditions: smoothness (Theorem 2) and the PolyakŁ ojasiewicz (P.L.) condition (Theorem 3). Both results guarantee that CD-GraBisfasterthandistributed random reshuffling (D-RR) . Our proofs rely on Corollary 7 from Dwivedi and Mackey [10], which shows that, with high probability, RandomizedBalance (subroutine in Algorithm 2, from Alweiss et al. [2]) guarantees a ˜ O (1) bound to the signed herding objective (4). 9

To begin, we restate this result to cohere with our framework, for which the vectors z j are gradients in an optimization context:

Theorem 1 ( Corollary 7, Dwivedi and Mackey [10] ) . Consider any vectors { z j } N j =1 ( z j ∈ R d ) with ‖ z j ‖ 2 ≤ 1 supplied as input to the RandomizedBalance subroutine in Algorithm 2. Then for any δ > 0 , with probability at least 1 -δ , RandomizedBalance outputs a sequence of signs { s j } N j =1 ∈ { -1 , 1 } that satisfy max k ∈ [ N ] ∥ ∥ ∥ ∑ k j =1 s j z j ∥ ∥ ∥ ∞ ≤ ˜ A , where ˜ A = √ 2log ( 4 d δ ) log ( 4 N δ ) = ˜ O (1) .

Algorithm 3 CD-GraB Workers

require:

m

workers,

n

:=

N

m

ex. per worker

input:

initial

w

1

1

, epochs

T

, learning rate

α

1:

receive:

initial permutations

2:

for

t

:=1

...T

do

/triangleright

Run in parallel for workers

i

=1

...m

3:

for

example

j

:=1

...n

do

4:

compute:

g

i

j

←∇

f

i

(

w

j

t

,

π

t,i

(

j

))

j

-th stochastic grad.

g

i

j

- - - - - - - - - - - - - - - - - - - - - - →

5:

send:

g

i

j

avg.

j

-th stochastic grad.

¯

g

j

←--- - - - - - - - - - - - - - - - - -

6:

receive:

¯

g

j

7:

update:

w

j

+1

t

w

j

t

-

α

¯

g

j

8:

end for

9:

receive:

next permutation

10:

update:

w

1

t

+1

:=

w

n

+1

t

11:

end for

12:

return:

w

T

+1

:=

w

1

T

+1

Algorithm 4 CD-GraB Parameter Server

require:

m

workers,

n

:=

N

m

ex. per worker

input:

epochs

T

{

π

1

,i

}

i m

=1

1:

send:

initial permutations

{

π

1

,i

}

i m

=1

2:

for

t

:=1

...T

do

3:

initialize:

running sum

h

=

0

; empty list

S

4:

for

example

j

:=1

...n

do

5:

receive:

{

g

i

j

}

i m

=1

from all workers

i

6:

compute:

avg. gradient:

¯

g

j

1

m

∑

m

i

=1

g

i

j

7:

send:

¯

g

j

to all the workers

8:

for

worker

i

:=1

...m

do

9:

if

j

mod 2=0

:

10:

h

,s

i

j

-

1

,s

i

j

PairBalance

(

h

,

g

i

j

-

1

,

g

i

j

)

11:

S

.

append

(

s

i

j

-

1

)

;

S

.

append

(

s

i

j

)

12:

end for

13:

end for

/triangleright

Call Alg. 1 for

i

=1

...m

on

π

t,i

and

S

14:

compute:

next permutations

{

π

t

+1

,i

}

i m

=1

π

t

+1

,i

15:

send:

{

π

t

+1

,i

}

i m

=1

to each worker

i

16:

end for

To integrate this result with our parallel setting, we need some additional assumptions that are standard in the literature on distributed optimization - that the variance of the per-example gradients on each worker is uniformly bounded (Assumption 1), and that the variance between worker-specific gradients is similarly bounded (Assumption 2). More precisely, following the distributed setup in (7), we denote the global loss gradient to be ∇ f ( w ) , each i -th worker's local loss gradient to be ∇ f i ( w ) ( ∀ i ∈ [ m ] ), and each i -th worker's per-example loss gradients to be ∇ f i ( w ; j ) ( ∀ j ∈ [ n ] ). We assume:

∥ ∥ Assumption 2 ( Bounded Data Heterogeneity ) . There exists a constant ς > 0 such that ∀ i ∈ [ m ] , ∥ ∥ ∇ f i ( w ) -∇ f ( w ) ∥ ∥ 2 2 ≤ ς 2 .

Assumption 1 ( Bounded Gradient Variance ) . For all i ∈ [ m ] there exists a constant σ > 0 such that for all j ∈ [ n ] and for all w ∈ R d , it holds that ∥ ∇ f i ( w ; j ) -∇ f i ( w ) ∥ 2 2 ≤ σ 2 .

Lastly, we include one additional assumption from the original GraB paper [24]: we assume a cross norm L 2 , ∞ (which can be easily adapted to L 2 -smoothness by setting L 2 , ∞ to be √ dL 2 ).

Assumption 3 ( Smoothness ) . There exists constant L 2 , ∞ > 0 such that for any w , v ∈ R d , any i ∈ [ m ] , and any j ∈ [ n ] , it holds that ∥ ∥ ∇ f i ( w ; j ) -∇ f i ( v ; j ) ∥ ∥ 2 ≤ L 2 , ∞ ‖ w -v ‖ ∞ . Given these assumptions, we can prove a convergence guarantee for CD-GraB:

Theorem 2. Suppose that Assumptions 1,2 and 3 hold. For any δ > 0 , if we set learning rate α to be where F 1 = f ( w 1 ) -inf w ∈ R d f ( w ) and ˜ A comes from Theorem 1. Then, with probability at least 1 -T δ ,

Wecan also prove an accelerated rate for CD-GraB if we additionally assume the P.L. condition:

Assumption 4 ( P.L. Condition ) . We say the loss function f fulfills the P .L. condition if there exists µ> 0 such that for any w ∈ R d , 1 2 ‖∇ f ( w ) ‖ 2 2 ≥ µ ( f ( w ) -inf v ∈ R d f ( v )) .

Theorem 3. Suppose that Assumptions 1, 2, 3, and 4 hold. For any δ > 0 , we set constants ˜ W and C 3 to be where ˜ A comes from Theorem 1, F 1 is from Theorem 2, and W 0 is the Lambert-W function. If we set learning rate α = 2 ˜ W Tnµ and if the number of epochs T satisfies

We prove Theorems 2 and 3 in the Appendix. Together, they show that CD-GraB exhibits a linear speedup in the number of workers m over GraB [24]'s convergence rates ( ˜ O (( nT ) -2 / 3 ) and ˜ O (( nT ) -2 ) , respectively). 10 under both smoothness and the P.L. condition. Further, CD-GraB's convergence rate of ˜ O (( mnT ) -2 ) is faster than many previous rates, 11 such as the high probability bound of ˜ O (( mn ) -1 T -2 ) for D-RR in Yun et al. [48].

5 CD-GraB in Practice: Distributed and Simulation Experiments

We next verify CD-GraB's accelerated convergence on a variety of empirical tasks. 12 For ease of comparison, we follow the experimental plan from the original GraB paper, 13 and add some additional large-scale logistic regression experiments. We also run an 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

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 3a), Long Short-Term Memory (LSTM) [17] on the WikiText-2 dataset [29] (Figure 3b), and autoregressive Multi-Layer Perceptron (MLP) on the M4 Weekly dataset [25] (Figure 3c). 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 (Figures 3a and 3b), we launch m =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 8 process workers. We use NCCL as the distributed communication backend [33] for the mortgage application and WikiText-2 tasks, and GLOO [1] as the distributed communication backend for the M4 task.

AsshowninFigure3,wecompareCD-GraB'sconvergencetothestandarddistributed-trainingexampleordering method: random reshuffling (D-RR). From all subfigures in Figure 3, we observe that CD-GraB outperforms the D-RR baseline significantly and consistently: CD-GraB exhibits better training loss and test metrics, measured against both the number of epochs and wall-clock time. We also note that the

(a) Logistic regression on mortgage application (NY 2017 subset) [7]

(b) LSTM [17] on WikiText-2 [29]

(c) Autoregressive MLP on M4 Weekly [25]

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 a 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 a 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 design 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(Section2.1), 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, is a side effect of our solution to develop CD-GraB, not our main contribution. Therefore, we run a simulation ablation study to disentangle the relative importance of each of CD-GraB's efficiency benefits over GraB. To do so, we compare the convergence of CD-GraB to two additional baselines in the distributed setting, beyond D-RR: (1) ID-GraB (Bal) , where each independent worker runs GraB locally using RandomizedBalance (subroutine in Algorithm 2) to perform gradient vector balancing; (2) ID-GraB (PairBal) , where each independent worker runs GraB locally using PairBalance .

Figure 4 summarizes the results, with convergence curves for m ∈ { 4 , 8 , 16 , 32 , 64 } workers training LeNet on CIFAR-10. We choose this task and architecture to cohere with the experiments done in the original GraB paper. For these experiments, we denote B to be the aggregated minibatch across all the workers, which refers to the number of stochastic examples used for an overall optimization step; each worker thus has a subset of this minibatch - an equivalently-sized subset of B examples. 15 We 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 m increases, the parallel herding bound (8) will increase linearly if there is no coordination. Second, as we scale up to larger m , the convergence curves of ID-GraB (Bal) and ID-GraB (PairBal) gradually approach

the curve for D-RR: at larger scales, herding-based example ordering will be no better than randomly permuting the dataset. Both observations give strong evidence that coordination (i.e., running online PairBalance on the server to coordinate per-worker permutations) is critical for accelerating training.

Wenote that all of these experiments use SGD, since both the theoretical results of the original GraB paper and our results for CD-GraB here are for SGD. In the Appendix, we additionally include results for training GPT-2 on WikiText-103, for which we use AdamW as the optimizer. We find that GraB with AdamW works in practice; however, our theory results do not directly apply to these experiments. Weadditionally include results on memory usage in the Appendix, which show that CD-GraB results in negligible overhead in practice.

6 Conclusion and Future Work: Toward an Order Server Architecture

Weelevate the benefits of provably faster, permutation-based example ordering to the contemporary MLdistributed-training setting. We focus on reformulating the online Gra dient B alancing algorithm (GraB) [24] because, even though it is the provably optimal permutation-based example-ordering method [6], it is limited by design to centralized settings (Section 3.1). To overcome these limitations, we redesign GraB'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 [3, 10, 11]. PairBalance operates on ordered pairs of vectors to do balancing , which enables our full-stack, low-overhead, C oordinated and D istributed online CD-GraB algorithm. We give a full specification of our online CD-GraB algorithm (Section 3.3), provide convergence rate guarantees regarding its speedups on both 1) smooth non-convex and 2) P.L. objectives (Section 4), and verify these speedups in practice on single-node distributed tasks and a simulated ablation study (Section 5).

Both our theory and experiments demonstrate that CD-GraB really shines when there are multiple training epochs (Appendix). This is another reason that we do not emphasize experiments involving finetuning pre-trained models like GPT-2, as fine-tuning can be achieved in 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, permutation-based example ordering more useful for base-model pre-training.

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 in such large-scale settings, future work should investigate moving beyond the single-node setup that we present. Notably, to train larger models, our results suggest a novel distributed training architecture. The ordering operation performed by the server (Algorithm 4) is not very latency sensitive; the server has the duration of the entire epoch t to compute the new permutations for the next, t +1 epoch. Given this relaxed latency requirement, and the success of our algorithmic results, it would be an exciting direction for future ML-systems research to invest in building an OrderServer architecture. Such an architecture, which could be composed with traditional parameter servers, would afford the scalability benefits of CD-GraB to a host of massive-scale ML applications.

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 a gift from SambaNova Systems.

References

Collective communications library with various primitives for multi-machine training, 2023. URL https://github.com/facebookincubator/gloo .

RyanAlweiss, Yang P Liu, and Mehtaab Sawhney. Discrepancy minimization via a self-balancing walk. In Proceedings of the 53rd Annual ACM SIGACT Symposium on Theory of Computing , pages 14-20, 2021.

Alessandro Barp, Carl-Johann Simon-Gabriel, Mark Girolami, and Lester Mackey. Targeted separation and convergence with kernel discrepancies. arXiv preprint arXiv:2209.12835 , 2022.

Dimitri P. Bertsekas. Incremental Gradient, Subgradient, and Proximal Methods for Convex Optimization: A Survey. In Optimization for Machine Learning . The MIT Press, 2011.

Léon Bottou. Stochastic gradient descent tricks. In Neural networks: Tricks of the trade , pages 421-436. Springer, 2012.

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.

A. Feder Cooper, Katherine Lee, Solon Barocas, Christopher De Sa, Siddhartha Sen, and Baobao Zhang. Is My Prediction Arbitrary? Measuring Self-Consistency in Fair Classification. arXiv preprint arXiv:2301.11562 , 2023.

Christopher De Sa. Random Reshuffling is Not Always Better. In Advances in Neural Information Processing Systems , 2020.

Jacob Devlin, Ming-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 1 (Long and Short Papers) , pages 4171-4186, 2019.

Raaz Dwivedi and Lester Mackey. Kernel thinning. arXiv preprint arXiv:2105.05842 , 2021.

Raaz Dwivedi and Lester Mackey. Generalized Kernel Thinning. In Tenth International Conference on Learning Representations , 2022.

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.

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.

Mert Gürbüzbalaban, Asu Ozdaglar, and Pablo A Parrilo. Why random reshuffling beats stochastic gradient descent. Mathematical Programming , 186(1):49-84, 2021.

Jeff Z. HaoChen and Suvrit Sra. Random Shuffling Beats SGD after Finite Epochs. In Proceedings of the International Conference on Machine Learning , volume 97, pages 2624-2633, 2019.

Nick Harvey and Samira Samadi. Near-Optimal Herding. In Proceedings of The 27th Conference on Learning Theory , volume 35, pages 1165-1182, 2014.

Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation , 9 (8):1735-1780, 1997.

Kun Huang, Xiao Li, Andre Milzarek, Shi Pu, and Junwen Qiu. Distributed Random Reshuffling over Networks. arXiv preprint arXiv:2112.15287 , 2021.

Hakan Inan, Khashayar Khosravi, and Richard Socher. Tying word vectors and word classifiers: Aloss framework for language modeling. arXiv preprint arXiv:1611.01462 , 2016.

MuLi, David G. Andersen, Jun Woo Park, Alexander J. Smola, Amr Ahmed, Vanja Josifovski, James Long, Eugene J. Shekita, and Bor-Yiing Su. Scaling Distributed Machine Learning with the Parameter Server. In Proceedings of the 11th USENIX Conference on Operating Systems Design and Implementation , OSDI'14, page 583-598, USA, 2014. USENIX Association. ISBN 9781931971164.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101 , 2017.

Yucheng Lu, Si Yi Meng, and Christopher De Sa. A General Analysis of Example-Selection for Stochastic Gradient Descent. In International Conference on Learning Representations , 2021.

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 7145-7155. PMLR, 2021.

Yucheng Lu, Wentao Guo, and Christopher De Sa. GraB: Finding Provably Better Data Permutations than Random Reshuffling. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. URL https://openreview.net/forum?id=nDemfqKHTpK .

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: https://doi.org/10.1016/j.ijforecast.2019.04.014. URL https://www.sciencedirect.com/science/article/pii/S0169207019301128 . M4 Competition.

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.

Tambet Matiisen, Avital Oliver, Taco Cohen, and John Schulman. Teacher-student curriculum learning. IEEE transactions on neural networks and learning systems , 31(9):3732-3740, 2019.

Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efficient learning of deep networks from decentralized data. In Artificial intelligence and statistics , pages 1273-1282. PMLR, 2017.

Stephen Merity, Nitish Shirish Keskar, and Richard Socher. Regularizing and optimizing lstm language models. In International Conference on Learning Representations , 2018.

Konstantin Mishchenko, Ahmed Khaled, and Peter Richtárik. Random Reshuffling: Simple Analysis with Vast Improvements. In Advances in Neural Information Processing Systems , 2020.

Amirkeivan Mohtashami, Sebastian Stich, and Martin Jaggi. Characterizing & Finding Good Data Orderings for Fast Convergence of Sequential Gradient Methods. arXiv preprint arXiv:2202.01838 , 2022.

Deanna Needell, Rachel Ward, and Nathan Srebro. Stochastic Gradient Descent, Weighted Sampling, and the Randomized Kaczmarz algorithm. In Advances in Neural Information Processing Systems , pages 1017-1025, 2014.

NVIDIA. NVIDIA Collective Communication Library, 2023. URL https://https: //developer.nvidia.com/nccl .

PyTorch Contributors. DataLoader API, 2023. URL https://pytorch.org/docs/stable/ data.html .

Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.

Shashank Rajput, Kangwook Lee, and Dimitris Papailiopoulos. Permutation-Based SGD: Is Random Optimal? In International Conference on Learning Representations , 2022.

Benjamin Recht and Christopher Ré. Toward a Noncommutative Arithmetic-geometric Mean Inequality: Conjectures, Case-studies, and Consequences. In Conference on Learning Theory , volume 23, pages 11.1-11.24, 2012.

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.07021 , 2022.

Mark Schmidt, Nicolas Le Roux, and Francis R. Bach. Minimizing finite sums with the stochastic average gradient. Mathematical Programming , 162(1-2):83-112, 2017.

Samuel L Smith, Pieter-Jan Kindermans, Chris Ying, and Quoc V Le. Don't Decay the Learning Rate, Increase the Batch Size. In International Conference on Learning Representations , 2018.

Petru Soviany, Radu Tudor Ionescu, Paolo Rota, and Nicu Sebe. Curriculum learning: A survey. International Journal of Computer Vision , pages 1-40, 2022.

Evangelos Spiliotis, Andreas Kouloumos, Vassilios Assimakopoulos, and Spyros Makridakis. Are forecasting competitions data representative of the reality? International Journal of Forecasting , 36(1):37-53, 2020.

Merity Stephen, Xiong Caiming, Bradbury James, and Richard Socher. Pointer sentinel mixture models. Proceedings of ICLR , 2017.

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.18653/v1/W18-5446. URL https://aclanthology.org/W18-5446 .

Max Welling. Herding dynamical weights to learn. In Proceedings of the 26th Annual International Conference on Machine Learning , pages 1121-1128, 2009.

Bicheng Ying, Kun Yuan, Stefan Vlaski, and Ali H. Sayed. On the performance of random reshuffling in stochastic learning. In 2017 Information Theory and Applications Workshop (ITA) , pages 1-5. IEEE, 2017.

Binhang Yuan, Yongjun He, Jared Davis, Tianyi Zhang, Tri Dao, Beidi Chen, Percy S Liang, Christopher Re, and Ce Zhang. Decentralized training of foundation models in heterogeneous environments. Advances in Neural Information Processing Systems , 35:25464-25477, 2022.

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.

Chulhee Yun, Suvrit Sra, and Ali Jadbabaie. Open Problem: Can Single-Shuffle SGD be Better than Reshuffling SGD and GD? In Conference on Learning Theory , 2021.

Peer reviews:
comments: SUMMARY:
This paper is devoted to the development of the Gradient Balance for the centralized distributed setting.
The 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.

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.

2. They prove that the convergence rate for CD-GraB exhibits a linear speedup over GraB, which is non-trivial.

3. 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.

2. The paper uses a different smoothness assumption. The authors seem to want to hide the degradation from the dimension $d$.

3. The technical difficulty of convergence analysis can be presented as roadmap for the reader.

QUESTIONS:
What is the main technical difficulty of convergence analysis for distributed GraB.

Minor: line 32: "Gradiant Balance algorithm"  should be "Gradient Balance algorithm"

LIMITATIONS:
Yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-24

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nPropose 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\n\nSTRENGTHS:\n**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.\n\nWEAKNESSES:\n**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?\n\nQUESTIONS:\nPlease see questions described in-line in weaknesses section.\n\nLIMITATIONS:\nLimitations of the work are not adequately addressed.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

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

The paper is very well and tightly written with clear limitations. The authors verified their claims both experimentally and theoretically.

STRENGTHS:
* The authors have clearly motivated the problem and discussed it in detail how their work compare to other prior works

* The experiments are well executed and clearly demonstrate the benefit of CD-GRAB. 

* The authors discussed the proof in detail in the appendix and organized it well.
 
* All limitations are discussed.

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 ?

* 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 ?

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1.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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nThe 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\nLIMITATIONS:\nNo.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nOne 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nNo 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-26"}