Asynchronous Local-SGD Training for Language Modeling

Authors: bo  liu
rachita  chhaparia
arthur  douillard
satyen  kale
andrei alex rusu
jiajun  shen
arthur  szlam
marcaurelio  ranzato
Venue: ICML_cc_2024_Workshop_WANT
Published: 2024-06-02 00:03:56

Abstract: Local stochastic gradient descent (Local-SGD), also referred to as federated averaging, is an approach to distributed optimization where each device performs more than one SGD update per communication.  
This work presents an empirical study of asynchronous Local-SGD for training language models; that is, each worker updates the global parameters as soon as it has finished its SGD steps.  We conduct a comprehensive investigation by examining how worker hardware heterogeneity, model size, number of workers, and optimizer could impact the learning performance.  We find that with naive implementations, asynchronous Local-SGD takes more iterations to converge than its synchronous counterpart despite updating the (global) model parameters more frequently. We identify momentum acceleration on the global parameters when worker gradients are stale as a key challenge.  We propose a novel method that utilizes a delayed Nesterov momentum update and adjusts the workers' local training steps based on their computation speed. This approach, evaluated with models up to 150M parameters on the C4 dataset, matches the performance of synchronous Local-SGD in terms of perplexity per update step,  and significantly surpasses it in terms of wall clock time.

Main paper: 
 introduction: 
 Large language models (LLMs) have revolutionized many applications, transforming the way machines interact with human language. The cornerstone of this revolution is training these models at massive scale. To manage such largescale training in reasonable amounts of time, it has been necessary to distribute computations across multiple devices. However, the standard approaches to this distributed training uses co-located devices with fast interconnects.One might hope to be able to effectively harness a broader range of computational resources, perhaps geographically distant from each other, to build even more powerful large models. However, utilizing numerous distant devices faces a significant hurdle: communication latency. When devices focus solely on computing gradients before sending them back to a central server, communication time may exceed computation time, creating an efficiency bottleneck.Local Stochastic Gradient Descent (Local-SGD) is a collection of optimization methods that can reduce communication bottlenecks. 1 1 These methods involve each device performing multiple local gradient steps before syncing their parameter updates with a parameter server. While Local-SGD enhances training efficiency by reducing communication frequency, it can suffer from the straggler effect caused by heterogeneous devices. For instance, faster devices are idle waiting for slower ones to catch up, undermining the overall efficiency of the system. Moreover, all devices are forced to communicate at the same time requiring high bandwidth connection with the parameter server. Asynchronous Local-SGD presents a more viable solution (illustrated in Figure 1 ), as it allows the server to update the model as soon as the updates of a worker are available, thereby enhancing computational utilization and minimizing communication bandwidth requirements. Figure 2 . Comparative evaluation of language models using sync. and async. Local-SGD methods with 4 heterogeneous workers on a 20M parameter model. The state-of-the-art sync. Local-SGD method, DiLoCo (Douillard et al., 2023) , employs AdamW and Nesterov momentum as the worker-side and server-side optimizers, respectively. This optimizer combination remains the strongest for async.Local-SGD training (See Figure 5 ), yet underperforms DiLoCo significantly. By integrating Delayed Nesterov (DN) (Algorithm 2) for outer optimization and Dynamic Local Updates (DyLU) (Section 5), we significantly bridge the performance gap in terms of perplexity versus updates between sync. and async. training in language modeling. Moreover, the proposed method significantly surpasses DiLoCo in terms of perplexity versus wall clock time.In this study, we explore the viability of asynchronously training LLMs using Local-SGD. We expand upon previous works that have attempted to alternate steps on subsets of workers or randomly drop certain subset of workers during synchronous Local-SGD (Ryabinin et al., 2021; Douillard et al., 2023) . The main content is structured into three parts:1. Framework (Section 3). The first part introduces our high-level design for the asynchronous training framework. We discuss how each worker determines which data shard to train on, for how many steps, with what learning rates, and how the server updates models asynchronously.2. Optimization Challenge (Section 4). In the second part, we conduct an empirical study of various existing optimization strategies suitable for asynchronous Local-SGD. This includes both worker-side optimization (inner optimization) and server-side optimization (outer optimization). We uncover a key challenge in utilizing momentum effectively. Notably, while adaptive momentum methods accelerate convergence of both inner and outer optimizations, their efficacy in asynchronous Local-SGD is comparatively reduced when both optimizations employ momentum techniques, when contrasted with the synchronous implementation.3. Proposed Solutions (Section 5). We introduce two simple and effective techniques: the Delayed Nesterov momentum update (DN) and Dynamic Local Updates (DyLU). These techniques, when combined and evaluated on training language model, allow asynchronous Local-SGD to approach synchronous Local-SGD in terms of perplexity versus the total number of local updates, and further improve asynchronous Local-SGD vs. synchronous Local-SGD in terms of perplexity versus wall-clock (Figure 2 ). 
