EDiT: A Local-SGD-Based Efficient Distributed Training Method for Large Language Models

Response to Reviewer eSB4 (Part 1/2)

Dear Reviewer, We appreciate your thorough evaluation of our paper. Below, we address your concerns and provide our responses. All revisions will be included in the forthcoming rebuttal version of the paper. **Q1: The novelty of this paper is limited.** A1: Thank you very much for your questions. Our responses are as follows: 1. The "modern distributed strategies" mentioned here refer to strategies such as Fully Sharded Data Parallel (FSDP), ZeRO, or 3D parallelism used for training LLMs. We also realized that the term "modern distributed strategies" is vague. Therefore, we have revised and specified it as "model sharding strategy," which refers to that shard model parameters to prevent out of memory. This terminology is more appropriate for the methods discussed in our paper and avoids potential confusion. 2. Although eliminating significant anomalous workers, averaging and clipping gradients are not new to deep learning community, to our knowledge, we are the first to propose using these strategies within the context of Local SGD to mitigate training instability caused by highly diverse data and worker variability. These strategies integrate elegantly with the Local SGD method, and we also provide experimental evidence to show that they are necessary when training LLMs with complex data. In this paper, we introduce EDiT, an innovative efficient distributed method that combines a tailored Local SGD approach with model sharding techniques to enhance large-scale training efficiency. EDiT performs layer-wise parameter synchronization during forward pass and employs a prefetch strategy to overlap computation and communication, thereby reducing the additional communication and GPU memory overhead introduced by parameter synchronization. Besides, EDiT employs a novel pseudo gradient penalty strategy to suppress loss spikes caused by suboptimal corpus across workers, which ensures training stability and improve performance. Additionally, we propose an asynchronous variant of the EDiT method named A-EDiT to deal with the consistent stragglers in heterogeneous clusters. We also provide a large-scale verification of asynchronous pre-training for LLMs, along with an extensive analysis of convergence, generalization, acceleration, scalability, and stability. Based on these points, we believe that our work is indeed novel. **Q2: It may need to wait for the backward gradients, leading to significant idle time and slowing down the training.** A2: Thank you for your question. In the model shard group, the model is sharded horizontally rather than vertically (i.e., pipeline parallelism). Therefore, there is no issue of one worker needing to wait for another worker to return gradients. All workers in a model shard group compute the gradients for a given layer simultaneously. **Q3: The improvement over the existing work since quite insignificant as shown in Figure 4.** A3: Thank you very much for your question. 1. In our experiments, we compared EDiT with the Baseline and a series of classical or state-of-the-art methods, including Post Local SGD, DiLoCo, and CO2. It is worth noting that most methods that introduce periodic synchronization to improve training throughput has the problem of model performance decline. However, our proposed EDiT method not only improves training throughput but also achieves training and testing results that are consistent with or even better than the Baseline, which convincingly demonstrates the effectiveness of our approach. 2. Regarding the work [Federated learning with buffered asynchronous aggregation](https://arxiv.org/abs/2106.06639) that you mentioned, we have cited it in the related work section. This work is based on a Parameter-Server architecture for federated learning that is not suitable for training LLMs. EDiT and other compared methods all adopt an All-Reduce architecture. The compared methods in our work are consistent with the state-of-the-art work CO2/DiLoCo. **Q4: No theoretical results on the convergence of the proposed approach.** A4: Thank you very much for your constructive suggestion. We have added the convergence analysis for our proposed EDiT method. For details, please refer to "Section 5: Theoretical Analysis" in the main text and "Section A.4: Proof of Theorem 1" in the appendix.

Response to Reviewer eSB4 (Part 2/2)

**Q5: The proposed hierarchical distributed training method on a 2-D device mesh should be better explained.** A5: The 2-D device mesh is composed of model shard groups and model sync groups, which can be understood as the row and column dimensions, respectively. 1. In each model shard group, model parameters are evenly sharded among all workers. All-gather communication is performed among these workers to aggregate the parameters before forward and backward computations. After the backward computation, reduce-scatter communication is conducted to synchronize gradients among these workers. 2. In each model synchronization group, all workers maintain the same parameter shards and periodically synchronize parameters. This involves performing All-Reduce communication to average the parameters, followed by computing pseudo gradients and updating parameters with the outer optimizer. Due to the frequent communication required in model shard groups, efficient communication links are necessary. Normally, all workers within a single node form a communication shard group. In contrast, model sync groups require less frequent communication, so high-efficiency communication links are not essential. Typically, all workers with the same rank across different nodes are gathered to form a model sync group. We have updated Figure 1, Section 1 Paragraph 3, Section 3.1, and Algorithm 1 in the main text for a better presentation of the idea and we hope this would remove the confusion you have. **Q6: Is there any design in the proposed method addressing the concern that each worker within model shard group may need to wait for the backward gradients of each batch from the next worker in the same group?** A6: Referring to the response A2: In the model shard group, the model is sharded horizontally rather than vertically (i.e., pipeline parallelism). Therefore, there is no issue of one worker needing to wait for another worker to return gradients. All workers compute the gradients for a given layer simultaneously. **Q7: How this proposed hierarchical distributed training method accelerates the distributed training? It should be evaluated in the numerical results.** A7: Thank you for pointing this out. The hierarchical distributed training method we propose leverages the efficient communication links within the model shard groups and the low communication frequency within the model sync groups to achieve training acceleration. Besides, the proposed layer-wise sync method further reduces the time of model synchronization. In Section 4.3, we have added a speed comparison for different methods, and the numerical results demonstrate the acceleration effect of our method. **Q8: In Section 4.5, it seems suggesting that DiLoCo scheme with the pseudo gradient penalty mechanism is the proposed scheme EDiT?** A8: Thank you for pointing this out. Indeed, we incorporated a pseudo gradient penalty strategy into our method to enhance training stability and model performance. However, our main contribution lies in integrating Local SGD with model sharding strategy for training LLMs. Unlike DiLoCo and other methods that globally synchronize all parameters after the optimizer updates, EDiT synchronizes the sharded parameters layer by layer during the forward pass, addressing the memory and communication issues present in previous methods. To avoid any misunderstanding, we have revised the statement in Section 4.5 to "EDiT w/o ALL." In fact, our ablation experiments were indeed conducted based on the EDiT framework.

Response to Reviewer zyrD (Part 1/4)

Dear Reviewer, Thank you for your insightful and constructive feedback, which has significantly contributed to the improvement of our paper. Bellow, we address your concerns in detail. All the revisions will be included in the forthcoming rebuttal version of the paper. **Q1: Claiming novelty on the distributed implementation of a local SGD algorithm seems to strong.** A1 (Part 1): Thank you very much for your helpful comment. As you mentioned, the sharding strategies and implementations of the ZeRO series/FSDP themselves are not novel. However, combining them with Local SGD methods, especially effectively integrating model sharding into Local SGD, had not been explicitly explored prior to our work. Moreover, EDiT is a method specifically designed for these distributed training frameworks (ZeRO/FSDP), leveraging the characteristics of both Local SGD and the distributed frameworks, and thus go beyond a straightforward combination of ZeRO and LocalSGD. Specifically, most previous Local SGD methods synchronize all parameters globally after optimizer updates, leading to remarkable communication overhead. Additionally, when using outer momentum, previous methods introduce extra memory usage (twice the size of the model), and CO2, due to its one-step-delay strategy, incurs even greater memory usage (four times the size of the model), making them less preferrable to scale to larger LLMs. In contrast, EDiT first proposes layer-wise synchronization of sharded parameters during forward computation, achieving the following advantages: 1. Significantly reduced communication and data transfer burden. For instance, when training the 32-layer 1B model in the paper on 8 A100 GPUs, EDiT's communication volume per step is $1/(8 \times 32)$ of other methods. Furthermore, EDiT only requires communication within each model sync group, making it more efficient than previous methods' global communication. 2. Greatly reduced extra memory overhead. In EDiT, each GPU only needs to store sharded extra parameters and outer momentum, which can be further saved to CPU memory, incurring no additional memory overhead. In the layer-by-layer synchronization mode, CPU-GPU data transfer can also be overlapped with computation. 3. The prefetch strategy naturally achieves overlap of computation and communication. Using prefetch, EDiT can synchronize and aggregate the next layer’s parameters while computing the current layer, thus overlapping computation with communication and GPU-CPU data transfer. We have supplemented speed comparison experiments in the paper. As shown in Table 2 and Figure 9, EDiT achieves throughput comparable to CO2 while handling larger models. The profiling results in Figure 9 also demonstrate that EDiT can achieve nearly 100% overlap of computation and communication when training the 1B model. EDiT better utilizes the characteristics of current distributed frameworks, providing a new solution for model synchronization, which we believe is novel.

Response to Reviewer zyrD (Part 2/4)

A1 (Part 2): Regarding SlowMo and CO2, as you pointed out, both mention the use of advanced distributed frameworks, and CO2 specifically mentions integrating ZeRO series optimizers to adapt to large model training. However, with detailed study of their papers, codes, and author responses, we found the following issues: 1. The open-source codes of CO2 are primarily modified based on the open-source SlowMo codes. First, both of them are variants to FairScale's DataParallelDistributed (DDP) implementation, where each GPU stores the complete gradients, optimizer states, and model parameters, facing large memory burden. The extra parameters and outer momentum further increase memory pressure, making them unsuitable for training LLMs. Although they provide a memory-efficient option, this operation can only shards extra parameters and outer momentum, and does not address the memory usage of the model, optimizer, and gradients. Meanwhile, this operation introduces extra communication overhead. EDiT’s 2D device mesh combines an efficient model sharding strategy with low-frequency communication method to take advantage of the interconnect topology. This addresses the problem of memory/communication efficiency and enlarge the local worker batch size at the same time. While methods like SlowMo/CO2's hierarchical local sgd variants can only solve the latter problem. 2. We failed to find any open sourced implementation of the aclaimed CO2-ZeRO integration. From our perspective, CO2 is not readily compatible with ZeRO. In the cases of ZeRO-1 and ZeRO-2, the gradients/optimizer states are sharded, with each GPU responsible for updating only its corresponding portion of parameters. If the parameters are synchronized after each internal update, the overall training process degenerates into a synchronized mode, in which case CO2 will only save the gradient synchronization communication and lose the advantage of Local SGD in alleviating issues such as stragglers. However, if the model is not synchronized after each internal update, meaning each GPU updates the corresponding parameters based only on local sharded gradients, it may lead to unstable training. In the case of ZeRO-3, the authors of CO2 mentioned that in an [OpenReview response](https://openreview.net/forum?id=ZO5cn4IfaN&noteId=ynHafwj2OM): "As a result, communication operations for CO2 with ZeRO-3 involve only two all-gather collective communication primitives at the beginning of the forward and backward passes, respectively, which differs from the original implementation of ZeRO-3." In this scenario, the two all-gather operations introduced by CO2 also cause training to degenerate into a synchronized mode, eliminating the performance gain. In summary, we believe that existing methods like SlowMo and CO2 are not designed for large scale distributed training frameworks and thus do not integrate well with them, while our method offers significant improvements. Thus in the main text we have put a detailed explanation of the fine integration of EDiT with the distributed training framework to highlight our method's improvement, and provide better transparency to readers less familiar with large scale distributed training. We also acknowledge that our original paper had some misinterpretations and lacked some related experiments, so we have made the following revisions to the paper: 1. Emphasized the integration of our method with model sharding strategies, as well as the superiority of our approach in synchronizing sharded parameters layer by layer during forward computation. 2. Added more introductions to CO2 and SlowMo in Section 2 (Related Work), highlighting the distinctions and advantages of our method. 3. Supplemented Section 4.3 and A.3.2 with speed comparisons and profiling results of different methods, quantitatively and intuitively demonstrating the advantages of EDiT.

Response to Reviewer zyrD (Part 3/4)

**Q2: No comparison with CO2.** A2: Thank you very much for your suggestion. We have added comparisons with CO2, including convergence, generalization, and acceleration. The experimental results can be found in Sections 4.2, 4.3, and A.3.2. **Q3：The impact of introduced hyper-parameters is not discussed.** A3: Thank you for your question. In this paper, $\alpha$ and $\delta$ are hyperparameters used in anomaly elimination, while $\phi$ is the hyperparameter to control the pseudo gradient clip. A larger $\alpha$ means the moving average is more closely aligned with the current situation, a larger $\delta$ means a lower probability to identify a worker as an outlier, and a larger $\phi$ means looser constraint on pseudo gradient clip. These three parameters only take effect when abnormal pseudo gradients appear, so they minimally impact the normal training process and convergence trend. Therefore, their settings are not strict. 1. Through extensive experiments, we found that $\alpha=0.02$ yields good experimental results. Therefore, this setting can be retained for other tasks as well. 2. It is generally considered that a point can be identified as an outlier if it deviates by three standard deviations from the mean in z-test. Hence, $\delta$ is typically set to 3. 3. Similar to gradient clipping in synchronous training, the setting of $\phi$ usually depends on the specific training task. When the training task is unstable, a smaller $\phi$ can be set to stabilize the training. When the training task is relatively stable, a larger $\phi$ can be set to accelerate convergence. In our experiments, we set $\phi=10$. For other training tasks, we recommend testing $\phi=1$ and $\phi=10$, respectively. **Q4: Lack of clarity and imprecision in the writing.** A4: We sincerely appreciate your thorough and detailed review of our paper and for pointing out the issues. Below are our responses to each of the questions you raised. **Q4.1: N models were data parallel sharded model replicas and not model-parallel splits.** A4.1: We have added a note in the third paragraph of the Introduction section stating that all workers operate in data parallelism. Additionally, we have included a schematic diagram in Figure 1 to illustrate the model sharding within each worker. Together with the detailed explanation of the models in the model shard groups and the model sync groups at the beginning of Section 3.1, and the Algorithm 1 in appendix, we hope to remove the confusion you raised. Besides, as mentioned in the last paragraph of Section 3.1, our method can also be combined with 3D parallelism, i.e., model parallelism, and the corresponding code will be open-sourced upon the acceptance of the paper. **Q4.2: line 21: "ensure training stability and improve performance." this is a bit cryptic.** A4.2: We have revised the expression in line 21 to: "Besides, EDiT employs a pseudo gradient penalty strategy to suppress loss spikes, which ensures training stability and improves performance.", making it easier for readers to understand how the pseudo gradient penalty strategy takes effect. Additionally, in the second and third paragraphs of the Introduction, we have stated "low-quality datasets may introduce instability" and "the pseudo gradient penalty strategy can address this instability issue". **Q4.3: line 50: "Current Local SGD methods do not integrate well with modern distributed strategies". Why?** A4.3: We have revised the statement to a clearer and less ambiguous one: "Existing Local SGD methods do not handle model sharding well". Additionally, we have added a paragraph in Section 2 (Related Work) to explain this point. In summary, most current methods are based on the DDP architecture, where a complete copy of the model is stored on each GPU, and an additional global communication step is used to synchronize all model parameters, resulting in significant memory and communication overhead. Although CO2 mentions that it can be combined with ZeRO3, this approach involves using all-gather to aggregate parameters before each forward and backward pass, making it revert to a synchronous mode. For a more detailed explanation, please refer to our previous response in Answer A1. **Q4.4: lines 73-75: the problem described is exactly the one the CO2 paper aims at solving, so saying "current local-SGD method" seems to be an overstatement here.** A4.4: We have changed "current" to "most" to prevent overstatement. **Q4.5: line 116: Can you clarify in which way DiLoCo extends SlowMo in this context?** A4.5: As you pointed out, the original statement was not accurate. We have revised the wording to emphasize that the key difference between DiLoCo and SLowMo is that DiLoCo demonstrates the suitability of the Nesterov optimizer as an outer optimizer.

Response to Reviewer zyrD (Part 4/4)

**Q4.6: line 170: What does "model synchronization" mean?** A4.6: The term "model synchronization" here refers to the synchronization of parameters among workers within each model sync group, as already stated in the same paragraph: "parameters are synced in model sync groups, as outlined in lines 7 to 9 in Algorithm 1." For detailed procedures, please refer to Section 3.2, Figure 2, and Algorithm 2. **Q5: Intra node for model shards: What size of model can be loaded & how many intra-node model replicas does it lead to on standard GPU clusters?** A5: In a typical 8 GPU node, we treat all GPUs in this node as a model shard group because these GPUs have highest bandwidth interconnectivitiy, meaning a model is sharded across 8 GPUs. Considering the sequence length is 4096, the vocabulary size is 79,800, the number of layer is 32, the batch size is 1, and using activation recomputation, the maximum model size that can be hosted on one node with EDiT is approximately 40B. Assuming using multiple connected clusters, we can further expand the scope of the model shard group by treating a single cluster as one model shard group, thereby extending the training scale across multiple clusters. **Q6: line 168: Do you observe a different behavior of warm up phase?** A6: As you mentioned, during our initial exploratory experiments, we found that the warm-up phase did not contribute to the final convergence results. However, the warm-up phase can enhance early training stability, particularly when the model is randomly initialized, effectively preventing training divergence. Therefore, we retained the warm-up phase, as described in Section 3.1, solely to improve early training stability. **Q7: Tab.1 & Fig.4: Why not experiment with A-EDiT on the in-house dataset?** A7: We apologize for not completing this experiment before the manuscript submission deadline. We have now added the results of A-EDiT on the in-house dataset in Table 1 and Figure 4. **Q8: Fig 5.a: Can you explain the reason why A-EDiT and EDiT exhibit the same throughput in the random-straggler scenario?** A8: Thank you for your question. In the random-straggler scenario, we randomly select a worker to slow down at each step. Given the uniformity, over a synchronization interval of 128 steps, all workers reach the synchronization point at roughly the same time. This results in minimal idle waiting periods for any worker, leading to similar throughput for both EDiT and A-EDiT. **Q9: In ablation Fig 7.a: Is it really necessary to keep Gradient Clip along with Weighted Average and Anomaly Detection?** A9: Thank you very much for pointing out this problem. We have revised our ablation study to include results where gradient clip and weighted average were individually removed. As shown in Section 4.5, the removal of gradient clip leads to noticeable spikes, and the validation perplexity (PPL) is consistently higher than that of EDiT. This demonstrates that gradient clip plays an essential role in improving both stability and performance.

Gentle Reminder of the Discussion Deadline

Dear Reviewer zyrD, Thank you once again for your time and review of our manuscript. We understand that you have a busy schedule, and we kindly wish to remind you that the discussion deadline is approaching. If there are any remaining concerns or suggestions on how we can further improve our manuscript, we would greatly appreciate your feedback. Additionally, if you feel that we have successfully addressed the points raised in your initial review, we would be grateful if you could consider revisiting the score assigned to the paper. Sincerely, The Authors of Paper 1502

Response to Reviewer zyrD

Dear Reviewer zyrD, Thank you very much for your response. During parameter synchronization in EDiT, the prefetch strategy allows for overlapping **two communication operations** with computation: one is the parameter synchronization (All-Reduce) among nodes (referred to as "model sync group" in the paper), and the other is the parameter aggregation (All-Gather) within nodes (referred to as "model shard group" in the paper). As you mentioned, the intra-node communication is often several orders of magnitude faster than the inter-node communication, thus the overhead from parameter aggregation within nodes is negligible. However, the overhead from parameter synchronization between nodes cannot be ignored. As shown in Section A.3.2, synchronizing a 1B model on two 8-GPU A100 nodes introduces an extra overhead of 160ms with the non-overlapping method (Post Local SGD), whereas EDiT introduces only an extra overhead of 19ms, which is due to the communication time of the first module that cannot be overlapped. From the experimental results in CO2, we can infer that this gap will widen further with an increase in the number of nodes, deterioration in inter-node communication, and larger model scales. The experimental results in Section 4.3 also demonstrate that EDiT can achieve high throughput even with small synchronization intervals, while the performance of Post Local SGD and DiLoCo is relatively poor. In fact, we first tried the straightforward "ZeRO+local SGD" method, but we found that the additional non-overlapping communication it introduced could affect the overall speed in practice. Therefore, we leveraged the advantages of hierarchical network topology and sharded parameters and proposed the layer-wise synchronization method. In summary, we believe that using the prefetch strategy to overlap forward computation with parameter synchronization communication is necessary.

Response to Reviewer 8feR

Dear Reviewer, We appreciate your thorough evaluation of our paper. Below, we address your concerns and provide our responses. All revisions will be included in the forthcoming rebuttal version of the paper. **Q1: How should we think about and choose $\alpha$, $\phi$, and $\beta$ on a new training task?** A1: Thank you for your question. In this paper, $\alpha$ is the decay of the exponential moving average in anomaly elimination, and $\phi$ is the threshold of pseudo gradient clip. $\beta=min(\frac{\phi}{G+\epsilon},1)$ is determined by $\phi$ . A larger $\alpha$ means the exponential moving average reflects the current trends more. A larger $\phi$ implies a looser constraint on pseudo gradient clip. These two hyperparameters only take effect when abnormal pseudo gradients occur and generally do not significantly impact the normal training process or convergence trends. Therefore, their settings are not stringent. 1. Through extensive experiments, we found that setting $\alpha = 0.02$ achieves good results, and this setting can be applied to new training tasks. 2. Similar to gradient clipping in synchronous training, the setting of $\phi$ is generally related to the specific training task. When the training task is unstable, a smaller $\phi$ can be set to stabilize the training. On the contrast, a larger $\phi$ can be set to accelerate convergence. In our experiments, we set $\phi = 10$. For new training tasks, we recommend testing $\phi=1$ and $\phi=10$, respectively. **Q2: Will ignoring examples with large gradient norms create bias in training?** A2: Thank you for the thoughtful question. Our conclusion regarding this issue is that it does not create training bias to ignore samples with large gradient norms. 1. We have analyzed the training process. In most cases, all workers are normal and have similar pseudo gradient norms, so anomaly elimination and weighted averaging do not take effect. We empirically note that only those anomalies that would otherwise cause local loss spikes in workers will trigger pseudo gradient penalty. To prevent misjudgment, we have set a warm-up process and an exponential moving average update strategy in anomaly elimination. As shown in Figure 7c, only about 0.8% tokens were flagged as outliers and excluded in the early stage of training. From Figure 4c, we can observe that as training progresses, the number of outliers decreases, resulting in a total outlier data ratio much lower than 0.8%, which has minimal impact on the overall data distribution for training. Additionally, the results in Figure 4 and Table 1 further confirm that ignoring samples with large gradient norms does not cause training bias. Notably, it has also been mentioned in some LLM technical reports that skipping data points can be a useful technique to suppress loss spikes, such as [PaLM](https://arxiv.org/abs/2204.02311). 2. In the paper [Self-Influence Guided Data Reweighting for Language Model Pre-training](https://arxiv.org/abs/2311.00913), the authors also found that using gradient norms for weighting in the mid-to-late stages of training can improve model performance and robustness. In fact, it can be seen as an approximate RMSProp optimization method to utilize pseudo gradient norms as a constraint metric of pseudo gradients, where the second moment is used to weight the gradients. This allows all workers to contribute equally to the update direction, thereby increasing the likelihood that the model updates along the correct direction. **Q3: Provide convergence analysis for the proposed EDiT method.** A3: Thank you very much for your constructive suggestion. We have added the convergence analysis for our proposed EDiT method. For details, please refer to "Section 5: Theoretical Analysis" in the main text and "Section A.4: Proof of Theorem 1" in the appendix.

Gentle Reminder of the Discussion Deadline

Dear Reviewer 8feR, Thank you once again for your time and review of our manuscript. We understand that you have a busy schedule, and we kindly wish to remind you that the discussion deadline is approaching. If there are any remaining concerns or suggestions on how we can further improve our manuscript, we would greatly appreciate your feedback. Additionally, if you feel that we have successfully addressed the points raised in your initial review, we would be grateful if you could consider revisiting the score assigned to the paper. Sincerely, The Authors of Paper 1502

Final Gentle Reminder - Discussion Period Ending Soon

Dear Reviewer 8feR, We wanted to reach out one final time regarding our manuscript, as the discussion period is drawing to a close very soon. We greatly value your expertise and initial feedback, and we hope our rebuttal has adequately addressed your concerns. If you have any remaining questions or if there are points that require further clarification, we would be most grateful to hear from you before the discussion period ends. Alternatively, if you feel satisfied with our responses, we would appreciate your consideration in updating the manuscript's evaluation. Best regards, The Authors of Paper 1502