This work presents BAdam, an optimization method that leverages the block coordinate descent (BCD) framework with Adam's update rule. BAdam offers a memory efficient approach to the full parameter finetuning of large language models. We conduct a theoretical convergence analysis for BAdam in the deterministic case. Experimentally, we apply BAdam to finetune the Llama 3-8B and Llama 3-70B models using a single RTX3090-24GB GPU and 4 A100-80GB GPUs, respectively. The results confirm BAdam's efficiency in terms of memory usage, running time, and optimization capability. Furthermore, the downstream performance evaluation based on MT-bench and math benchmarks shows that BAdam outperforms existing memory efficient baselines such as LoRA. It also demonstrates that BAdam can achieve comparable or even superior performance compared to Adam. Finally, the ablation study using SGD's update rule illustrates the suitability of BCD for finetuning LLMs. Our code can be easily integrated into any PyTorch-based codebase and is available at https://github.com/Ledzy/BAdam.
Paper
Similar papers
Peer review
Summary
This paper proposes BAdam, a novel optimization method for memory-efficient full parameter finetuning of large language models. BAdam leverages the block coordinate descent framework with Adam as the inner solver. It partitions the model parameters into blocks and updates one block at a time using Adam steps.
Strengths
1. The paper addresses the highly relevant and important problem of enabling full parameter finetuning of large language models under memory constraints. BAdam provides an original and creative solution by combining block coordinate descent with Adam. 2. The theoretical convergence analysis lends credibility to the proposed method, even if limited to the deterministic case. The proof seems sound and complete. 3. The experiments are thorough and convincing, clearly demonstrating BAdam's advantages in memory efficiency, running time, convergence behavior, and downstream performance compared to strong baselines like LoRA and LOMO across multiple models and datasets.
Weaknesses
1. While sufficient for an initial proposal, the theoretical analysis is limited to the deterministic case. Extending the convergence results to the stochastic setting would further strengthen the paper. 2. The paper focuses on applying BAdam to the finetuning stage of large language models. It would be interesting to explore and discuss whether the proposed method could also be applied during the pretraining phase. While I understand that conducting experiments on pretraining may be prohibitive due to time constraints, providing some conceptual discussion on the feasibility, potential benefits, and challenges of using BAdam for pretraining could broaden the scope and impact of the work. For example, the authors could comment on whether the block-wise update scheme of BAdam would remain effective and efficient when dealing with the larger datasets and longer training horizons typically involved in pretraining. Addressing this aspect, even briefly, would give readers a more comprehensive view of BAdam's potential across different stages of the model development pipeline.
Questions
See my weakness part.
Rating
8
Confidence
4
Soundness
3
Presentation
3
Contribution
3
Limitations
The authors have discussed key limitations such as the theoretical analysis being restricted to deterministic gradients and comparisons with Adam being limited to medium-sized models due to resource constraints. Suggestions are provided to address these limitations in future work.
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.
Authors' Response
Thank you for thoroughly reviewing our response and tracking the feedback of other reviewers. We are pleased to know that most of your concerns have been addressed. We have provided further clarification regarding the score gap of GSM8K and MATH, please see our response to reviewer XVND. We would like to thank the reviewer once again for supporting our work. Should you have any additional questions, please let us know.
Comment on Scores of GSM8K and MATH
Dear reviewer JX1Q, Thank you for your thoughtful participation in discussion. We would like to provide additional clarifications regarding the few-shot benchmark scores for GSM8K and MATH, as previously elaborated in our response to Reviewer XVND. We recognized that these benchmarks are highly sensitive to the chain-of-thought (CoT) prompting technique, where minor variations in the prompts can lead to heavy fluctuations in the score. Consequently, the impact of the optimization algorithms diminishes when CoT samples are employed, as we have empirically verified. Based on this observation, we believe that evaluating the efficacy of BAdam solely based on these specific few-shot scores—without accounting for the CoT's impact—may not provide a convincing assessment of its performance. Additionally, we remark that it would be inappropriate to disregard BAdam's performance on the zero-shot setting, as it provides a measure that excludes the CoT's effects. We hope that our clarification addresses your concern regarding the benchmark scores. Should you have any additional question, please let us know. Best, Authors
Thank you
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.
Thank You
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.
Summary
The paper introduces BAdam, a memory-efficient optimization method for fine-tuning large language models (LLMs) by leveraging block coordinate descent (BCD) with Adam as the inner solver. BAdam aims to reduce memory consumption while maintaining or improving performance. The paper presents theoretical convergence analysis and experimental results showing BAdam's effectiveness compared to existing methods like LoRA and LOMO, particularly in terms of memory efficiency and downstream performance.
Strengths
- **Extensive Theoretical Proof of Convergence**: The paper provides substantial theoretical evidence to support the convergence of the proposed method. - **Detailed Method Analysis**: The analysis of the method is thorough, covering aspects like memory consumption and computation time.
Weaknesses
- **Need for More Quantitative Results**: The evaluation of 7B and 8B models requires more quantitative results, such as mathematical and world knowledge benchmarks (e.g., GSM8K and MMLU). Relying solely on MT-bench, which is scored by GPT-4, is not sufficiently objective. - **Block-wise vs. Layer-wise Updates**: The paper's core discussion revolves around block-wise updates, but the actual implementation uses layer-wise updates. Other formats of block-wise updates should be explored. - **Similarity to Existing Work**: The motivation of this paper is similar to "LIFT: Efficient Layer-wise Fine-tuning for Large Model Models"[1] which also discusses "learning one layer/block at a time" in Section 3.2. This similarity needs to be addressed and differentiated. - **Verbose Section 3.1.2**: The discussion in Section 3.1.2 is overly verbose. Experiments with learning rates and other hyperparameters could be moved to the ablation studies or the appendix, while the main text should focus on the core experimental results. [1] https://openreview.net/forum?id=u0INlprg3U
Questions
N/A
Rating
5
Confidence
4
Soundness
3
Presentation
2
Contribution
2
Limitations
N/A
Summary
The paper introduces memory-efficient optimizer BAdam, which combines the concepts of block coordinate descent (BCD) and Adam's update rule. BAdam demonstrates that, with moderate memory consumption—more than LOMO—it can surpass LoRA and significantly outperform LOMO in fine-tuning Llama 2-7B and Llama 3-8B. Additionally, BAdam shows similar fine-tuning performance to Adam when applied to the medium-sized masked language model RoBERTa-large.
Strengths
The paper's contributions and strengths are as follows: 1. It proposes using the well-known optimization technique BCD for the task of fine-tuning large language models while being memory efficient. 2. Empirical evidence highlights BAdam's potential in fine-tuning language models. It can outperform LoRA and significantly exceed LOMO in instruction-tuning Llama 2-7B and Llama 3-8B. Moreover, BAdam demonstrates comparable fine-tuning performance to Adam when used with the medium-sized masked language model RoBERTa-large. 3. It provides theoretical convergence analysis for the deterministic case.
Weaknesses
The paper's weaknesses are summarised as follows: 1. In extremely memory-limited settings, such as when there is only enough memory for inference, where MeZO or LOMO can apply, BAdam cannot apply due to the additional requirement of storing block-wise optimizer states. 2. The paper does not provide theoretical or practical insights into why BCD is effective for language model fine-tuning. 3. Compared to LoRA, BAdam requires storing full parameter checkpoints instead of a small number of adapters. As the scale of language models continues to grow, the issue of storing full parameters for fine-tuning becomes much more intolerable in practice.
Questions
1. In Table 1, why LOMO update precision is only limited to Float16. Can't it be Float32? 2. In Table 5, using the same learning rate for SGD and AdamW is not convincing. Typically, SGD requires a larger learning rate than AdamW to achieve good performance. Have the authors tried conducting a grid search to determine the optimal learning rate for LOMO? 3. When presenting the results in Table 5, what parameters were fixed? Batch size, memory, etc.? I noticed that the batch sizes are not consistent across methods; for example, LOMO uses a batch size of 8, while LoRA and BAdam use a batch size of 16. Comparing different optimization methods with varying batch sizes is not ideal. 4. I believe more ablation studies could be conducted. For example, why is Adam used for the intermediate steps instead of SGD? How much performance loss would occur if SGD were used instead?
Rating
5
Confidence
3
Soundness
2
Presentation
3
Contribution
3
Limitations
The author has addressed the limitations.
Summary
The authors proposed fine-tuning of LLMs with a block coordinate descent based Adam optimizer. They presented results on convergence analysis, memory and run time profiling, and the quality of resulting fine-tuned models.
Strengths
+ The proposed idea is clearly stated. + The memory usage analysis is comprehensive. + The authors conducted experiments on recent LLMs.
Weaknesses
- The novelty of the proposed technique is BCD optimization for fine-tuning LLMs, using the Adam optimizer. In this sense, the comparison against LOMO is justified, but the empirical advantage is not very significant. The author did not perform ablation study showing the necessity of the two ingredients: BCD and Adam--another important condition to compare with LOMO and BAdam would be BCD with SGD. - The comparison against PEFT such as LoRA is not scientifically justified, on the other hand. Because BCD and parameter-efficient reparameterization are not mutually exclusive, and could be complementary. - In addition to optimizer (e.g. LOMO and this work) and PEFT, there is a third class of memory-efficient LLM fine-tuning techniques: gradient compression (e.g. Galore, arXiv:2403.03507). Like LoRA, this is orthogonal and potentially complementary to BCD optimization as well. There is no comparison here.
Questions
* See above weaknesses on empirical results. * A crucial motivation of the proposed method, is a conjectured unique suitability of BCD in LLM fine-tuning: as the authors put it, "the finetuning process boils down to an optimization problem that needs to handle a huge number of trainable model parameters, while the number of training data points are relatively small. This setting matches exactly the advantage of the BCD scheme". Unfortunately, however, this remained a conjecture. There is a lack of adequate theoretic or empirical support to this interesting and important open question.
Rating
3
Confidence
4
Soundness
2
Presentation
2
Contribution
2
Limitations
Yes.
Looking forward to your feedback (if any)
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.
Thank you for the authors' response. We have reviewed your feedback and noted that most of our concerns have been addressed. However, we observed an issue with your Table 2 and Table 3 in the supplementary PDF. According to the official reports for LLaMA 3 and LLaMA 3.1 [1,2], the baseline performance that you report is lower than the official figures, especially for GSM8K and MATH in Table 2. The official tech report [1] indicates that LLaMA 3.1 8B achieves **84.5** on GSM8K (8-shot, CoT), whereas your report shows **17.8**. Similarly, the official MATH (0-shot, CoT) score is **51.9**, while your report indicates **8.6**. For Table 3, the tech report [2] shows that LLaMA 3 8B scores **66.7** on MMLU (5-shot). Based on this information, I believe your supplementary PDF lacks credibility and may negatively impact your paper. > Reference: > [1] https://ai.meta.com/blog/meta-llama-3-1/ > [2] https://huggingface.co/meta-llama/Meta-Llama-3.1-8B-Instruct
Authors' Reponse
Thank you for reading our rebuttal and confirming that most of your concerns have been addressed. We now provide clarifications in response to your further questions regarding the gap in evaluation scores. **A. Lower baseline GSM8K and MATH (in Table 2) than Meta's official values.** In the sequel, [1] and [2] refer to the two references the reviewer provided. The values reported in [1] are based on the finetuned model, i.e., Llama 3.1-8B-Instruct. This can be confirmed by comparing several benchmark scores in [1] with those of the instruct version in [2]. Namely, Meta just uses Llama 3.1-8B to represent Llama 3.1-8B-Instruct in [1]. In our provided Table 2 in the supplementary PDF, the "Base model" refers to the *pretrained base model*, i.e., Llama 3.1-8B. In addition, all of the reported scores in our Table 2 are obtained using *0-shot* prompting, which can also be different from that of 8-shot. Hence, we have different baseline performance from Meta's official values. In our Table 2, we use the evaluation code open-sourced by the MathInstruct paper [3] and use the same evaluation setup for all the models (base and finetuned models), ensuring a fair comparison between different optimization methods. **B. Lower baseline MMLU (in Table 3) than Meta's official value.** The few-shot MMLU score heavily depends on the *prompt engineering* and *evaluation approach*, as revealed by the Hugging Face report [4]. In this report, it explains in detail why open LLM leaderboard has a much lower MMLU score of the Llama model compared to the official one released by Meta. Since Meta does not open-source their prompts and evaluation methodology that yields their reported MMLU score, the open-source evaluation code may produce a lower MMLU score than Meta's official value. In our provided Table 3 in the supplementary PDF, we use the open-sourced MMLU evaluation code from Llama-Factory [5] and use the same evaluation setup for all the models (base and finetuned models). We would like to thank the reviewer once again for reading our rebuttal and raising further questions. We hope that the above response clarifies our results. If the reviewer has any additional questions, please let us know. **References:** [3] Yue et al. "Mammoth: Building math generalist models through hybrid instruction tuning", ICLR 2024. [4] Fourrier et al. "What's going on with the Open LLM Leaderboard?", Hugging Face Blog. [5] Zheng et al. "Llamafactory: Unified efficient fine-tuning of 100+ language models", ACL 2024.
Thank you for your detailed response. A: According to the Llama3 paper [6], Table 12, the Llama3-8B (the pre-trained model) achieves a score of **57.2** on GSM8K and 20.3 on MATH. Consequently, the Llama3.1-8B should exhibit performance that is at least equal to or greater than these values, which are still higher than the scores you have reported. B: I appreciate your clarification on this point and am satisfied with your response regarding the MMLU scores. Thank you once again for your thorough explanations. > References: > [6] The Llama 3 Herd of Models, https://arxiv.org/pdf/2407.21783
Authors' Response
Thank you for your careful reading of our response. We are pleased to see that your concern regarding the MMLU score has been addressed. We now provide further clarifications regarding the gap in the baseline GSM8K and MATH scores. Let us first remark that all the values reported in Llama 3 paper are for Llama 3.1 models, as indicated in [6, paragraph 3 of page 1]. Hence, the reported values in its Table 12 are for Llama 3.1-8B instead of Llama 3-8B. Importantly, their 57.2 GSM8K score is achieved using the *8-shot chain-of-thought prompting* technique, and their 20.3 MATH score is obtained through the *crafted 4-shot math prompting* technique. These setups can be found in [6, paragraph "Experiment setup" of page 29], in which they refer to their GitHub evaluation page. In summary, they use tailored techniques for almost each task to obtain the officially reported values. Instead, we use the default *0-shot standard prompting* setting for all the benchmarks, i.e., no additional prompts are added to the original input. The use of carefully designed chain-of-thought prompting may significantly enhance the benchmark scores. For instance, as displayed in the seminal chain-of-thought paper [7, Figure 2], the GSM8K score increases from 0.18 to 0.57 after the chain-of-thought prompting technique is applied to the same model. We would like to thank the reviewer once again for your thorough review of our results. We hope the response clarifies the gap in evaluation scores. If the reviewer has any additional questions, please let us know. **References:** [7] J Wei, et al. "Chain-of-Thought Prompting Elicits Reasoning in Large Language Models", NeurIPS 2022.
Thank you for your thorough response. If possible, could you please conduct a new inference experiment that aligns with the official benchmark settings? I believe this process would not take much time and would provide results that are more fair and convincing. When I first raised this concern, it should have prompted a reassessment of the experiment's validity and whether it was fair and meaningful, rather than attempting to convince the reviewer of an evidently questionable result.
Authors' Response
| Shot | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | Average | |------------|-------|-------|-------|-------|-------|-------|-------|-------|-------|----------------| | Base model | 17.8 | 46.9 | 45.0 | 53.0 | 54.2 | 51.2 | 55.6 | 54.9 | 56.7 | 48.4 | | LoRA | 48.7 | 61.0 | 57.2 | 59.0 | 59.7 | 58.7 | 60.1 | 59.4 | 59.8 | 58.2 | | BAdam | 49.6 | 60.5 | 61.9 | 59.6 | 59.4 | 59.1 | 61.6 | 61.0 | 59.2 | 59.1 | **Table 1: GSM8K score.** --- | Shot | 0 | 1 | 2 | 3 | 4 | Average | |-------|-------|-------|-------|-------|-------|---------| | Base | 8.6 | 15.7 | 19.7 | 19.7 | 19.8 | 16.7 | | LoRA | 13.7 | 22.8 | 23.2 | 23.0 | 22.8 | 21.1 | | BAdam | 17.2 | 23.2 | 22.8 | 22.5 | 22.8 | 21.7 | **Table 2: MATH score.** --- | Shot | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | Average | |-------------|-------|-------|-------|-------|-------|-------|-------|-------|-------|----------------| | Base model | 17.8 | 46.9 | 54.8 | 53.9 | 51.7 | 55.6 | 54.2 | 55.6 | 53.4 | 49.3 | | LoRA | 48.7 | 61.0 | 64.9 | 63.6 | 62.5 | 61.9 | 62.2 | 62.5 | 60.7 | 60.9 | | BAdam | 49.6 | 60.5 | 64.1 | 61.5 | 61.0 | 62.9 | 64.4 | 62.7 | 62.5 | 61.0 | **Table 3: GSM8K score after swapping the second and eighth CoT examples.** We apologize for misunderstanding the reviewer's request and greatly appreciate the suggestion regarding the experiment setup. We use the same prompt examples and template from Meta's GitHub instructions to conduct evaluations on the GSM8K and MATH benchmarks. The results are presented in Table 1 and Table 2. The scores for *all shots* are displayed for a comprehensive study. The base model achieves a 56.7 GSM8K score with 8-shot CoT, and a 19.8 MATH score with 4-shot CoT, which nearly match Meta's official values of 57.2 and 20.3, respectively. Although BAdam slightly outperforms LoRA on average, the scores of all three models vary across the $n$-shot results ($n \geq 1$), particularly for GSM8K. This indicates that more shots might not necessarily lead to better results, and adding one more CoT example (i.e., one more shot) may even negatively influence the outcome. Such variability leads us to question the stability of the metric and motivates us to examine the influence of different prompting strategies. To that end, we conduct a small ablation study on the GSM8K benchmark. Specifically, we alter the ordering of the prompts by swapping the second CoT example with the eighth, without changing any other content of these CoT examples. Surprisingly, *this seemingly minor change results in substantial score fluctuations*, as shown in Table 3. Note that this change is not due to randomness, as we have fixed the random seed and use greedy decoding rather than sampling during inference. Therefore, two tests should generate exactly the same scores once the evaluation setup is unchanged. Compared to Table 1, the overall performance of all three models clearly increases, as verified by the average scores. Additionally, LoRA performs almost as well as BAdam in this case. **Conclusions.** Based on these additional results, we make the following *preliminary* conclusions: * Compared to the base model, proper finetuning can improve the scores of benchmarks like GSM8K and MATH. However, under the few-shot CoT setting, we suspect that the dominant factor determining the score is the CoT examples and even the ordering of these examples, rather than the different finetuning optimization methods. That is, the effects of different optimization methods (such as LoRA and BAdam) diminish when few-shot CoT is utilized. This can be partly indicated by comparing the performance of model tuned by LoRA or BAdam with that of the base model under the 0-shot and few-shot settings. * Since the 0-shot setting excludes the effect of CoT, it could provide a simpler setup for comparing different optimization methods. In this case, BAdam appears to perform better than LoRA. We deeply thank the reviewer for the insightful discussion and constructive feedback, which helped us better understand these benchmarks. We hope these additional results address the reviewer's concerns. If the reviewer has any additional questions and/or suggestions, please let us know.
Thank you for your detailed response. Given the limited performance improvement, I have decided to maintain my current score.
The reviewer thanks the authors for the discussion. I decide to keep my score.
Decision
Accept (poster)