background: 
 In this study, we focus on the distributed optimization of shared model parameters θ across k data shards, denoted as D = {D 1 , . . . , D k }, with k workers. 2 The primary goal is described by the following equation:EQUATIONwhere ℓ(•; θ) represents the loss function (for instance, crossentropy loss for next token prediction in language modeling), and | • | indicates the set size.We extend the definition of Local-SGD in this work to include not just the original Local-SGD method, but also its variants that incorporate advanced optimization techniques. We particularly focus on DiLoCo (Douillard et al., 2023) , which sets the standard for synchronous Local-SGD in language modeling. DiLoCo's methodology is detailed in Algorithm 3. Each worker i performs H local updates using an inner optimizer on their data shard D i before sending the parameter change (pseudo-gradient) δ(t) i = θ (t-1) -θ (t) iback to the server. The server then computes the aggregated outer gradient∆ (t) = 1 k k i=1 δ (t)i , and applies an outer op-timizer with ∆ (t) to update θ. A key insight from DiLoCo is the optimal use of AdamW and Nesterov Momentum as the best inner and outer optimizers, respectively. 
async. local-sgd framework: 
 This section outlines the asynchronous Local-SGD pipeline design, where we assume a central server controls all workers and asynchronously aggregates their updates.Data Shard Sampling Unlike in the federated learning setting where each device is attached to its own data, in distributed optimization, the user has the right to choose which data shard is assigned to which worker, even dynamically. To balance the learning progress on different data shards (as workers are heterogeneous), whenever a worker is ready to start a new local optimization round, we sample a data shard inversely proportional to its "learning progress". Specifically, define n i as the number of learned data points in D i , then we sample a shard i sampled according to:i sampled ∼ p,whereEQUATIONIn other words, we sample a data shard only when it is "under-sampled" (i.e., ni j nj ≤ |Di| j |Dj | ). The degree to which a shard is under-sampled determines its sampling rate. By doing so, we ensure that the data shard with slower progress is more likely to be sampled for training, therefore balancing the learning progress across shards.Learning Rate Scheduling In contrast to synchronous training methods like DiLoCo, asynchronous training can lead to uneven progress across different data shards, especially when workers are allowed varying numbers of training steps. This raises the question of how to effectively schedule learning rates. In our approach we assign each data shard its own learning rate schedule. Specifically, we implement a linear warmup combined with a cosine learning rate decay, where T represents the target total training iterations for each data shard:EQUATIONIn practice, asynchronous training may conclude with different final iteration counts (t end ) for each data shard. Since we cannot predetermine t end due to the unpredictability of asynchrony, we set the minimum learning rate (η min ) to a small positive value. This ensures continued progress even if t exceeds T . Additionally, we adjust T -t warmup to be nonnegative and ensure that the ratio t-twarmup T -twarmup remains within the range of [0, 1]. This helps maintain effective learning rate adjustments throughout the training process.Grace Period for Model Synchronization In asynchronous training, the completion time of each worker's tasks can vary. For example, if worker B completes training shortly after worker A, it might be beneficial for A to wait briefly until the server processes updates from both workers before receiving the updated model for its next training task. However, this waiting period should be minimal and occur only when necessary. Specifically, if no other worker completes its task within the grace period while worker A is synchronizing with the server's model, A should promptly commence its new training task using the server's current model. For a visual representation of this process, please refer to Figure 3 . 0) . Note that the train() function implements the data sampling technique and performs the learning rate scheduling mentioned before. In line 6, we reset the starting time of the grace period τ sync to ∞. This is because we want to synchronize with a worker only when its completion time is within τ sync + τ grace . The main asynchronous Local-SGD training loop is provided in line 6-19. Within the loop, we first attempt to get a completed worker w (line 7). We retrieve the earliest completed worker that we have not yet processed yet, as long as its completion time is still within the grace period (e.g., w.completed time ≤ τ sync + τ grace ).If no such workers exist, get worker() will return null. In line 10-15 where such a worker w is found, we synchronize its update with the server model θ. In line 17-20 when no such workers are found, we assign new training jobs for all completed workers and empty the list of completed workers.For the detailed pseudocode of the train() and get worker() functions, please refer to Appendix E. In practice, for the sake of reproducibility of research, we implement a determininistic version of Algorithm 1 with faked training time based on real-world device statistics. We validate the cor-Algorithm 1 Async. Local-SGD Task Scheduling.1: Require: Initial pretrained model θ (0) 2: Require: k workers 3: Require: Grace period τ grace 4: Require: Total local updates t max 5: t local = 0 6: θ ← θ (0) 7: W = [init worker() for i in [k]] 8: W completed = [] 9: train(W, θ) //τ sync = ∞ 23: train(W completed , θ) 24: W completed = [] 25:end if 26: end while rectness of the training pipeline by simulating synchronous updates using the asynchronous framework. 
