Question & Answer

Thank you for your patience in reading the above part. Then we clarify your concerns one by one.

Q1: The novelty of SEPARATE and the theory

We think the discussions in the Important Discussion can answer this part of your concerns. In total, the simplicity of the method is necessary for the acceleration of wall-clock training time. We design such an exquisite and simple method to realize it in training, and the theoretical analysis is also beyond.

Q2: The moving average error feedback section

Thank you for pointing out the unclear part of our writing in this section. We have refined this part in the update version. You can refer to refined Section 4.2 for details. To provide a brief overview, we emphasize the utilization of random vectors for projection may engender large deviation, which implies the potential for abrupt fluctuations in the error term $e_n^k$. This is particularly pronounced under multiple consecutive iterations with a series of random vectors that possess significantly divergent directional properties. Such a phenomenon can precipitate the entire training process into converging towards an alternative suboptimal region, as graphically illustrated in Figure 3(b) in Section 6.3.

Furthermore, we introduce the moving average error feedback, which, in theory, can regulate the accumulated error to a magnitude on the order of the compression error of the current step. This technique exhibits commendable performance in practical applications. Notably, for training from scratch, the absence of error feedback fails to converge, while the absence of moving average error feedback leads to convergence towards an alternative suboptimal region. These findings show the significance of our method in ensuring the stability and efficacy of the training process.

Q3: Experiment and reproducibility concerns

Thank you for proposing your concerns about the experiments. In fact, we have proposed the detailed experimental setting of our experiments in Appendix E, including the details of methods being compared in Table 4 in Appendix E.3, and the hyperparameter setting of all the methods in Appendix E.1. Specifically, for training from scratch, our global batch size was 8 $\times$ 512. We set the learning rate at 6.0e-4 with the cosine decay down to the minimum 6.0e-5 after 2000 iterations of warm-up. We also used the gradient accumulation and set the gradient accumulation step at 4. We used the global gradient norm clipping of 1 and set the Adam with $\beta_1 =0.9$ and $\beta_2 = 0.95$. For finetuning tasks, we set the data parallelism at 4 and model parallelism at 2 without pipeline parallelism. We set the global batch size at 32, the learning rate of 2e-5, gradient clipping of 2 and gradient accumulation step at 8. To be fair, these hyperparameters are set the same and we follow the default setting in pre-training and fine-tuning corresponding models.

The improvement of error feedback is remarkable in training from scratch as we talk in the response of Q2. Without error feedback, the pre-training is difficult to work. For fine-tuning tasks, the impact of error feedback on the results is small, but still positive. In essence, we did not tune the hyperparameters specifically for SEPARATE; instead, we utilized the same configurations to maintain consistency and comparability with established practices. We have discussed these in detail in updated Appendix A.

Q4: What the "top-heavy" eigenvalues have to do with gradient compression

We think the discussion in the second part of Important Discussion can answer this question. "Top-heavy" indicates the trace of Hessian is bounded and thus the variance of our estimate is limited and we obtain the main theorem by it.

Resopnse to Reviewer nVYQ

We sincerely thank the reviewer nVYQ for the valuable and constructive comments. We have updated the submitted paper and try to clarify your concerns. We start with an important discussion to address your concerns about our method and theory in total. Then we clarify the weaknesses and questions one by one.

Important Discussion

We begin with expressing our gratitude for your insightful observation regarding the application of related techniques to diverse tasks, exemplified by GaLore and Flora. However, what we need to clarify is that communication-efficient training has specific challenges that make our "simple but efficient" method necessary.

First, it has indeed come to our attention during our investigation that updating the optimizer state within a low-dimensional subspace is perceived as a strategy to mitigate memory overhead. However, it is imperative to clarify that these methods cannot be directly applied to the reduction of communication overhead. The reduction of memory and the mitigation of communication overhead need to be recognized as distinct challenges for efficient training. GaLore and our approach each address different sides of efficient training. While there exists the potential for a synergistic combination of these two, the additional computational burden is deemed more intolerable for communication reduction. Specifically, for memory reduction, the methods introduce some frequent heavy computation (like SVD) to maintain the accuracy of the estimate, but for communication reduction, to reduce the wall-clock time, lightweight computation is necessary. The primary goal of communication cost reduction is to reduce the overall wall-clock training time. To achieve this, our "simple but efficient" method is currently the only solution. Under frequent communication rounds, the computationally intensive nature of SVD and analogous operations becomes prohibitive. A more detailed discussion of this topic has been presented in updated Appendix A.

