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

Authors: qijun  luo
hengxu  yu
xiao  li
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-12 08:58:53

Abstract: 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.

Main paper: 
 introduction: 
 Large language models (LLMs) such as GPT-4 [1] and Llama 3 [15] have shown its strong ability in language understanding, generation, reasoning, translation, etc [5, 73, 72, 60] . Due to its strong applicability, LLMs have been regarded as a feasible approach towards artificial general intelligence [6] . Finetuning or adaptation has become an important step in applying pretrained LLMs to follow human instructions or perform specific downstream tasks [44, 63] .Backgrounds. When GPU memory (RAM) is not a major limitation, full parameter finetuning methods-such as applying Adam to the entire set of parameters of LLMs-often offer the most flexibility for parameter search. However, executing such a full parameter training method typically requires a significant amount of GPU memory. For instance, to finetune an LLM with M billion parameters, Adam [23] necessitates roughly 18M GB of GPU memory for successful training, and this estimate does not even account for the storage of activations used in the backpropagation (BP) process; see Section 2.2.1 for a detailed analysis. This requirement poses challenges for computational resources as models scale up, given the fact that GPU memory is often limited in practical settings.Parameter efficient finetuning (PEFT) methods such as low-rank adaptation (LoRA) [22] , Adapter [21] , prompt-and prefix-tuning [29, 26] , among others, play a critical role in finetuning large language models under memory resource constraints. The principal idea of PEFT is to represent the parameter updates in a much lower-dimensional subspace and, consequently, the memory consumption is significantly reduced. Despite the success of PEFT methods, finetuning within a substantially lower-dimensional subspace may potentially limit downstream performance; see, e.g., [62] .Adam [23] 18M✓ ✓ Float32 ✓ LOMO [37] 2M + 2M D ✓ ✗ Float16 ✗ LoRA [22] 2M + 36rM m ✗ ✓ Float32 ✓ BAdam 2M + 16M D ✓ ✓ Float32 ✓Table 1 : Algorithm feature summary. Here, M represents that the model to be trained has M billion number of parameters, r is the LoRA rank, m is the weight matrix dimension (here, we consider square weight matrices for simplicity), D is the number of transformer layers in LOMO or the number of partitioned blocks in BAdam. BAdam performs full parameter mixed precision update, while only requires memory that is comparable to LOMO and LoRA.The observations outlined above motivate us to explore a memory efficient full parameter optimization method without imposing low-rank constraint on the parameter update.Main results. In this work, we have the following main contributions:(C.1) We propose a block coordinate descent (BCD)-type optimization method with Adam's update rule, termed BAdam; see Section 2.1 for the detailed description. This method partitions the entire set of model parameters into D blocks, updating one block at a time using Adam's efficient update steps. BAdam offers a memory efficient solution to the full parameter finetuning of LLMs. For example, by partitioning a model with M billion parameters into D equal-sized blocks, BAdam requires only about 2M + 16M D GB of GPU memory for successful mixed precision training; see Section 2.2.1 for detailed analysis. This leads to a significant reduction in memory demands compared to full parameter finetuning using Adam. Theoretically, we provide a convergence analysis for BAdam in the deterministic case, demonstrating that leveraging the BCD framework and Adam's update rule yields a convergent scheme; see Theorem 2.1. (C.2) 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. Specifically, we present in Section 3.1 BAdam's efficiency in both memory consumption and running time. In Section 3.2, we empirically verify BAdam's optimization capability via its fast convergence and the high rankness of its learned perturbations. We further evaluate the downstream performance of different methods using MT-bench and several math benchmarks; see Section 3.3. The results illustrate that BAdam generally outperforms existing memory efficient baselines such as LoRA. Importantly, BAdam achieves comparable average performance with Adam on math benchmarks and even surpasses Adam in instruction-following tasks evaluated by MT-bench score. Moreover, we conduct ablation study using SGD's update rule (BSGD) in Section 3.4. The results show that BCD variants maintain optimization capability compared to their full counterparts and even exhibit better downstream performance. It also demonstrates that BSGD can achieve similar downstream performance to BAdam, illustrating the effectiveness and suitability of BCD for finetuning LLMs.We compare BAdam with several representative methods in Table 1 . In summary, we believe that BAdam may serve as a viable alternative optimization method to state-of-the-art methods such as LoRA in scenarios with limited computing memory. 