optimization challenge: 
 Effect of InnerOpt + OuterOpt To study how optimization affects the language modeling performance in asynchronous Local-SGD, we first experiment with different combinations of the inner and outer optimizers (we use A+B to denote A and B as the inner and outer optimizer, respectively): SGD+Nesterov, SGD+Adam, AdamW+SGD, AdamW+SGD Momentum, AdamW+Adam, AdamW+Nesterov. The hyperparameters for each combination are tuned separately, for AdamW as InnerOpt we kept the default values. We assume there are k = 4 workers, whose device speed is shown in Figure 4 . Then we apply asynchronous Local-SGD finetuning on a 20M-parameter language model for 64,000 steps per worker (256,000 local steps in total), where the initial model checkpoint was pretrained for 24,000 steps with Adam without distributed training. We choose finetuning with Local-SGD as it has been observed that Local-SGD style methods work well in finetuning but is less efficient from scratch (Lin et al., 2018) , though others have also observed that Local-SGD works well even for training from scratch (Douillard et al., 2023) .The learning rate scheduling and task scheduling follow the procedures described in Section 3. We use inner steps = 50 across all workers in all experiments by default. The result is shown in Figure 5 . Observation The analysis reveals that combining AdamW as the inner optimizer with Nesterov momentum as the outer optimizer yields the best results, aligning with findings from synchronous training, like the DiLoCo method. Notably, using AdamW as the outer optimizer is less effective. This may be because AdamW, derived from Adam, introduces a normalization effect, which can be counterproductive in Local-SGD where pseudo-gradients tend to be larger than true gradients, potentially slowing convergence. When AdamW is used in the inner optimization, SGD, SGD Momentum, and Nesterov show comparable performance. However, Nesterov not only stabilizes the learning curve but also slightly improves final performance. This can be attributed to its update mechanism (here we abuse the notation and let t denote t server ):EQUATION)where ϵ is the learning rate, m t is the momentum, g t the gradient at time t, and β ∈ (0, 1) the decay factor. The key difference between Nesterov and SGD Momentum is in how Nesterov adjusts the weightings, reducing the momentum component (β 2 instead of β) and increasing the gradient component (1 + β instead of 1). This suggests that momentum plays a crucial yet intricate role in Local-SGD.To delve deeper into the momentum term's impact on the outer optimizer, we conducted comparative analyses between AdamW+SGD and AdamW+Nesterov under both synchronous and asynchronous training settings. These experiments were carried out under identical conditions as described earlier. The results are reported in Figure 6 . Observation The figure clearly shows that in asynchronous Local-SGD, AdamW+SGD, which lacks a momentum term, leads to better final perplexity and learning efficiency than its synchronous counterpart. However, incorporating Nesterov momentum into the OuterOpt significantly boosts the performance of synchronous Local-SGD, outperforming the asynchronous version. It's noteworthy that asynchronous AdamW+Nesterov remains the best performer across all tested combinations of inner and outer optimizers (as seen in Figure 5 ). This observation indicates that while momentum is beneficial in asynchronous Local-SGD for language modeling, its effect is more pronounced in synchronous settings.Is Staleness the Cause? We further apply the asynchronous DiLoCo algorithm with homogeneous devices. By doing so, we maximally diminish the staled gradient problem in Local-SGD, which refers to the fact that we are using an outdated outer gradient to update the server model.In particular, this means if we have k workers, all of them will return the computed outer gradient back to the server at the same time. Therefore, the only staleness comes from the fact that we are sequentially applying the individual updates instead of aggregating them together and apply it once. Results are summarized in Figure 7 .Observation Figure 7 reveals a notable finding: even with homogeneity among workers, asynchronous DiLoCo sig- nificantly lags behind its synchronous counterpart. This suggests that the inherent staleness from sequentially applying simultaneous updates leads to considerable performance drops. To elucidate this effect, let's consider a scenario with k = 4 workers providing identical outer gradients (denoted as g). The standard Nesterov momentum update is outlined in Equation ( 4). In a sequential application of pseudo-gradients:EQUATIONFrom this, we observe that sequential application results in a more rapidly decaying momentum term but amplifies the actual change in parameter θ. Consequently, a higher β maintains more recent momentum but may lead to greater changes in parameters, and vice versa. Note this imbalance cannot be easily removed by reducing the learning rate.Baselines We consider several synchronous baselines from the literature, and their naive application to an asynchronous setting: 1) Finetune 1 worker (4xbatch): This involves finetuning a single worker with a larger batch size, equating to synchronous SGD. 2) DiLoCo (Douillard et al., 2023) : This synchronous Local-SGD method combines AdamW with Nesterov. 3) Async. DiLoCo: The asynchronous version of DiLoCo.Existing Fixes We investigated potential fixes from the asynchronous Local-SGD literature to address observed challenges. The following methods were considered: 1) Async. DiLoCo + Poly (Xie et al., 2019) : Extends Async. DiLoCo by downweighting the pseudo-gradient with g ← (1 + staleness) -0.5 g. 2) Async. DiLoCo + PolyThres: Adds a threshold to discard updates with staleness beyond 10.3) Async. DiLoCo + Delay Comp. (Zheng et al., 2017) : Introduces delay compensation (Delay Comp.) to approximate true pseudo-gradients. Denote the gradient function at θ as g(θ), then the main idea of delay compensation is to approximate the true gradient g(θ t ) by a stale gradient g(θ t ′ ) (t ′ < t) with the first-order Taylor approximation, e.g., g(θ t ) ≈ g(θ t ′ ) + ∇g(θ t ′ )(θ t -θ t ′ ). In practice, the Hessian ∇g(θ t ′ ) is approximated by diagonalized gradient outer product, e.g., ∇g(θ t ′ ) ≈ λg(θ t ′ ) ⊙ g(θ t ′ ), where λ ∈ R + is a hyperparameter. In our setting, we apply the delay compensation technique to pseudogradients instead of gradients. 4) Async. Buffer: Accumulates and averages all gradients in a First-In, First-Out fashion before applying Nesterov updates; a variation of the original FedBuff algorithm (Nguyen et al., 2022) , using AdamW+Nesterov.The results are provided in Figure 8 . Observation Polynomial discounting of the pseudogradient shows marginal benefits. Thresholding and delay compensation techniques don't offer much improvements. Again, the fact that delay compensation is not working well points out the difference between asynchronous SGD and asynchronous Local-SGD. The Async. Buffer method excels at convergence but exhibits instability early in training. Crucially, none of the methods match the performance of the synchronous DiLoCo baseline.To facilitate future research and expedite the prototyping of novel ideas, we provide a minimal toy example (a one-file Python code), that reproduces the optimization challenge observed in asynchronous Local-SGD (refer to Appendix C) 