Second, in the subsequent discourse, we delve into a more detailed sketch of our theoretical analysis. We start with the variance analysis shown in Lemma 4.2. This lemma serves as a foundation in our theoretical analysis framework, elucidating that the variance inherent in our estimate can be effectively bounded by the trace of the Hessian, as we discuss in Remark 4.3. The observation of the "top-heavy" of the Hessian is significant. It is a pivotal property that ensures a bounded trace, as rigorously defined in Assumption 5.4. This boundedness of the trace is a critical factor in substantiating the bounded variance, which in turn underpins the main Theorem 5.5. Within this theorem, we achieve a convergence rate of $\Omega(d^{1/2}\epsilon^{-4})$ in the standard setting. The result transcends a mere adaptation of existing methods. It represents a significant improvement. This improvement is particularly noteworthy when compared with the conventional convergence rate of gradient compression techniques $\mathcal{O}(d\epsilon^{-4})$, as we discuss in Remark 5.6. It signifies a substantial advancement in communication-efficient optimization algorithms. This improvement is a testament to the efficacy of our approach in mitigating the computational complexities associated with high-dimensional data, thereby offering a more streamlined path to convergence. This work is a cornerstone of our contribution, offering a deeper insight into the interplay between the Hessian's structure, the trace's boundedness, and the consequent impact on the convergence of optimizers.

Thank you for you positive feedback and for adjusting the rating. We genuinely appreciate your insightful and constructive comments. We will refine our work according to your comments and investigate remaining problems in future work. Your engagement has greatly enhanced the quality of our work.

Response to Reviewer acpN

We sincerely thank the reviewer acpN's approval of our work and the valuable and constructive comments proposed. We have updated the submitted paper and try to clarify your concerns. We try to answer your questions one by one, and hope our clarification could help solve your concerns and improve our submission.

Some Related Work

Thank you for sharing the related work with us, which will undoubtedly assist our readers in finding relevant techniques. We have cited Galore in the new version of our submission. Moreover, it is worth noticing that these methods for memory-efficient training have significant differences from ours. The key challenge we focused on is a simple but effective strategy aimed at reducing communication overhead. Even minor increases in computational cost are undesirable during the frequent communication process. We have further discussed the differences between Galore and ours in updated Appendix A.

Report of Data Transfer

Thank you for your suggestions about reporting the data transfer to demonstrate reduction in data exchanged. we report the model flops utilization (MFU) and throughout (tokens per second) on each GPU when pre-training GPT-2 345M from scratch to prove the data exchanged reduces as below. The results show that we have reduced the data exchanged.

We sincerely appreciate your constructive suggestions and believe the discussion and additional experiments significantly improve the quality of our submission. We hope this provides sufficient reasons to raise the score.

Response to Reviewer Notb

We sincerely thank the reviewer Notb's for the valuable and constructive comments. We have conducted some experiments to add some adaptability to our method. We try to adaptively set the compression ratio and $\beta$ in moving average error feedback. We propose the preliminary experimental results and hope these could help solve your concerns and improve our submission.

Robustness of Our Method

We are grateful for your concerns regarding the robustness of our method and would like to address them by clarifying three key aspects. First, our variance analysis and the main theorem demonstrate that, provided the selection of the compression ratio aligns with the conditions in our theoretical analysis (Theorem 5.5 and 5.7), our method can exhibit effective performance. This theoretical underpinning ensures that our method is robust within the specified compression ratio. Second, as a result, in practical applications, especially for models with a substantial number of parameters, such as those in the millions or billions, we can defaultly set the compression ratio as 16, 32, or 64. As illustrated in Table 3, the performance differences between compression ratios of 16, 32, or 64 times for gradient information are minimal. Consequently, the choice of compression ratio is more influenced by the user's device constraints rather than the intrinsic characteristics of the model. Third, in our experimental setup, all hyperparameters are derived from the default settings of the respective models as trained within their corresponding frameworks. We do not tune the hyperparameters specifically for SEPARATE. Instead, we utilize the same configurations to maintain consistency and comparability with established practices. We have discussed these in detail in updated Appendix A.