the badam method: 
 Block coordinate descent (BCD) method has a long history in optimization society, which can be traced back to the very origins of the discipline; see, e.g., [43, 36, 4, 55, 41, 58] . At each iteration, BCD maintains the majority of the optimization parameters at their up-to-date iteration values, while it approximately optimizes the objective function over the remaining parameters, resulting in a much lower dimensional problem.BCD is known to be efficient for huge-scale problems where the number of optimization parameters is extensive [41] , particularly when it significantly exceeds the number of data points / component functions. Based on this main feature, we reveal an interesting link between BCD and the finetuning of LLMs. Namely, 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, providing the possibility to release the requirement on large GPU memory. We refer to Sections 3.2 and 3.4 for empirical verification on the power and suitability of BCD for finetuning LLMs.In this subsection, we propose BAdam, a block coordinate descent method embedded with Adam's update rule. The method is displayed in Algorithm 1 and illustrated in Figure 1 . Formally, let us consider an abstract formulation for finetuning LLMs min θ L(θ) = 1 n n j=1 ℓ j (θ). Here, θ ∈ R d represents the concatenation of the vectorized parameters of the model, n is the number of training data points, and ℓ j is the negative log-likelihood loss function for language modeling on the j-th training data point.Block partition and block coordinate descent framework. At the t-th block-epoch, BAdam first generates an ordered block partition π = {π 1 , . . . , π i , . . . , π D }, which splits the whole model parameters θ ∈ R d into D blocks, i.e., θ = {θ π1 , . . . , θ πi , . . . , θ π D } with θ πi ∈ R di and D j=1 d j = d. The partition π can be very flexible and is a unified representation. Given a large language model, one natural block partition is by transformer modules. Apart from this partition, one can also choose a small part of parameters from each transformer module and regard these parameters as one block θ πi . Note that π may be either a deterministic or a random partition, as long as the aggregation of all the blocks {θ πi } forms the whole set of parameters θ. For instance, if we partition by consecutive transformer modules, we may list the blocks in π in ascending order (e.g., from the input to the output module), descending order (e.g., from the output to the input module), or random reshuffling order.We now present the optimization framework of BAdam. Our core idea is to adopt the main spirit of BCD. Namely, we approximately optimize over only one active block θ πi at one time, given that the other inactive blocks are fixed at their up-to-date values. Mathematically, at the t-th block-epoch, updating the current active block θ πi amounts to solving the following subproblem:EQUATIONWhen this approximate minimization becomes exact, this scheme is also known as Gauss-Seidel method or alternating minimization. Even with exact minimization, some literature still refers to it as BCD. One can see that subproblem (1) fixes the inactive blocks at their most recent values, and hence Algorithm 1: BAdam: A block coordinate descent method with Adam's update rule. // compute the block stochastic gradient it is a much lower dimensional optimization problem compared to min θ 1 n n j=1 ℓ j (θ), providing the possibility to implement the method in situations with limited GPU memory. Solving subproblem (1) sequentially for i = 1, . . . , D moves the block-epoch from t to t + 1.9 g t,k πi ← stochastic approx. of ∂ ∂θπ i L(θ t+1 π1 , . . . , θ t+1 πi-1 , θ t,k-1 πi , θ t πi+1 , . . . , θ t π D ); 10 m t,k πi ← β 1 m t,k-1 πi + (1 -β 1 )g t,k πi , v t,k πi ← β 2 v t,k-1 πi + (1 -β 2 )(g t,k πi ) 2 ; 11 mt,k πi ← m t,k πi /(1 -β k 1 ), vt,k πi ← v t,k πi /(1 -β k 2 ) ; 12 θ t,k πi ← θ t,k-1 πi -α mt,k πi / vt,k πi + ε ; // AdamUpdate using Adam steps. Similar to most of the concrete BCD methods, we propose to implement the approximate minimization subproblem (1) using several gradient-based steps starting at θ t πi . Abstractly, BAdam executes the update θ t+1 πi ← A(θ t+1 π1 , . . . , θ t+1 πi-1 , θ t πi , θ t πi+1 , . . . , θ t π D ).(EQUATIONWe choose the algorithmic procedure A in (2) to be K Adam steps [23] starting at θ t πi , in order to efficiently decrease the objective function. To specify the concrete Adam steps, we first note that the gradient of the objective function can be correspondingly decomposed asEQUATIONWe call ∂L ∂θπ i in (3) the block gradient of the objective function L over block θ πi . According to the main spirit of stochastic optimization methods, we select a batch of data points to compute a block stochastic gradient g πi using the up-to-date iterates for approximating the block gradient, as outlined in Line 9 of Algorithm 1. With g πi , we construct the momentum and second moment for the active block θ πi as shown in Line 10 -Line 11. Finally, we implement Adam update in Line 12. One may also invoke decoupled weight decay [34] into Line 12. In summary, Line 6 -Line 15 concretely implement the BCD update (1) .It is important to note that BAdam differs from existing BCD with momentum approaches [41] , which often maintain dense momentum vectors. BAdam is specifically designed for memory efficiency, and hence it clears the optimizer states in Line 15. Additionally, we do not offload the optimizer states, as they will no longer correspond to the updated block parameters in the next block-epoch. Thus, clearing the states in Line 15 and starting with zero initial states for every new active block in Line 6 are crucial for ensuring convergence, stability, and memory efficiency of our method.The number K in Line 7 of Algorithm 1 is the only additional hyperparameter introduced by BAdam, compared to Adam. We provide a detailed discussion on selecting K in Appendix B.2.Convergence result. We provide a convergence analysis for BAdam in the deterministic case, aiming to establish that combining the block coordinate descent framework with Adam's update rule results in a convergent scheme. We consider the extension to the stochastic case as future work. Indeed, combining the analysis for Adam with stochastic gradients, as in [65, 27, 57] , with our analysis for the block coordinate descent framework could be a feasible direction for such an extension. The informal theorem is presented below, while the formal theorem and proofs are put in Appendix D.Theorem 2.1 (informal). BAdam using deterministic gradients is a descent method, under certain commonly utilized conditions for analyzing block coordinate descent method and Adam. That is, after one block-epoch of updates for the whole model, we haveEQUATION)Consequently, BAdam finds a δ-approximate stationary point within O(δ -2 ) iterations.We conclude this section by noting that BAdam is essentially a block coordinate descent method, in which the BCD framework achieves low memory consumption. Apart from the chosen Adam's update rule, it is possible to propose other efficient optimization procedures for concretely implementing (1); see Section 3.4 for an ablation study where we also employ SGD's update rule.We analyze the memory consumption of BAdam, caused by storing the model parameters, gradient, and optimizer states. Let us consider a large language model with M billion parameters. We will use GB as the unit of GPU memory in the sequel.We first analyze the memory cost of Adam with mixed precision training. One needs to store the FP16 model parameters for the BP process, which costs 2M memory. For a more precise update, the optimizer also maintains a master copy of a FP32 model, which costs 4M memory. Then, it comes to store the gradient (converted to FP32), momentum, and second moment in FP32 precision, costing 4M + 4M + 4M = 12M memory. In total, Adam needs roughly 18M memory.In terms of BAdam, it needs to store the up-to-date model parameters (see Figure 1 ) in FP16 precision, which costs 2M memory. Importantly, since BAdam only updates the active block at one time, we can store the model parameters, gradient, momentum, and second moment only for the active block θ πi in FP32 precision, where the FP32 model parameters and gradient of the active block can be obtained by transforming their FP16 versions to the FP32 versions. Let us consider the simple case where the partitioned D blocks are equal-sized. Then, BAdam only needs in totalEQUATIONNote that the above analyses do not account for the memory required to store activations, as this is associated with the BP process rather than the optimization method itself. Furthermore, gradient checkpointing [11] can be employed to reduce the memory requirement needed for storing activations. We display the actual memory consumption for finetuning the Llama 3-8B model in Section 3.1.We consider the specific case where the partitioned D blocks {θ πi } are D consecutive transformer modules of LLMs. Thanks to the property of backpropagation, BAdam can reduce the computation time of BP compared to Adam and LoRA under the same amount of data utilization.Let us consider one block-epoch of BAdam, meaning that it has trained with K • D data batches, where K is defined in Algorithm 1. We consider that each data point has the same sequence length and each transformer module has the same amount of parameters, in order to ease the analysis. Recall that a BP process consists of a forward pass and a backward pass. For the forward pass, BAdam has almost the same computational load as that of Adam, while LoRA requires more forward computation due to its extra low-rank adapters. Hence, it remains to consider the number of unit-backward-pass after utilizing KD data batches, where the unit-backward-pass is defined as a backward pass of a single data batch through a single transformer module. Importantly, BAdam only updates the active block, and hence the number of unit-backward-pass largely depends on the depth of the active block. For instance, if the input module or output module is the current active block, we need D unit-backward-pass or only 1 unit-backward-pass, respectively. Thus, after one block-epoch (i.e., utilizing KD data batches), BAdam requiresEQUATIONHowever, Adam and LoRA need to backward for all the D transformer modules, thus requiring KD 2 unit-backward-pass after utilizing KD data batches.Apart from saving the number of unit-backward-pass, some of the unit-backward-pass of BAdam may even take less computational time compared to that of Adam. Let us take the backward pass of the input module as an example. BAdam does not require explicit stochastic gradient computation of the model parameters of the intermediate modules ∂z l /∂θ l , where {z l } are the activations of the intermediate modules and {θ l } are the trainable model parameters of these modules. However, Adam needs to compute these quantity explicitly. We refer to Table 4 for an experiment illustration.In summary, BAdam with consecutive module-based block partition saves computational load of the BP process compared to Adam and LoRA, after training with the same amount of data. We demonstrate this through experiments detailed in Section 3.1. If the module-based block partition is not consecutive, for instance, when one block consists of modules (such as matrices) from different transformer layers, we still expect that BAdam can reduce BP time to some extent, though not as significantly as indicated by ( 6). 