proposed solutions: 
 In addressing the optimization challenges outlined in Section 4, we developed two strategies. // default to c = 0 5: Require: Buffer size N 6: t = 0 7: m 0 = 0 // momentum 8: ∆ = 0 // aggregated gradient 9: while not finished do 10:Receive the pseudo-gradient g t 11:// sync. step in Alg. 1.12:∆ ← ∆ + g t 13:if (t + 1) % N == 0 then 14: 8 ). Additionally, our analysis revealed that asynchronous training with AdamW+SGD, sans outer momentum, outperforms synchronous methods (Figure 5 ). Based on these insights, we propose the Delayed Nesterov (DN) strategy, which represents the sync() function in Algorithm 1. This approach involves using the Nesterov update intermittently-every N server updates. Between Nesterov updates, we aggregate pseudo-gradients in a buffer ∆ and update the model parameters using gradient descent (or gradient descent plus a small fraction of the old momentum). To balance gradient and momentum-based descent, we introduce a parameter c ∈ [0, 1/N ]. A c value of 0 indicates pure gradient descent between Nesterov updates, while c equal to 1 evenly distributes the momentum term over N updates. The specifics of this algorithm are detailed in Algorithm 2. Unlike the Async. Buffer (Nguyen et al., 2022) , which updates model parameters only once in N periods, the Delayed Nesterov continuously updates using gradients, incorporating a fraction of the old momentum, and updates the momentum term once every N server updates.m t+1 ← βm t + ∆/N 15: θ t+1 ← θ t -ϵ (1 -cN + c)βm t+1 + g t /The Delayed Nesterov approach addresses the momentum challenge in the OuterOpt by buffering pseudo-gradients and strategically delaying momentum updates. An alternative perspective considers synchronous training as a solution, where pseudo-gradients from all workers are synchronized. However, the diversity in device capabilities often hinders simultaneous pseudogradient returns, if each worker executes the same number of local training steps. A viable workaround involves customizing local training steps (e.g., w.steps) based on the processing speed of each device. In particular, denote v(w) as the training speed (in terms of steps per second) for worker w, we can compute a worker's desired training steps as:EQUATIONwhere H denotes the number of local training steps the fastest worker runs and ⌊x⌋ denotes the largest integer not greater than x. 3 We name this approach the Dynamic Local Updates (DyLU). This adjustment allows slower workers to execute fewer steps, aligning the completion times across different workers. Incorporating a grace period for model synchronization in this setup further reduces the impact of stale gradients, improving overall performance. 