General Version of SEPARATE

Thank you for the interesting suggestion on adaptive extension of our method. We have carefully considered your requirements and extended our algorithm with adaptive compression ratio and $\beta$. We show the algorithm as below. You can refer to the formal version in updated Appendix E.5.

Algorithm: G-SEPARATE: General Version of Simple Low-rank Projection

Input:

• Initialization model parameters with $L$ layers, $N$ nodes, layer-wise communication ratio $\lbrace m_l \rbrace _ {l=1}^{L}$, layer-wise $\lbrace \beta_l\rbrace_{l=1}^{L}$, $\beta_l=0.95$, $\forall l \in [L]$, error reset frequency $T_e$, adaptive update frequency $T_a$, a common Gaussian random number generator, initialize $e^0 \in \mathcal{B}(\mathbf{0},c_1)$

While $k \leq K$ do:

1. In each node $n$ compute stochastic gradient $g_{n,l}^k$ and $h_{n,l}^k = g_{n,l}^k + e_{n,l}^k$;
2. Generate fresh i.i.d. common random Gaussian vectors $\xi_1,\cdots,\xi_{m_l^k}\sim N(0,I_d)$ and compute $[p_{1,n,l},\cdots,p_{m_l^k,n,l}]$ with $p_{i,n,l} = \langle h_{n,l}^k,\xi_i\rangle$ as the low-dimension projection of $h_{n,l}^k$;
3. Do all-reduce and obtain global projected gradient $[\tilde p_{1,l},\cdots,\tilde p_{m_l^k,l}]$;
4. Compute $\tilde h_{n,l}^k = \frac{1}{m_l^k} \sum_{i=1}^{m_l^k} \tilde p_{i,l} \cdot \xi_i$ and use $\tilde h_{n,l}^k$ for model weight update in node $n$;
5. Update error:
   • $e_{n,l}^{k+1} = (1-\beta_l^k)e_{n,l}^k + \beta_l^k (h_{n,l}^k -\tilde h_{n,l}^k)$ if $k$ % $T_e \not= 0$;
   • $e_{n,l}^{k+1} = 0$ if $k$ % $T_e =0$ (error reset);
6. Layer-wisely update compression ratio and $\beta_l$:
   • $m_l^{k+1} = {\rm int}\left(1 + m_l^k \cdot \left(1 + \frac{\langle \tilde h_{n,l}^k, g_{n,l}^k \rangle}{\Vert \tilde h_{n,l}^k \Vert \cdot \Vert g_{n,l}^k \Vert} \right)\right),$
   • $\beta _ l^{k+1} = \max\left\lbrace\min\left\lbrace\beta_l^k \cdot\left(1 + \frac{\langle \tilde h_{n,l}^k, g_{n,l}^k \rangle}{\Vert \tilde h_{n,l}^k \Vert \cdot \Vert g_{n,l}^k \Vert} \right), 0.99 \right\rbrace, 0.90 \right\rbrace$ if $k$ % $T_a = 0$;
   • $m_l^{k+1} = m_l^k$,
   • $\beta_l^{k+1} = \beta_l^k$, if $k$ % $T_a \not= 0$;

End While

Response to Reviewer Kod5

We sincerely thank the reviewer Kod5's approval to our work and the valuable and constructive comments proposed. We have updated the submitted paper and try to clarify your concerns. We try to answer your questions one by one, and hope our clarification could help solve your concerns and improve our submission.

Question & Answer

Q1: Sharing random generator between nodes might introduce a layer of complexity

Thank you for pointing out this matter and for your suggestions. We appreciate your feedback on areas where we could improve our explanations and clarifications. First, in practice, there is no additional communication cost involved, because we can initialize the random variable generator with the same specific seed on each node, rather than generating it in one node and then broadcasting it to others. This ensures efficiency and avoids extra communication overhead. We have discuss this in updated Appendix A. Second, we have conducted an assessment of the introduced randomness with the aim of evaluating the influence of random seeds on our results, as detailed in Section 6.3. A more comprehensive discussion on this topic is provided in the response to Q7.

