BAdam: A Memory Efficient Full Parameter Optimization Method for Large Language Models

Thank you very much for the constructive and helpful comments. We address the reviewer's concerns in a point-by-point manner below. All additional experiment results (figures and tables) are put in the one page supplementary PDF of the global rebuttal.

A. Convergence result under stochastic setting. Our approach of obtaining complexity under the deterministic setting consists of two main steps: (1) establishing the descent property for the inner solver within one block and (2) aggregating the descent inequalities across all blocks for one block-epoch. Step (2) is a standard technique in the complexity analysis of BCD-type methods under reasonable conditions. To extend to stochastic setting, the main difficulty lies in extending step (1). The complete proof is beyond the scope of this rebuttal stage; we provide a proof sketch here. We adopt the same assumptions on the smoothness of the objective function and stochastic gradient errors as in [1].

Notation modification. We distinguish the true gradient $g_{i}^{t,k}$ and the estimated stochastic gradient $\tilde{g}\_i^{t,k}$. The difference term $e_{i}^{t}$ between the true gradient $g_{i}^{t,k}$ and the scaled momentum $\hat{m}_{i}^{t,k}$ now includes stochastic gradient errors:

where the momentum is updated with stochastic gradients $m\_{i}^{t,k} = \beta\_{1} m\_{i} ^{t,k-1} + (1- \beta\_{1} ) \tilde{g}\_{i} ^{t,k}$.

Assuming stochastic gradients with a uniform almost sure bound $\sigma$, similar to [1], $||e_{i} ^{t,k}||$ can be bounded by the norms of the true gradients with probability at least $1- \delta$: $\sum_{k=1}^{K} ||e_{i} ^{t,k} ||^2 \le \Theta (\lambda) \sum_{k=1}^{K} ||g_i^{t,k} || ^2 + \mathcal{O}\left(\sigma ^2 \log (\delta ^{-1} ) ( 1/ \beta_1 + \beta_1 K) \right).$

Approximate descent inequality (Inner solver). Substituting the above probabilistic bound of error terms into the following approximate descent inequality for each step of the inner solver:

$\mathcal{L}( \theta _{i}^ {t,k}) - \mathcal{L}(\theta _{i} ^{t,k-1} ) \le - \Theta (\alpha ) ||\nabla _{i} \mathcal{L} (\theta _{i} ^{t,k-1} ) || ^2 + \Theta (\alpha) || e\_{i} ^{t,k-1} || ^2.$

Then combining a telescoping argument with our technique of bounding gradient differences in the inner loop $||g_{i} ^{t,k} - g_{i} ^{t,1} ||, \forall k \le K$ by $||g_{i} ^{t,1} ||$:

$\sum_{k=1}^{K} || g _{i} ^{t,k} - g _{i} ^{t,1} || \le \mathcal{O} \left(\alpha K^2 \right) || g _{i} ^{t,1} || + \mathcal{O} \left(\alpha K^{3 /2} \sigma \sqrt{(1 / \beta _{1} + \beta _{1} K) \log (\delta ^{-1} ) } \right).$

We obtain the following approximate descent inequality for each block under the stochastic setting:

$\mathcal{L} (\theta\_{i} ^{t} ) - \mathcal{L}(\theta\_{i-1}^{t} ) \le - \Theta ( \alpha K) || \nabla_i \mathcal{L}(\theta_i^{t}) ||^2 +\mathcal{O}\left(\alpha ^3 K ^3 \sigma ^2 \log (\delta^{-1} ) ( 1/ \beta_1 + \beta_1 K) \right).$

Approximate descent inequality (BCD epoch). By aggregating the above approximate descent inequality across different blocks, similar to our arguments in the submitted manuscript, we derive the approximate descent inequality for one block-epoch:

$\mathcal{L}(\theta ^{t+1} ) - \mathcal{L}(\theta ^{t} ) \le - \Theta (\alpha K) ||\nabla \mathcal{L}(\theta ^{t} )|| ^2 + \mathcal{O}\left(\alpha ^3 K ^3 D\sigma ^2 \log (\delta ^{-1} ) ( 1/ \beta_1 + \beta_1 K) \right).$

Through standard manipulations, we should obtain a complexity result for finding an $\varepsilon$-stationary point.

[1] Li, H., Rakhlin, A., & Jadbabaie, A. (2023). Convergence of Adam Under Relaxed Assumptions. NeurIPS.