experiments: 
 This section details experiments conducted to assess the efficacy of our two proposed methods, Delayed Nesterov (DN) and Dynamic Local Updates (DyLU). Additionally, ablation studies explore the effectiveness of these methods as we vary the number of workers and model sizes.Evaluating Delayed Nesterov (DN) + Dynamic Local Updates (DyLU) Figure 2 compares async. Local-SGD with DN+DyLU against baselines such as single worker finetuning and DiLoCo, using the same setup as in Figure 8 .Observation The results demonstrate that DN combined with DyLU significantly reduces perplexity, surpassing the synchronous DiLoCo's performance over updates. Additionally, DN+DyLU outperforms in terms of time efficiency, avoiding delays from slower workers.We examine how both the proposed DN+DyLU method and vanilla asynchronous DiLoCo fare under varying degrees of worker device heterogeneity, as shown in Figure 9 (perplexity curve) and Table 1 (final perplexity). Observation DN+DyLU consistently excels across all heterogeneity levels. 4 Interestingly, even with homogeneous devices, vanilla asynchronous DiLoCo struggles, suggesting that the issue partly lies in the sequential application of pseudogradients. This underscores the importance of delayed momentum updates. Additionally, a periodic oscillation in performance is observed in certain device groupings, further highlighting the lack of robustness of the original asynchronous algorithm. Ablation with Different Model Sizes Lastly, we apply DN+DyLU to models of varying sizes (20M, 60M, 150M), with results summarized in Figure 11 (perplexity curve) and Table 3 (final perplexity). Observation Both synchronous and asynchronous Local-SGD methods outperform the approach of finetuning a single worker with an increased batch size. Notably, this advantage becomes more pronounced during the later stages of convergence, aligning with findings from previous research that highlight Local-SGD's superior generalization capabilities (Gu et al., 2023) . Additionally, our proposed DN+DyLU method demonstrates consistent efficacy across various model sizes. It's important to note that the performance disparity between synchronous and asynchronous DiLoCo does not diminish even as the model size increases.We apply c ∈ {0, 0.1} in Async. DN+DyLU with varying k (4, 8, 16) and model sizes (20M, 60M, 150M), with the 4 "very" heterogeneous workers. This is because when the level of heterogeneity is small, using different c will have smaller difference (e.g., when there is no heterogeneity, any c results in the same algorithm). Results are summarized in Table 4 . Observation Empirically, we observe no significant difference between c = 0 and c = 0.1, indicating that adding slight momentum at intermediate steps does not help too much. As a result, we set c = 0 as the default value in Algorithm 2, which corresponds to performing SGD updates between two consecutive Nesterov updates. Note that setting the value of c does not introduce any overhead to the overall algorithm. 