Q2: About previous estimator and ours

Thank you for sharing the related work with us, which will undoubtedly assist our readers in finding relevant techniques. We have incorporated this related work into the new version of our submission in line 143 and 144. Moreover, we have provided a clarification of our approach compared with other previous estimator for memory-efficient training. The key challenge we focused on is a simple but effective strategy aimed at reducing communication overhead. Even minor increases in computational cost are undesirable during the frequent communication process, as we have detailed shown in Introduction and updated Appendix A. To achieve this, we have developed an improved moving average error feedback technique, which is particularly effective when training from scratch. Our strategy is independent of the optimizers and can be seamlessly integrated into FSDP. In contrast, previous methods were either challenging to implement or were limited to DDP applications. We have discussed these in updated Appendix A.

Q3: Some concerns of Table 1

Thank you for bringing this issue to our attention, which we acknowledge was not previously explained with sufficient clarity. We have taken your feedback and have provided a detailed explanation of this in line 413 and 414 in the updated version of our submission. To clarify, when we refer to "performance," we are indicating the scores achieved on the relevant benchmarks. The practice of taking the average score is a widely accepted method for evaluating LLMs [1,2]. In line with your suggestion, we have also taken into account the average rank and have included this metric in Table 1 for a comprehensive assessment. We show the results below.

Q4: Regarding regularization by projection

We appreciate your observation regarding the limitations in our previous discussion. It is important to note that our initial remarks constituted a preliminary hypothesis and interpretation of the observed phenomenon. We acknowledge that this is not a straightforward issue that can be succinctly elucidated under with study; therefore, we intend to conduct a more in-depth investigation in future research endeavors. Consequently, in the current version of our work, we have elected to delete the aforementioned sentence to avoid any potential misinterpretation or oversimplification of the complex problem at hand.

Q5: Motivation V.S. new discovery

We extend our sincere gratitude for highlighting the ambiguity present in our submission. We have meticulously refined this part within the updated version, with a clear distinction between the motivation and our contributions in the Introduction. We underscore that the observation of the "top-heavy" Hessian is the cornerstone of our research motivation, and we present both theoretical and empirical evidence of this phenomenon in Section 3. Subsequently, we propose our contributions about designing an efficient algorithm for communication-efficient training, which encompasses the common random projection technique, an improved moving average error feedback mechanism, theoretical analyses, and corroborative experimental results.

Q6: Some explanations and support

We appreciate your suggestions regarding the areas in our previous submission that required further elucidation. We have revised the description of error feedback in the updated version. First, we would like to emphasize a critical aspect of single-step iteration within our method. Specifically, the use of random vectors for projection may introduce significant deviation (although this doesn't always happen). This implies the potential for abrupt fluctuations in the error term $e_n^k$, particularly under consecutive iterations with a sequence of random vectors that exhibit markedly different directions. Such fluctuations can precipitate the entire training process into converging towards an alternative suboptimal region, as shown in Figure 3(b) in Section 6.3. In practical applications, we have noted a significant improvement in the training process when moving average error feedback is incorporated, especially when training from scratch. Referred to Section 6.3, the absence of error feedback renders the pre-training phase challenging to execute effectively.

Q7: Concerns of ablation study in Table 2

Thank you for providing this suggestion for the report on the ablation study. What we want to emphasize is that for the baseline, our method and other methods to be compared, we use the greedy sample for the generation process. The results fluctuate less so we do not report as average + std form in the original version. We take your suggestion into consideration and change the sampling method, but it takes time to modify all the evaluation experiments due to equipment constraints. We report the results of the experiments in Table 2 that have been carried out below, and we will update it in the coming days. As stated below, the difference is small enough.

We sincerely appreciate your constructive suggestions and believe the discussion, analysis, explanations and additional experiments significantly improve the quality of our submission. We hope this provides sufficient reasons to raise the score.

[1] Edward Beeching et al. Open LLM Leaderboard (2023-2024).

[2] Abhimanyu Dubey et al. The Llama 3 Herd of Models.