B. Initial continual pretraining experiment using BAdam. Pretraining from scratch is out of scope of this work; we have conducted an initial continual pretraining (CPT) experiment for Llama 3-8B-Instruct on the StarCoder dataset using BAdam. We evaluate its performance using the online CPT loss; see Figure (f). We are only able to finish about 0.1 epoch (1.5B tokens) training due to time constraints. The online CPT loss shows that the model is effectively learning new domain specific (code-related) knowledge. In the following, we list several possible reasons why BAdam may become a candidate optimizer for LLM continual pretraining.

BAdam exhibits as high rank update as that of Adam. In the pretraining stage, model acquires massive knowledge over different domains by training on billions of tokens. Intuitively, new factual information (not appeared in the pretraining corpus) might not be encoded by low rank update. As shown in Figure (e), BAdam achieves almost the same high rank update as Adam through all modules of different transformer layers, partially ensuring BAdam’s learnability.

Possible ability of avoiding general knowledge forgetting. Perhaps one of the most challenging tasks in CPT is to avoid forgetting of general knowledge. We suspect that BAdam might be good at avoding forgetting compared to Adam. This conjecture stems form the fact that the BAdam uses each data batch to update only one block of parameters, thereby better preserving the general knowledge of the model during CPT. This is partly confirmed by that Llama 3-8B-Instruct has a 67.7 MMLU score (0-shot), while the MMLU score of our CPTed checkpoint after training on 1.5B tokens only decreases to 66.7 (without using any model merging techniques).

We hope that our response is satisfactory to the reviewer and that all concerns have been addressed appropriately. We will properly add the above continual pretraining experiments (with additional comparisons with LoRA and Adam) and try to add convergence result under stochastic setting (if our sketch does not contain vital error) to our next manuscript version. If the reviewer has any additional concerns, please let us know during the reviewer-author discussion period.

Thanks for the author's detailed response. It addressed most of my concerns. However, I share the same issue as Reviewer XVND regarding the GSM8K and MATH scores. I hope this can be addressed properly.

Thank you for following up and the detailed response. I appreciate the effort authors put into addressing this issue. Your explanation has increased my confidence in my comment, and I have raised my confidence level from 3 to 4 and score from 6 to 8. My expertise lies more in pretraining language models, and our team also struggles with dealing with the memory issues brought about by the Adam optimizer. Being able to optimize Adam is a good design. From a practical perspective, I believe this paper can bring relatively significant benefits to the field. I hope the authors can open-source their code in a high-quality manner to facilitate language model community use.

The authors would like to express their deepest gratitude to the reviewer for acknowledging our work and responses. We also greatly appreciate the reviewer for increasing the confidence level from 3 to 4 and the score from 6 to 8.

For pretraining, we will complete the continual pretraining experiment to show the efficiency of BAdam in this pretraining-related setting. Regarding open-sourcing, we promise that our implementation code that can reproduce our paper's results will be made publicly available and include the following features:

1. Easy to use. The integration of the BAdam and BCD framework into user's own codebase will be straightforward, necessitating minimal changes to the existing code.

2. Distributed training support. Our code will support both data-parallel and model-parallel training, based on DeepSpeed ZeRO-3, allowing for the efficient finetuning / training of truly large-scale models (e.g., 70B). We will also make the implementation of distributed training straightforward.

3. Memory efficiency. We will ensure that actual memory usage is consistent with the values reported in our paper. For instance, one will be able to train a Llama 3-8B model using a single RTX3090 and a Llama 3-70B model with just $3\times$ A100-80GB GPUs (based on our distributed training implementation).

Once again, we deeply thank the reviewer for your kind words and support of our work.

Thank you very much for the constructive and helpful comments. We address the reviewer's concerns in a point-by-point manner below. All additional experiment results (figures and tables) are put in the one page supplementary PDF of the global rebuttal.

A. More quantitative results on MMLU and math benchmarks. Following the reviewer's suggestion, we have tested the MMLU scores of the instruction-tuned models; see Table 3. We can observe that BAdam performs as good as Adam and outperforms LoRA.

In Table 2, we have also conducted math-instruction tuning for Llama 3.1-8B on the MathInstruct dataset using BAdam and LoRA, and tested several math benchmarks including GSM8K, MATH, NumGLUE, SVAMP, MMLU-Math, and SAT-Math. The results demonstrate that BAdam clearly outperforms LoRA. We have not conducted the same experiments for Adam due to time constraints, but we will add such experiments in the next version.