related work: 
 This section provides a concise overview of the literature on federated learning and local-SGD style distributed optimization, particularly focusing on the asynchronous settings.Local-SGD and Distributed Optimization Local-SGD is a specific distributed optimization technique designed to reduce communication frequency (Stich, 2018; Zhang et al., 2016; Bijral et al., 2016; McDonald et al., 2010; Coppola, 2015; Zinkevich et al., 2010) . The core principle of Local-SGD is to let each worker execute several local training iterations before engaging in global synchronization. This technique was later applied to the federated learning setting, leading to the development of the FedAvg method (McMahan et al., 2017) , which aims to reduce communication overhead. Unlike Local-SGD, federated learning also addresses user privacy issues and typically involves heterogeneous devices. To further minimize communication overhead, Fe-dOpt integrates adaptive optimization methods like SGD momentum and Adam (Reddi et al., 2020) . However, as worker heterogeneity increases, the convergence rate often deteriorates. Methods like SCAFFOLD (Karimireddy et al., 2020) and MIME (Karimireddy et al., 2021) are introduced to adapt these methods for heterogeneous environments.Asynchronous training was developed to mitigate the "straggler effect" observed in synchronous distributed optimization, where learning efficiency is bottlenecked by the slowest worker (Koh et al., 2006; Recht et al., 2011; Dean et al., 2012; Lian et al., 2015; 2018; Diskin et al., 2021b) . A significant challenge in asynchronous optimization is the staled gradient problem, which occurs when an outdated gradient is applied to a recently updated model. Asynchronous SGD with delay compensation (Zheng et al., 2017) addresses this issue by approximating the true gradient using the old gradient. Asynchronous methods have also been explored in federated learning contexts (Xie et al., 2019) . Despite the challenge, asynchronous training has demonstrated success for language modeling as well (Diskin et al., 2021b) , by using devices world-wide.Local-SGD for Language Modeling The concept of local-SGD (or FedAvg) has previously been applied in the realm of language modeling. Cross-device federated learning, for instance, has been utilized to pretrain and finetune language models (Hilmkil et al., 2021; Ro et al., 2022; Ryabinin et al., 2021; Diskin et al., 2021a; Presser, 2020; Borzunov et al., 2022) . More recently, DiLoCo has extended the local-SGD methodology to encompass larger language models, specifically proposing the use of AdamW + Nesterov momentum as the InnerOpt + OuterOpt pairing. In asynchronous settings, the FedBuff (Nguyen et al., 2022) algorithm buffers pseudogradients from clients, updating the server model only after accumulating a sufficient number of pseudogradients. TimelyFL (Zhang et al., 2023) aims to reduce asynchrony by allowing slower devices to train only parts of the model. 