experiment results: 
 In this section, we evaluate the proposed BAdam on finetuning LLMs. Selected baselines include LOMO (essentially SGD) [37] , LoRA [22] , Galore [66] , and Adam [23] . All BAdam experiments for training the Llama 2-7B and Llama 3-8B models are conducted on a single RTX3090-24GB GPU, whereas BAdam experiments for the Llama 3-70B model use 4×A100-80GB GPUs. Experiments for the baseline methods are conducted using either a single RTX3090 or multiple A100 GPUs, depending on their memory requirements. Our implementation is based on Llama-Factory [69] . Detailed experiment setup can be found in Appendix B.1.In this subsection, we present the empirically measured memory consumption and wall-clock running time of BAdam and baseline methods. All the measurements in this subsection are based on finetuning the Llama 3-8B model on Alpaca-GPT4 dataset [46] using a single RTX3090-24GB GPU.Memory consumption. We report the actual memory consumption of BAdam and the baseline approaches in Table 2 for finetuning the Llama 3-8B model, in which the memory consumption of Adam is estimated rather than tested. This result indicates that all of LOMO, LoRA (with a reasonable rank), and BAdam can finetune the Llama 3-8B model using a single RTX3090. It can be observed that all the methods have nearly the same memory cost for storing the model parameters, while LoRA requires slightly more memory due to its low-rank adapters. Furthermore, LOMO, LoRA, and BAdam significantly reduce memory consumption regarding the storage of the gradient and optimizer states compared to Adam. Moreover, it is easy to see that the total memory consumption (the last column of Table 2 ) is higher than the sum of the listed quantities. The additional memory costs arise from storing activations and training data, pre-allocated memory caches by PyTorch, and other buffers for intermediate computing results. Indeed, our tests show that BAdam can successfully finetune the Llama 3-8B model with input sequences of length 1024 using a batch size of 2, or input sequences of length 2048 using a batch size of 1, with a single RTX3090-24GB GPU.Wall-clock running time comparison. We conduct experiments on finetuning the Llama 3-8B model for 3 epochs with each method and report the averaged wall-clock time per epoch; see In Table 4 , we conduct tailored experiments to further support our analysis in Section 2.2.2. It can be observed that: 1) backward for "Output module only" is almost time free, as it requires only 1 unit-backward-pass; 2) backward for "All modules" takes significantly more time, as it has to implement D unit-backward-pass; and 3) backward for "Input module only" takes less time than D unit-backward-pass (i.e., backward for "All modules"), since the former scheme does not need to compute the stochastic gradients of the intermediate modules' parameters. We verify the optimization capability of BAdam through both the training loss convergence and the effective rank of the learned perturbations. Experiments in this subsection correspond to exactly the same training process of the lower block of Table 5 for finetuning the Llama 3-8B model.Loss convergence. In the left figure of Figure 2 , we display the online training loss. From a pure optimization perspective, namely, in terms of driving the training loss lower, BAdam demonstrates better convergence behavior than LoRA when using 1e-5 as the initial learning rate. If the initial learning rate is set to 1e-6, BAdam initially converges slightly faster, but the two methods soon align as the learning rate becomes too small to make substantial progress.Effective rank of the learned perturbations. We empirically measure the learning and optimization capability of BAdam through the effective rank of its learned perturbations, i.e., the difference between the learned weight matrix and the pretrained base weight matrix ∆W := W K -W 0 . The cumulative explained variance "cvar" and the effective rank of matrix A ∈ R m×n are defined as:cvar(r) := r i=1 σ i (A) 2 min{m,n} j=1 σ j (A) 2, effective_rank(A) := min {r : cvar(r) ≥ 0.9} , where σ i (A) is the i-th largest singular value of A.In the middle figure of Figure 2 , we display cvar of BAdam's update for the 25th layer's up-proj matrix. This result shows that BAdam's update has a heavy tailed singular values distribution and is far away from a low-rank update. In the right figure of Figure 2 , we plot the effective rank of the learned perturbation by BAdam and Adam through all modules of different transformer layers. Notably, BAdam achieves almost the same high rank update as Adam, which partly justifies BAdam's learning and optimization capability.In this subsection, we conduct supervised finetuning for the Llama 2-7B [54] , Llama 3-8B, and Llama 3-70B models on the Alpaca-GPT4 and MathInstruct [61] datasets. The setting of hyperparameters is deferred to Appendix B.2.MT-bench results. To illustrate the models' downstream performance, we report the MT-bench scores of the instruction-tuned models obtained by different methods for 3 epochs. We utilize two initial learning rates, 1e-5 and 1e-6, with a cosine scheduler for all methods. The results are displayed in Table 5 . (a) Adam with learning rate 1e-5 for finetuning Llama 3-8B overfits the Alpaca-GPT4 dataset and achieves MT-bench scores that are even lower than that of the base model. Hence, we omit this outlier. Math benchmarks. We also finetune the Llama 3-8B and Llama 3-70B models on the MathInstruct dataset for 3 epochs, and evaluate the trained model using math benchmarks across different domains.The results are shown in Table 6 . In terms of average score, BAdam outperforms all the memory efficient baselines, and even slightly surpasses the benchmark score of Adam while requiring significantly less memory consumption compared to Adam. In particular, for the experiments on finetuning Llama 3-8B, BAdam outperforms LoRA in 4 out of 6 tasks, and surpasses LOMO and Galore in all the tasks by a large margin. For finetuning Llama 3-70B, BAdam beats LoRA in 5 out of 6 tasks. In this subsection, we conduct ablation study to consider SGD's update rule in our BCD framework, leading to BCD with SGD (BSGD). Then, we compare the performance of BAdam, BSGD and their full counterparts, i.e., Adam and SGD, to illustrate the power of BCD in LLMs finetuing.Optimization. In the left and middle figures of Figure 3 , we display the training loss of BAdam, BSGD, and their full counterparts. 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 only updates one block of parameters per data batch, demonstrating the strong optimization ability of BCD for LLMs finetuning.Downstream performance. In the right table of Figure 3 , we test the MT-bench scores of the four methods. 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. We suspect that the superiority of the 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 improved downstream performance of BCD compared to their full counterparts further illustrate its suitability for LLM finetuning.We provide more experiment results in Appendix C. Here are the summarized results: 1) In Appendix C.1, we conduct an ablation study on the ordering scheme of the partition π in BAdam, considering random reshuffling, ascending, and descending orders. In terms of convergence behavior, these three choices are competitive. We also provide an ablation study on the hyperparameter K in BAdam, with K being chosen from {10, 50, 100, 200}. The results indicate that these four choices of K perform similarly in terms of convergence behavior. However, we observe that the convergence speed of choosing different K can vary across different models. We refer to Appendix B.2 for a detailed discussion on the selection of K. In summary, our experiment results in Section 3 demonstrate that BAdam has the potential to serve as a competitive optimization method for finetuning LLMs when the GPU memory is limited, compared to state-of-the-art memory efficient methods such as LoRA. 