We believe that these additional quantitative results, together with the MT-bench results, illustrate the capability of BAdam for LLM finetuning.

B. More block partition schemes. As requested by the reviewer, we have tested the ratio-based block partition scheme, where one block is formed by selecting a small part of parameters from each matrix module. We can observe that ratio partition performs nearly as well as the layer partition in terms of optimization ability as shown in Figure (c). In addition to this ratio partition, we can also treat each attention / mlp module as one block, and our test shows that this scheme also has similar optimization ability to the layer partition. We will add ablation study on these additional block partition schemes in the next version.

C. Fundamental differences from LIFT. Let us first mention that LIFT is properly cited in our Section 4. Though LIFT utilizes a layer-wise update of Adam, it differs from BAdam fundamentally in the following aspects.

Different principles in algorithm foundation. LIFT has two types of iteration policies as stated in their Appendix D, including cyclic and grouped iteration policies. The cyclic policy updates a layer for one iteration and then offloads Adam's optimizer states to the CPU (see the last paragraph of their Section 3.2). This is not a valid BCD with Adam scheme because the gradient of the same layer in the next round of update does not match the offloaded Adam's optimizer states (since other layers' parameters have been updated). On the other hand, their grouped policy updates one layer for the total number of scheduled iterations and does not revisit it, and hence there is no BCD loop at all. This makes their grouped policy a special type of block-restricted implementation of Adam rather than a general BCD scheme.

By contrast, BAdam is built on the BCD optimization framework, resulting in a fundamentally different algorithm. BAdam includes an outer BCD loop and properly accumulates the first and second momentum for updating a block ($K$ in our algorithm description). Importantly, when BAdam revisits a block, its first and second momentum start at 0 (no offloads), which is crucial for ensuring both theoretical and empirical convergence.

Theoretical convergence guarantee. Thanks to BAdam's foundation in the BCD framework, we can provide a convergence guarantee, adding reliability and predictability to its performance.

Practical code implementation. We have made significant efforts to develop a compatible and robust code implementation for BAdam, with special attention to memory management. For example, our submitted code allows efficient finetuning of Llama 3-8B using a single RTX3090-24GB GPU, whereas LIFT requires several A100 GPUs (according to their experiment description since they have not open-sourced their implementation). Indeed, implementing a direct block-restricted update of Adam could be trivial with sufficient resources, as one can simply set a certain layer to require gradients in the original Adam optimizer without needing complicated memory management and optimizer coding. Our practical implementation ensures broader applicability and benefits for practitioners.

D. Verbose experiment section. We fully agree with the reviewer's concern. We have conducted additional experiments, including more quantitative results (MMLU and math benchmarks), ablation experiments comparing Adam, SGD, BAdam, and BSGD (BCD with SGD), comparison with Galore, and continual pretraining; see the one page supplementary PDF of the global rebuttal. We will incorporate these experiments into our next manuscript version. Specifically, our plan is to move a large part of Section 3.1.1, the convergence verification in Section 3.1.2, and the entire Section 3.2 to the Appendix. Instead, we will add more downstream evaluations (such as MMLU and math tests), create an ablation study section, and a continual pretraining section.

We hope that our response is satisfactory to the reviewer and that all concerns have been addressed appropriately. If the reviewer has any additional concerns, please let us know during the reviewer-author discussion period.

Thank you very much for the constructive and helpful comments. We address the reviewer's concerns in a point-by-point manner below. All additional experiment results (figures and tables) are put in the one page supplementary PDF of the global rebuttal.

A. Ablation study & effectiveness of BCD for LLM finetuning. We have conducted ablation experiments on Adam, SGD (LOMO), BAdam, and BSGD (BCD with SGD) for finetuning Llama 3-8B. All optimization methods use grid-searched learning rates based on training and validation losses.

Convergence ability. In Figure (b), it can be observed that BCD variants converge slightly slower but soon exhibit similar convergence behavior in terms of running time compared to their full counterparts. It is worth mentioning that, unlike the full counterparts, BCD variants only update a block of parameters per data batch, demonstrating the strong optimization ability of BCD for LLM finetuning.

Downstream performance. In Table 1, we also test the MT-bench scores of the finetuned models. It is quite interesting to see that BSGD significantly outperforms SGD (almost as good as BAdam), even though they have almost the same optimization convergence behavior. The superiority of BCD variants over their full counterparts possibly arises from the fact that BCD uses each data batch to update only one block of parameters, thereby better preserving the general knowledge of the pretrained model during finetuning. These strong downstream performances of BCD further illustrate its suitability for LLM finetuning.

Learnability interpreted by high rankness. Unlike LoRA and Galore, BCD's memory efficiency is achieved without restricting its updates to a low rank space. In Figure (d), we display the cumulative explained variance of BAdam's update for the up-proj matrix, which is defined as sum of the $k$ largest squared singular value over the sum of all squared singular value. This result shows that BAdam's update has a heavy tailed singular values distribution and is far away from a low rank update. In Figure (e), we show that BAdam achieves almost the same high rank update as Adam through all modules of different transformer layers, partially ensuring BCD's learnability.

We believe that these experiment results demonstrate the effectiveness of BCD for LLM finetuning.

B. Handle the extremely memory-limited setting with BAdam. In the extremely memory-limited setting, we have the following two choices to apply BAdam.

Smaller block partition. Our BCD optimization framework (as well as our code implementation) provides a flexible block partition. Hence, we can choose a smaller block partition, further reducing the memory consumption as we only need to store the gradient and optimizer states for a smaller block of parameters. We have tested such a smaller block setting of BAdam (each layer is further partitioned to attention and mlp blocks), and it achieves a MT-bench score of 6.50, which is just slightly worse than the score of 6.67 reported in the manuscript.

Cheap CPU update. We can also store the optimizer states in CPU memory, and performs the cheap block update purely on CPU. In this case, only the block gradient needs to be communicated between CPU and GPU during update. We have tested this approach under the Llama 3-8B experiment setting and found that the operation only induces a 0.8-second overhead per update, resulting in roughly 16 minutes of additional time cost over 3 epochs of training. We will upload this part of new code implementation when uploading is allowed.

We remark that MeZO and LOMO also maintain the gradient of a certain block in memory when updating this block. Thus, our BCD scheme can be applied using one of the above schemes whenever MeZO and LOMO are applicable.

C. BAdam needs to store the full checkpoint. Storing the full parameter checkpoint appears to be inevitable when using full parameter optimization methods such as Adam and SGD. While LoRA has the advantage of storing different type of checkpoints, full parameter finetuning might achieve higher performance, as we have shown.

D. LOMO update precision issue. LOMO's update precision depends on the precision of the model weights. Consequently, updating the model using Float32 will double the memory consumption to $4M$. We follow the reported update precision (Float16) used in the LOMO paper.

E. LOMO's performance under grid searched step size. We have conducted a grid search for SGD (LOMO)'s learning rate among (5e-2, 1e-2, 1e-3, 1e-4, 1e-5, 1e-6) based on loss, and identified 1e-2 as the optimal one for MT-bench evaluation. In Table 1, it can be observed that SGD (LOMO) with learning rate 1e-2 achieves an MT-bench score of 5.82, which aligns with the score of 5.83 obtained by using learning rate 1e-6 in our manuscript.

F. Batch size is not consistent across different methods. Since LOMO performs update on the fly and does not store the full gradient, it does not support gradient accumulation. We choose the largest batch size under the memory constraint of RTX3090-24GB for LOMO. We also ensure each algorithm uses the same amount of data. Hence, we count twice updates of LOMO as one update in our manuscript, ensuring a fair comparison to LoRA and BAdam.

We hope that our response is satisfactory to the reviewer and that all concerns have been addressed appropriately. We will properly add the above experiments to our next manuscript version. If the reviewer has any additional concerns, please let us know during the reviewer-author discussion period.

Thank you very much for the constructive and helpful comments. We address the reviewer's concerns in a point-by-point manner below. All additional experiment results (figures and tables) are put in the one page supplementary PDF of the global rebuttal.

A. Ablation study on BCD, Adam, and SGD. As requested by the reviewer, we have conducted experiments on Adam, SGD (LOMO), BAdam, and BSGD (BCD with SGD as inner solver) for finetuning Llama 3-8B. All optimization methods use grid-searched learning rates based on training and validation losses.

Convergence ability. In Figure (b), it can be observed that BCD variants converge slightly slower but soon exhibit similar convergence behavior in terms of running time compared to their full counterparts. It is worth mentioning that, unlike the full counterparts, BCD variants only update one block of parameters per data batch, demonstrating the strong optimization ability of BCD for LLM finetuning.

Downstream performance. In Table 1, we also test the MT-bench scores of the finetuned models. It is quite interesting to see that BSGD significantly outperforms SGD (almost as good as BAdam), even though they have almost the same optimization convergence behavior. The superiority of BCD variants over their full counterparts possibly stems from the fact that BCD uses each data batch to update only one block of parameters, thereby better preserving the general knowledge of the pretrained model during finetuning. These strong downstream performances of BCD further illustrate its suitability for LLM finetuning.

These ablation results partly confirm the suitability of BCD for LLM finetuning and reveal that choosing either Adam or SGD as the inner solver has similar performance in terms of our downstream performance evaluation (i.e., MT-bench score for model trained on Alpaca-GPT4 dataset). We will add BSGD to our BCD framework and leave the comprehensive study of different inner solvers as future work.

B. The suitability of BCD in LLM finetuning. Let us elaborate on the suitability of BCD for LLM finetuning from the following three perspectives.

Performance. Based on the results of the above ablation study, we believe that we have demonstrated the strong performance achieved by BCD compared to its full counterpart. This directly indicates that BCD is suitable for LLM finetuning in terms of both optimization ability and downstream performance.

Low memory consumption. BCD not only achieves strong performance but is also memory efficient; it only stores the gradient and optimizer states of the active block, which is substantially lower than Adam's $18M$ cost. In particular, BSGD requires little additional memory beyond just storing the $2M$ LLM, while still exhibiting strong downstream performance. Given that memory is usually the bottleneck for finetuning LLM, the memory efficiency of BCD further demonstrates its suitability for this task.

Learnability interpreted by high rankness. Unlike LoRA and Galore, BCD's memory efficiency is achieved without restricting its updates to a low rank space. In Figure (d), we display the cumulative explained variance of BAdam's update for the up-proj matrix, which is defined as sum of the $k$ largest squared singular values over the sum of all squared singular values. This result shows that BAdam's update has a heavy tailed singular values distribution and is far away from a low rank update. In Figure (e), we show that BAdam achieves almost the same high rank update as Adam through all modules of different transformer layers, partially ensuring BCD's learnability.

C. BCD + LoRA. We first note that the high rank update of BAdam shown in Figures (d) and (e) partly illustrates the scientific difference between BCD and LoRA. Following the reviewer's suggestion, we also combined BCD with LoRA (B-LoRA) for LLM finetuning. Due to the advantages of BCD, B-LoRA can have a higher rank configuration under the same amount of memory budget and save computation time, compared to LoRA. B-LoRA has a slightly worse MT-bench score than LoRA (6.12 vs. 6.28) as shown in Table 1, while achieves a better 5-shot MMLU score than LoRA as shown in Table 3.

D. Comparison with Galore. In Figures (a)-(b) and Table 1, we report the results of Galore with rank 1024 and 8-bit Adam optimizer for finetuning Llama 3-8B. We set other hyperparameters according to their paper's suggestions. It can be seen that BAdam and BSGD outperform Galore in terms of MT-bench downstream performance, demonstrating the ability and suitability of BCD for LLM finetuning. In terms of convergence, BAdam also exhibits better convergence ability than Galore.

We have not conducted experiments on combining BCD and Galore due to time constraints, but we will add such experiments in our next manuscript version, as suggested by the reviewer.

E. Significance of advantage over LOMO. We note that the MT-bench score difference between GPT-3.5-turbo and GPT-4 on the LMSYS leaderboard is 1.05, which is close to the gap of 0.96 presented in our manuscript between the Llama 3-8B models tuned by LOMO and BAdam (5.69 versus 6.65). Therefore, the advantage appears to be significant.

We hope that our response is satisfactory to the reviewer and that all concerns have been addressed appropriately. We will properly add the above experiments to our next manuscript version. If the reviewer has any additional concerns, please let us know during the reviewer-author discussion period.

Dear reviewer FWuu,

We hope that our response justifies the design of BAdam and addresses most of your concerns. Since the deadline of the discussion phase is approaching, we would highly appreciate to receive feedback from you. If you have any further questions or remarks, please let us know so that we will have enough time to address your concerns. We will be more than happy to provide additional clarifications and details.

Best,

Authors.