conclusions and limitations: 
 This study examines asynchronous Local-SGD in language modeling, identifying that while momentum is key in the outer optimization loop, its effectiveness decreases in asynchronous settings compared to synchronous ones when applied simplistically. To address this, we propose a novel method using sporadic momentum updates with buffered pseudogradients alongside continuous stochastic pseudogradient updates. We find that adjusting local training steps to match each worker's computational speed significantly boosts performance. This study, while thorough, has several limitations. The exact cause of the optimization challenge is uncertain, highlighting a need for further theoretical research. Additionally, the effectiveness of the Local-SGD method appears to decrease as the number of workers increases, an issue that affects its scalability and requires more understanding. Lastly, despite DN+DyLU's enhanced performance, it lacks a theoretical convergence analysis, warranting additional exploration.A. Synchronous Local-SGD (DiLoCo (Douillard et al., 2023) ) PseudocodeHere, we provide the pseudocode for the synchronous Local-SGD training pipeline in Algorithm 3, using the optimization suggested by DiLoCo (Douillard et al., 2023) . θ (t) ← OuterOpt(θ (t-1) , ∆ (t) ) 17: end for 
 
 Appendix: We provide the plots of the performance (perplexity) of various methods for different ablation studies mentioned in Section 6 in the following. For the convenience of future research and quick prototyping of new ideas, we present a minimal toy example that replicates the observed optimization challenge in asynchronous Local-SGD (See Figure ??). 5 Network Architecture We displayed in Table 6 the architectural difference between the 20M, 60M, and 150M models. They are all transformer decoder-only, based on the Chinchilla family (Hoffmann et al., 2022) .Training Dataset We consider a language modeling task on the C4 dataset, a dataset derived from Common Crawl (Raffel et al., 2020) . The total number of steps is set to 88,000 for all models, with 24,000 steps of pre-training done without any federated learning methods, akin to post Local-SGD (Lin et al., 2020) .Hyperparameters In Table 5 , we outline the optimization hyperparameters considered for this study.Inner Optimizer States Following Douillard et al. (2023) , in all experiments, when worker B picks up the data shard worker A just finishes training on, we reset the AdamW's optimizer state. In other words, each local worker-side training is an independent training process with a new optimizer, and only the learning rate is adjusted according as described in Section 3.In this section, we provide the pseudocode for the train() and get worker() functions in Algorithm 1.Algorithm 4 train() in Algorithm 1. 1: Require: Available workers W 2: Require: Current server model θ 3: for w ∈ W do 4:Sample shard D ′ for w (Eq. 2).w.local updates = DyLU(D ′ ) (Eq. 6).Decide lr schedule (w.lr) (Eq. 3).w.update = train worker(w, D ′ , θ). 8: end forThe University of Texas at AustinGoogle DeepMind. Correspondence to: Bo Liu <bliu@cs.utexas.edu>.The term Local-SGD, sometimes also known as Federated Average (FedAvg), is used here to emphasize its roots in distributed optimization, where users have control over data allocation to different workers.We assume the number of workers (k) equals the number of data shards, though our methods are also applicable when there are fewer workers than data shards.We implicitly assume the device speeds are given. Otherwise, it is straightforward to estimate the device speed empirically.We notice that Async. DN+DyLU performs slightly better than DiLoCo when there is no heterogeneity, this is due to the numerical error, as the training curves match almost perfectly. 


Decision: We thank the authors for their time and contribution to WANT and we are pleased to share that after the reviewing process the paper has been accepted. Congratulations! We encourage the authors to consider reviewers' feedback for the improvement of the camera-ready version. We hope to see you in person at the workshop and brainstorm on efficient training research together!