brief literature review: 
 We briefly review several memory efficient finetuning methods in this section. A more comprehensive literature review is presented in Appendix A due to limited space.One major branch for memory efficient finetuning of LLMs is parameter-efficient finetuning (PEFT), which freezes the pre-trained weight and only trains the additional injected parameters. LoRA [22] , adapter [21] , and prefix-tuning [29] belong to this class and have been verified to be effective in finetuning LLMs. Another line of works focus on memory efficient full parameter finetuning. MeZO [38] performs zero-th order SGD update without calculating the stochastic gradient, thereby only requires the memory of performing inference. LOMO [37] efficiently performs on-the-fly SGD update during the backward pass without storing the stochastic gradient. However, LOMO's implementation design prevents it from using gradient accumulation technique. Galore [66] reduces memory consumption by projecting the gradient into low-rank space. It requires constantly performing SVD to obtain the low-rank projector. 
conclusion and discussions on limitations: 
 In this work, we have proposed the BAdam optimization method, which is built upon the block coordinate descent framework with Adam's update rule. We finetune the Llama 3-8B and Llama 3-70B models on the Alpaca-GPT4 and MathInstruct datasets by BAdam with a single RTX3090-24GB GPU and 4×A100-80GB GPUs, respectively. The results illustrated the efficacy of BAdam in terms of GPU memory consumption and running time. Empirically, BAdam exhibits better convergence behavior compared to LoRA and learns high rank update. Further downstream performance assessments have demonstrated BAdam's superior performance in instruction finetuning and math finetuning, in comparison to LOMO, LoRA, and Galore. Additionally, BAdam has on par or even better downstream performance compared to Adam. In summary, we believe that BAdam may serve as a viable alternative for finetuning LLMs with limited memory resources.Limitations. Our focus has been on applying BAdam for supervised finetuning. Extending its application to preference optimization represents another opportunity to demonstrate BAdam's capabilities. Moreover, our CPT experiment using BAdam is only preliminary. Exploring extensively BAdam's performance in the CPT setting is an interesting direction. We leave these directions for future improvements.Broader impacts. Our proposed method significantly lowers the barrier to full parameter finetuning of large models for a broader range of researchers. This is a technical algorithmic contribution that does not yield explicit negative societal impacts. However, it carries a risk of misuse. A More Related WorksWe present a review of the relevant literature below. Given the extensive and rapidly growing body of work in this field, it is important to note that the references we include here are not exhaustive.Block coordinate descent method. The block coordinate descent (BCD) method is a well-established algorithmic scheme in the field of optimization [43, 36, 4, 55, 41, 58] , which is especially efficient for problems with an exceptionally large number of trainable parameters. Some advanced analyses on its convergence in both convex and nonconvex cases can be found in, e.g., [50, 35, 42, 67, 7, 8] and the references therein. The works [52, 16, 51, 10, 39 ] also theoretically discuss the effects of different choices of the block update order in BCD and ADMM-BCD variants, which is related to the ordering scheme in the partition π of BAdam. BCD has been widely and practically applied in the machine learning area as well. For instance, a specific instance of BCD, known as layer-wise training, has been utilized for neural network training [20, 3, 2, 71, 45] . The sequential minimal optimization technique integrated into LIBSVM [9] is also a BCD-type method.Parameter efficient finetuning (PEFT). An effective strategy for finetuning LLMs is to train a small number of (possibly extra) model parameters, while keeping the majority of the pretrained parameters frozen. Numerous approaches have been proposed and studied along this line of research.For instance, adapter tuning only finetunes the inserted small modules between layers called adapters [21] . Prompt-tuning / prefix-tuning [26, 29] attaches additional trainable prefix tokens to the input and/or hidden layers, while remaining the base model unchanged. Another prevalent method is the low-rank adaptation (LoRA) [22] , which models the increment to the base model as a product of two significantly lower dimensional trainable low-rank matrices. Subsequent research on LoRA has aimed at extending its rank constraints [30, 59] , further reducing the number of trainable parameters [24, 25] , decreasing memory usage through quantization [14] , etc. Presently, LoRA-type methods are commonly employed for finetuning LLMs with limited memory resources. Interested readers are referred to [17] for a unified framework and comprehensive comparison of these PEFT methods.Memory efficient full parameter finetuning. To conduct full parameter finetuning of LLMs with limited memory, the work [37] proposes LOMO, which efficiently leverages the BP process to update parameters on the fly in the process of computing stochastic gradients. Consequently, LOMO helps to execute SGD for full parameter finetuning without physically storing the stochastic gradients, significantly reducing memory consumption. However, it is worth emphasizing that SGD generally converges more slowly and is often considered suboptimal compared to Adam. MeZO [38] is to approximate SGD by using only the forward pass. The idea of MeZO derives from zeroth-order optimization, which utilizes function value difference to approximate the stochastic gradients of the trainable model parameters. Galore [66] uses gradient low-rank projection, which largely reduces memory consumption for full parameter finetuning compared to Adam. Adam-mini [64] proposes to apply block-wise adaptive learning rate, which reduces the memory for storing the full second moment. Another popular approach for finetuning with limited memory is to perform CPU offloads to reduce the memory consumption caused by training data and optimizers; see, e.g., [49, 48, 33] .• For all the experiments of finetuning Llama models, we apply the gradient checkpointing technique [11] to reduce the memory cost of storing activations for all methods. In particular, we checkpoint the input of each transformer layer and re-forward to calculate the layer's activations during the backward phase. • The experiments on finetuning Llama models are implemented using Llama-Factory [69] .The SuperGLUE experiments are based on the implementation of jiant [47] . 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: We do not release any data or model in this work. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: The used datasets, codes, and models are properly cited based on their licenses. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.• For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [Yes] Justification: We opensource our code implementation publicly and provide detailed documentation on how to use our code.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA]Justification: The paper does not involve any crowdsourcing experiments or any research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA]Justification: The paper does not involve any crowdsourcing experiments or any research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: We summarize the choices of hyperparameters for SuperGLUE, instruction tuning, and math finetuning in Table 8 , Table 9 , and Table 10 , respectively. Some discussions follow:• BAdam introduces only one additional hyperparameter compared to Adam, namely, K in Algorithm 1. One natural choice is K = n BD , where n is the number of training data points, B is the effective batch size, and D is the number of blocks in BAdam. We round n BD to the nearest integer if it is a fractional. Such a setting ensures that after one block-epoch, all the n training data points are equally distributed to the D blocks for training. On another front, too small K may result in insufficient Adam steps, while too large K may over-optimize one block before moving to others. Therefore, one possible choice isEQUATIONThe suggested selection is merely a guideline. One may also choose other reasonable values for K based on n BD , once each of the D blocks is appropriately trained with a reasonable amount of data.• For applying LoRA to finetune the Llama 2-7B model, we set LoRA rank to be 100 so that it has as many trainable parameters as BAdam at each iteration. However, when finetuning the Llama 3-8B model, using rank 100 results in out-of-memory error, as shown in Table 2 .We observe that LoRA with rank 8 and rank 100 achieve similar MT-bench scores for finetuning the Llama 2-7B model, corroborating the conclusion in LoRA paper [22, Section 7.2] that the LoRA rank does not essentially affect its performance. Therefore, we report the performance of LoRA with rank 8 for finetuning the Llama 3-8B model. We display the online loss under different block switch strategies and inner Adam step choices for each block sub-problem in Figure 4 . The results are based on finetuning Llama 3-8B on Alpaca-GPT4 dataset.We compare three block ordering strategies including descending (from the output to the input module), ascending (from the input to the output module), and random reshuffling. As shown in Figure 4a , although the descending scheme initially converges more slowly, all three ordering schemes finally exhibit nearly identical convergence behaviors. As shown in Figure 4b , different choices of Adam steps K does not affect the convergence of online training loss evidently for the task of finetuning Llama 3-8B. However, we notice that the convergence speed may vary across different models for different values of K. One can refer to (7) for a detailed discussion on selecting K. We test the performance of BAdam on classification tasks by training RoBERTa-large [32] on the SuperGLUE benchmark [56] . The implementation is based on jiant [47] using a single RTX3090. The setting of hyperparameters is put in Appendix B.2. We display the test results on 6 tasks selected from the SuperGLUE benchmark. We choose these tasks to conduct experiments since they are selected in [37, 38] . The results can be found in Table 11 . It can be observed that BAdam outperforms LoRA in 5 out of the 6 tasks. Furthermore, BAdam demonstrates performance that is comparable to, or tied with, Adam. Based on these results, we can conclude that BAdam is capable of closing the performance gap with Adam more efficiently than LoRA. Consequently, we extrapolate that BAdam has the potential to perform nearly as well as Adam, even when finetuning larger models.We conduct a preliminary continue pretraining experiment on StarCoder-Python [28] dataset, which consists of repositories from GitHub, including GitHub issues and commits. We train Llama 3.1-8B-Instruct on Starcoder dataset using BAdam, with learning rate 1e-5 and batch size 160. The training loss is shown in Figure 5 . One can see that BAdam effectively decreases the CPT loss to around 0.89 after being trained for 10 billion tokens (about 1 epoch). We present the memory consumption for finetuning the Llama 2-7B model in Table 12 . We can see that all the methods successfully finetune Llama 2-7B within 24GB memory, and the cost for each part corroborates our interpretation in Section 2.2.We also display the time costs for finetuning the Llama 2-7B model in Table 13 and Table 14 . One can see that LoRA requires approximately twice the forward time cost when fine-tuning the Llama 2-7B model. This additional cost may be attributed to implementation-level overhead. Notably, finetuning the Llama 2-7B model requires a much longer training time compared to that of the Llama 3-8B model, as the latter utilizes grouped query attention that significantly improves both inference and backward efficiency. To establish the convergence of BAdam, we first prove a descent inequality for updates applied to one block by bounding the error terms introduced by Adam updates. Integrating these descent inequalities across different blocks shows that BAdam is a descent method and has a complexity bound of O(δ -2 ).For a compact and clear convergence analysis, we focus on BAdam with deterministic gradients, leaving the stochastic cases for future work.We make the following two assumptions. Assumption D.1 is standard for analyzing block coordinate descent-type methods [58] . Assumption D.2 is commonly used in the analysis of Adam [13] . We adopt this assumption for simplicity of presentation, noting that it can be provably ensured [27] .Assumption D.1. The loss function L is L-smooth. And when restricted on i-th block, it is L i -smooth. Mathematically,∥∇L(θ 1 ) -∇L(θ 2 )∥ ≤ L∥θ 1 -θ 2 ∥, ∂L ∂θ i θ 1 i - ∂L ∂θ i θ 2 i ≤ L i ∥θ 1 i -θ 2 i ∥, i = 1, . . . , D.We define parameters L = max i=1,...,D L i as the maximum smoothness constants across all blocks and L = min i=1,...,D L i . Assumption D.2. BAdam has bounded partial derivatives along its trajectory, i.e., there exists G > 1 such that ∥g t,k i ∥ ≤ G. Here, we adopt the notations as specified in Algorithm 1: t = 0, . . . , T are block epochs, i = 1, . . . , D are different blocks, k = 0, . . . , K are inner iterations over data for a certain block, and g t,k i = ∇ i L(θ t i ) = ∂L ∂θ t,1 2G I ≼ H t,k i ≼ 1 λ I.Proof. By definition of v, its elements (v t,k i ) j are exponential moving average of the history gradients' squared elements (g t,k i ) 2 j , so we have boundmax 1≤j≤d (v t,k i ) j ≤ max 1≤j≤d (g t,k i ) 2 j ≤ ∥g t,k i ∥ 2 ≤ G 2 .Therefore 1(v t,k i ) j + λ ≥ 1 G + λ ≥ 1 2G , ∀1 ≤ j ≤ d, then H t,k i -1 2G I is a positive semidefinite matrix. Similarly, since 1/( (v t,k i ) j + λ) ≤ 1/λ, we have H t,k i ≽ 1 λ I.Here we formally present the theorem for the convergence of BAdam in Section 2.1. This section will focus on providing the proof for the theorem. Theorem D.4 (descent method). Under Assumption D.1and Assumption D.2 and suppose that 0 < λ < 1. If the learning rate satisfies the inequality α ≤ λ 2 LK min 1 K , λ 12G , then BAdam with deterministic gradients has the following descent property after one block-epoch of updates:L(θ t+1 ) -L(θ t ) ≤ - αK 16G 1 + 2α 2 KL 2 D λ 2 4 L2 α 2 K 3 λ 6 + 1 ∥∇L(θ t )∥ 2 .Corollary D.5 (first-order complexity). Under the conditions in Theorem D. 4 T = 128D 1/4 LG(L(θ 0 ) -L * ) δ 2 λ max D 1/4 K, 6G λ gradient evaluations.Proof. With the above choice of learning rate and Theorem D.4, the descending property of one block epoch can be written asL(θ t+1 ) -L(θ t ) ≤ - αK 32G ∥∇L(θ t )∥ 2 .Sum over t, we obtain the bound for minimum gradient normmin t≤T ∥∇L(θ t )∥ 2 ≤ 1 T + 1 T t=0 ∥∇L(θ t )∥ 2 ≤ 32G(L(θ 0 ) -L * ) αK .Substitute the choice of α and we get the complexity result.To simplify notation, in the following we abuse notation L(θ L(θ t i ) -L(θ t i-1 ) ≤ - αK 2G 1 2 - 2L i αKG λ 2 ∥∇ i L(θ t i )∥ 2 + αG + L i α 2 K 2 ∥e t i ∥ 2 ,where we denote∥e t i ∥ = ∥ 1 K K k=1 1 1-β k 1 H t,k i (m t,k i -(1 -β k 1 )g t,1i )∥ as the difference between the updates of Adam and full GD scaled by coordinate wise adaptive step sizes.Proof. With Assumption D.1, we have the following descent inequality:L(θ t i ) -L(θ t i-1 ) ≤ ⟨∇ i L(θ t i ), θ t+1 i -θ t i ⟩ + L i 2 ∥θ t+1 i -θ t i ∥ 2 ≤ -α⟨∇ i L(θ t i ), Ke t i + K k=0 H t,k i g t,1 i ⟩ + L i α 2   K 2 ∥e t i ∥ 2 + K k=1 H t,k i g t,1 i 2   ≤ - αK 2G ∥∇ i L(θ t i )∥ 2 + αK 4G ∥∇ i L(θ t i )∥ 2 + αKG∥e t i ∥ 2 + L i α 2 K 2 ( 1 λ 2 ∥∇ i L(θ t i )∥ 2 + ∥e t i ∥ 2 ) = - αK 2G 1 2 - 2L i αKG λ 2 ∥∇ i L(θ t i )∥ 2 + αKG + L i α 2 K 2 ∥e t i ∥ 2 ,where the last inequality is because Corollary D.3 and Cauchy-Schwarz inequality.We further have the following lemma, which provides a bound for the above error term. Lemma D.7 (bound for error term). Under the conditions in Theorem D.4, we have bound for ∥e t i ∥:∥e t i ∥ ≤ 2L i αK λ 2 ∥g t,1 i ∥.Proof. By definition in Algorithm 1 we have the expression for inner updates:EQUATIONSo the error term can be written as:∥e t i ∥ = 1 K K k=1 1 1 -β k 1 H t,k i (m t,k i -(1 -β k 1 )g t,1 i ) ≤ 1 K K k=1 1 λ(1 -β k 1 ) ∥m t,k i -(1 -β k 1 )g t,1 i ∥ ≤ 1 K K k=1 1 λ ∥(1 -β 1 ) k j=1 β k-j 1 (g t,j i -g t,1 i )∥ ≤ 1 -β 1 Kλ K j=2 ∥g t,j i -g t,1 i ∥ K k=j β k-j 1 ≤ L i Kλ K-1 j=1 ∥θ t,j i -θ t,0 i ∥,where the first inequality is because Corollary D.3 and the second inequality is by definition of m t,k i in (8) .Denote ∆ t i = K-1 j=1 ∥θ t,j i -θ t,0 i ∥, now let's bound ∆ t i . ∆ t i = K-1 j=1 ∥θ t,j i -θ t,0 i ∥ = K-1 j=1 α j k=1 H t,k i mt,k i ≤ K-1 j=1 α j k=1 1 λ(1 -β k 1 ) (1 -β 1 ) k l=1 β k-l 1 g t,l i = K-1 j=1 α j k=1 1 λ(1 -β k 1 ) (1 -β 1 ) k l=1 β k-l 1 (g t,l i -g t,1 i ) + (1 -β k 1 )g t,1 i ≤ K-1 j=1 j k=1 α λ(1 -β k 1 ) (1 -β 1 ) k l=1 β k-l 1 L i ∥θ t,l i -θ t,0 i ∥ + (1 -β k 1 )∥g t,1 i ∥ ≤ K-1 j=1 j k=1 α λ(1 -β k 1 ) (1 -β 1 )L i ∆ t i k l=1 β k-l 1 + (1 -β k 1 )∥g t,1 i ∥ ≤ α λ L i K 2 ∆ t i + α λ K 2 ∥g t,1 i ∥.In the above, the first inequality is because Corollary D. EQUATIONProof. Substitute Lemma D.7 into Lemma D.6 and we haveL(θ t i ) -L(θ t i-1 ) ≤ - αK 2G 1 2 - 2L i αKG λ 2 - 8L 2 i α 2 K 2 G 2 λ 4 - 8L 3 i α 3 K 3 G λ 4 ∥∇ i L(θ t i )∥ 2 ≤ - αK 8G ∥∇ i L(θ t i )∥ 2 ,where the last inequality is because α ≤ λ 2 24 LKG . Now we are ready to prove Theorem D.4.Proof of Theorem D.4. Sum up (9) over i, we haveEQUATION)In the following, we use the right hand side D i=1 ∥∇ i L(θ t i )∥ 2 to upper bound the whole gradient vector of one block epoch ∥∇L(θ t )∥ 2 .EQUATIONwhere the second inequality is because the partial derivative vector ∇ i L(θ t i ) is part of the gradient vector ∇L( θt i ) and the last inequality is by L-smoothness of L. For the second term we have∥ θt i -θ t ∥ 2 = i j=1 ∥θ t+1 j -θ t j ∥ 2 = i j=1 α 2 Ke t j + K k=0 H t,k j g t,1 j 2 ≤ i j=1 2α 2 K 2 ∥e t j ∥ 2 + i j=1 2α 2 K 2 λ 2 ∥∇ j L(θ t j )∥ 2 ≤ i j=1 2α 2 K λ 2 ( 4L 2 i α 2 K 3 λ 6 + 1)∥∇ j L(θ t j )∥ 2 .Substitute into (11) and sum over, we have∥∇L(θ t )∥ 2 = D i=1 ∥∇ i L(θ t )∥ ≤ 2 1 + 2α 2 KL 2 D λ 2 4 L2 α 2 K 3 λ 6 + 1 D i=1 ∥∇ i L(θ t i )∥ 2 .Substitute back into (10) and we get our descent inequality for a whole block epoch.Answer • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).https://github.com/TIGER-AI-Lab/MAmmoTH 


Decision: This paper looks into the problem of how to perform memory-efficient fine-tuning of LLMs. LLM training is memory-intensive, due to the large volume of parameters and the memory for storing optimizer states. To address this issue, the paper proposes BADAM, a method that leverages the block coordinate descent framework with Adam as the inner solver. Basically, it partitions the model parameters into blocks and updates one block at a time using Adam steps. By doing this, the max memory footprint of fine-tuning LLMs is reduced. 

Overall, this is a borderline paper. While BAdam can potentially become a good alternative way for fine-tuning LLMs with limited memory capacity. There are a few concerns:
1) There have already been many memory-efficient fine-tuning methods developed, such as LoRA, GaLore, etc. Empirical comparison with these methods shows that BAdam does not have significant advantage over those methods despite it supports full-parameter optimization.
2) Full-parameter optimization requires saving checkpoints of the full model states, which is appearing only if the performance improvement is great. However, the improvement seems to be marginal.  
3) Optimizations are sensitive to hyperparameter choices. The paper uses pre-selected parameters for alternative baselines. Given that there are many hyperparameter options, it is hard to say the performance improvement is mainly from the BAdam approach itself. 
4) Even for full parameter optimization, there are many alternative systems optimizations, such as offloading optimizer states to secondary memory or layer-by-layer fine-tuning. There has been a lack of discussion of the proposed method wrt to those optimizations. 

Given these, the paper can be accepted only if there are additional space.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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.
limitations: limitations: Yes.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: The author has addressed the limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: - **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: questions: N/A
limitations: limitations: N/A
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: See my weakness part.
limitations: 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.

Comments: title: Thank You
Thank you
Comment on Scores of GSM8K and MATH
Looking forward to your feedback (if any)
Authors' Response
Authors' Response
Authors' Response
Authors' Response
Authors' Reponse
comment: 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 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.
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
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 your detailed response.

Given the limited performance improvement, I have decided to maintain my current score.
| 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 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.
Thank you for your careful review of our paper and your positive assessment. If you have any further questions, please let us know.
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.
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.
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.
The reviewer thanks the authors for the discussion. I decide to keep my score.
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
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 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