BAdam: A Memory Efficient Full Parameter Optimization Method for Large Language Models
Authors: Qijun Luo
Hengxu Yu
Xiao Li
Venue: NeurIPS 2024 poster (2024)
Keywords: block coordinate descent
large language models
Primary area: optimization
TL;DR: This work designs a memory efficient block coordinate descent optimization method for finetuning LLMs, which effectively finetunes Llama 3-8B and Llama 3-70B using a single RTX3090-24GB GPU and 4 A100-80GB GPUs, respectively.

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.

BAdam: A Memory Efficient Full Parameter Optimization Method for Large Language Models
Qijun Luo Hengxu Yu!
Xiao Li
'The Chinese University of Hong Kong, Shenzhen Shenzhen Research Institute of Big Data
{qijunluo www.ghnnggn 0link. cuhk edu. en, lixiao0cuhk, edu. on
Abstract
This work presents BAdam, an optimization method that leverages the block coordi- nate descent (BCD) framework with Adam's 's update rule. BAdam offers memory efficient approach 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 single RTX3090-2440 tively. The results confirm BAdam's efficiency terms of memory usage, running time, and optimization capability. Furthermore, the downstream performance evaluation based on MT-bench and math benchmarks shows that BAdam outper- forms 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 I.L.Ms. Our code can be easily integrated into any Py Torch-based codebase and is available https: httptttt.t../////////////t/e.eet thub. coom/eddy//BB /BAdam.
GPU and 4 x A100-80GB GPUs, respec-
1 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 applicabilityy LLMs have been regarded as a feasible approach towards artificial general intelligence [6]. Finetuning adaptation has become important step in applying pretrained L.LMs 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 parameters LLMs- offer the most flexibility for parameter search. However, executing such 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; Section 2.2.1 for 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 low-rank adaptation (LoRA) [22], Adapter [21], prompt- and prefix-tunill peeiinttunnng [29, 26), among others, play critical role in finetuning large language models under memory resource constraints. The principal idea of PEFT is 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 power-dimensional subspace may potentially limit downstream performance; see, e.g., [62].
"Corresponding Author
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Full parameter training
Momentum and second moment
Gradient
Method
Memory
Update precision
accumulation
Adam [23]
18M
Float32
LOMO [37]
2M +
2M
D
/
X
Float16
x
LoRA [22]
2M
36rM
X
Float32
m
BAdam
2M
16M
Float32
Table 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 memory efficient full parameter optimization method without imposing low-rank constraint on the parameter update.
Main results. In this work, 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 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 significant reduction in memory demands compared 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 convergent scheme; see Theorem 2.1.
(C.2) We apply BAdam to finetune the Ilama 3-8B and Llama 3-70B models using single RTX3090-24GB GPU and 100-800- GPUs, respectively. Specifically, we present in Section 3.1 BAdam's 's efficiency in both memory consumption and running time. In Section 3.2, we empirically verify BAdam's '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 Sec- 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-folloviii evaluated by MT-bench score. Moreover, we conduct ablation study using SGD's '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.
tasks
We compare BAdam with several representative methods Table 1. In summary, we believe that BAdam may serve as a viable alternative optimization method to state-of of-the-art methods such as LoRA scenarios with limited computing memory.
2 The BAdam Method
Block coordinate descent (BCD) method has a long history in optimization society, which can be traced back to the very origins 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

---

BAdam: A block coordinate descent method with Adam's update rule
Block-epoch -
0
0
0
6
0441
Subproblem the block fr.
Concrete implementation: Adam sleps starting 02
61-1
approx argmin 9e, 634
0141 Pn.
Figure Illustration of the proposed BAdam, which based on the block coordinate descent framework. Colors represent the states of the partitioned blocks in one block-epoch, including the active block, non-updated inactive blocks, and updated inactive blocks.
functions. Based on this main feature, we reveal an interesting link between BCD and the finetuning of LLMs. Namely, the finetuning process boils down optimization problem that needs to handle a huge number 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 the power and suitability of BCD for finetuning LLMs.
2.1 Algorithm Description and Convergence Result
In this subsection, we propose BAdam, a block coordinate descent method embedded with Adam's update rule. The method is displayed in Algorithm and illustrated in Figure 1. Formally, let us consider an abstract formulation for finetuning LLMs ming L(O) A Σÿ= (0). Here, 0 Rd represents the concatenation the vectorized parameters of the model, n is the number of training data points, and lj 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 π {TI..... , TD which splits the whole model parameters 8 E Rd into D blocks, i.e., 0 = {0..... (f.,........ Or A............................... 0r,} with fr, Rd and E dz 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 0z.- Note that æ may either deterministic or random partition, as long as the aggregation of all the blocks (fz,) forms the whole set of parameters f. 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 adopt the main spirit of BCD. Namely, we approximately optimize over only one active block Oz, al one time, given that the other inactive blocks are fixed at their up- up-to-date values. Mathematically, at the t-th block-epoch. updating the current active block 0x, amounts to solving the following subproblem:
ge+1 approx argmin On, ER4:
Σ 1=1
(00+
01+1 ,0..1 Te-1 TL
On Hr
6² 0z
(I)
When 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 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.
input: 31, B2, e, K, and learning rate o.
2 initialization: block-epoch index & <- and model parameters 09.
3 while stopping criterion not meet do
generate block partition स = { ; repeat block-epoch i < 1, I..... D
// BCD loop // Block initialization
k: 4 0; m <- O; u +- 0: R,
t K,
7
repeat for K Adam steps to update the active block 0n.
1: + k +
// compute the block stochastic gradient
g²,²" + stochastic approx. of 0 C(0+t) ...., G++1 0th-1 0t a0.
TT = 0 v B2
mik + -im== + (1- B1)g
10
((1 B²), i + L.L T.
aL. t v.. ;
11
/ (Vott +s):
12
0c. + 02 arii / m
1/ Adam update
13
end
gc+1 1 +- gt.K.
14
me
nt
15
IR. Ma Un; f None
// clear memory for grad and optim states
end
17
1 1;
18
end
19 return learned model parameters 0²
is a much lower dimensional optimization problem compared ming 1 Σ (0), providing n the possibility to implement the method in situations with limited GPU memory. Solving subproblem (I) sequentially for = I,..., D moves the block-epoch from f to é + 1.
Update 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 xx Abstractly, BAdam executes the update
g-+ A(
6 0- Te..
(2)
Ti-,
We choose the algorithmic procedure in (2) be K Adam steps [23] starting at 0%, 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 as
CC())=
Oc 30..
ac Jue
à 0 30,
2 T E. {((0)] 20
(3)
We call aC in (3) the block gradient of the objective function L over block 0... According to the 8e., main spirit of stochastic optimization methods, we select a batch of data points to compute a block stochastic gradient Sr, using the up-to-date iterates for approximating the block gradient, as outlined in Line 9 of Algorithm 1. With 9m., we construct the momentum and second moment for the active block 9., 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 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 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 of Algorithm is the only additional hyperparameter introduced by BAdam, compared to Adam. We provide detailed discussion 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 's update rule results

---

in 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 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 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 have
[(g)+))- - -(())) S -((aK))
(4)
Consequently, BAdam finds a 8-approximate stationary point within O iterations.
We conclude this section by noting that BAdam essentially block coordinate descent method, in which the BCD framework achieves low memory consumption. Apart from the chosen Adam's 's update rule, it possible to propose other efficient optimization procedures for concretely implementing (I);
see Section 3.4 for an ablation study where we also employ SGD's update rule.
2.2
Analysis of Memory Consumption and BP Time
2.2.1 Memory Consumption Analysis
We analyze the memory consumption of BAdam, caused by storing the model parameters, gradient, and optimizer states. Let us consider 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 FP32 model, which costs 4.M 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, can store the model parameters, gradient, momentum, and second moment only for the active block 0., 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 total
16M D
2M +
memory.
(5)
Note that the above analyses do not account for the memory required to store activations, as this associated with the BP process rather than the optimization method itself. Furthermore, gradient checkpointing [11] I] 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 1.
2.2.2 BP Time Analysis for Consecutive Module-based Block Partition
We consider the specific case where the partitioned D blocks (0,,) } 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 ease the analysis. Recall that BP process consists of forward pass and 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 K D data batches, where the mait-backward-pass is defined as backward pass of a single data batch through single transformer module. Importantly, BAdam only updates the active block, and hence the number of barebbbbwwwwddpass largely depends on the depth of the active block. For instance, if the input module output module is the current active block, we need D
5

---

mit-backward-pps -asss or only 1 unit-backward pass, respectively. Thus, after one block-epoch (i.e., utilizing K D data batches), BAdam requires
((1 + +... +..+++++ +) KD(D + +1)
unit-backward-pass.
(6)
2
However, Adam and LoRA need to backward for all the D transformer modules, thus requiring KD2 unit-backward-pass after utilizing KD data batches.
Apart from saving the number of unit-backward-past some of the unit- 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 example. BAdam does not require explicit stochastic gradient computation of the model parameters of the intermediate modules / where { are the activations of the intermediate modules and {0r 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 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).
3 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 A 100-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 Llama-Factory [69]. Detailed experiment setup can be found in Appendix B.1.
3.1 Memory Consumption and Wall-clock Running Time
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 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 reasonable rank), and BAdam can finetune the Llama 3-8B model using 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 Adam. Moreover, 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 Py 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 batch size of 1, with 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 Table 3. The forward time for three approaches are rather close. The slightly higher time cost for LOMO and LoRA attributes to additional operations for registering activations and the calculation of the low-rank adapters, respectively. Regarding backward time, BAdam reduces the time cost by nearly half compared to LoRA and LOMO, supporting the analysis in Section 2.2.2. It is important to note that the backward time for all methods includes the re-forward time due to gradient checkpointing, which diminishes the running time advantage of BAdam.

---

Method
Parameter
Gradient
Optimizer states
Memory consumption
Adam
I6.I
32.IGB
96.6GB
144.8GB+
LOMO
16.1GB
0.5GB
21.5GB
1oRA-rank
16.7GB
1.0GB
3.IGB
0OM
LoRA-rank8
16.2GB
0.1GB
0.3GB
22.3GB
BAdam
16.1GB
0.9GB
2.6GB
23.5GB
Table 2: Actual memory costs of applying mixed precision training to finetune Llama 3-8B with gradient checkpointing using single RTX3090. Note that LOMO only supports FP16 precision training. The maximum input sequence length is 728 and the batch size is 2.
Method
Forward
Backward
Update
Backward scheme
Backward time
LOMO
0.78 hours 3.70 hours
All modules
0.64 seconds
LoRA
0.83 hours
3.20 hours
136 seconds
Input module only
0.33 seconds
BAdam
0.71 hours
1.74 hours
142 seconds
Output module only
0.03 seconds
Table 3: Time spent per epoch on forward, back- ward, and update for finetuning Llama 3-8B using single RTX3090. The single pass batch size is 2. The results are averaged over 3 epochs.
Table 4: Time spent on different backward schemes with batch size for finetuning Llama 3-8B using single RTX3090. The re- sults are averaged over 100 backward passes.
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, it requires only 1 sitt-bacwward-passae 2) backward for "All modules" takes significantly more time, as has to implement D unit-backward -ass:; and 3) backward for "Input module only" takes less time than D unit-backward-pea (i.e., backward for "All modules"), since the former scheme does not need to compute the stochastic gradients of the intermediate modules' parameters.
3.2
Optimization Capability
Thaiiing
Bllecikeraak the pertubation
Lofo, 1res
0.5
3.2
2000
4
Admm zAdum
1500
1.0
500
2020 4000
800g
10000
Cata 20.8 of
Numvber of /irgu AOA
Figure 2: Optimization capability of BAAam for finetuning Llama 3-8B Alpaca-GPT4 dataset. Left: Online training loss LoRA and BAdam. Middle: Cumulative explained variance of BAdam's learned perturbation to the 25th layer's up-proj matrix. Right: Effective rank of Adam's and BAdam's m learned perturbations.
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 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 le-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 AW := =WK Wo. The
7

---

cumulative explained variance "cvar" and the effective rank of matrix A E R are defined as:
Ei (A)²
effectiie_rank((A) = min :var(r) 2 0.9), cvar(r) = Σj=1 Lmi oy(4)²²
where (A) is the i-th largest singular value of A.
In the middle figure of Figure 2, display evar of BAdam's update for the 25th layer's up-proj matrix. This result shows that BAdam's '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, 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 's learning and optimization capability.
3.3 Downstream Performance Evaluation
In this subsection, conduct supervised finetuning for the Llama 2-7B [54), Llama 3-8B, and Llama 3-70B models on the Alpaca-GPPT and Mathlnstruct [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-tune models obtained by different methods for 3 epochs. We utilize two initial learning rates, le-5 and le-6, with a cosine scheduler for all methods. The results are displayed in Table 5.
Model: Llama 2-7B (base model MTT-bench: 3.93)
Ir
le-5
le-6
Method
Adam
LOMO
LoRA
Galore
BAdam
Adam
LOMO
LoRA
Galore
BAdam
Epoch! I
4.41
4.01
4.77
4.70
4.79
4.62
3.99
4.59
4.12
4.71
Epoch 2
4.73
4.06
4.84
4.83
5.21
4.94
4.02
4.86
4.17
4.83
Epoch 3
5.16
4.11
4.01
4.88
4.87
5.13
4.06
4.81
4.26
4.88
Average
4.76
4.06
4.54
4.80
4.96
4.90
4.02
4.75
4.18
4.81
Model: Llama 3-8B (base model MT-bench: 5.46)
Ir
le-5
le-6
Method
Adam (a)
LOMO
LoRA
Galore
BAdam
Adam
LOMO
LoRA
Galore
BAdam
Epoch! 1
5.49
6.17
5.78
6.07
6.15
5.40
6.41
5.66
6.65
Epoch 2
5.62
6.36
5,80
6.19
6.26
5.85
6.19
5.77
6.64
Epoch 3
5.41
6.28
5,89
6.64
6.29
5.83
6.20
5.70
6.67
Average
5.51
6.27
5.82
6.30
6.23
5.69
6.27
5.71
6.65
Table 5: MT-bench scores of the instruction-tuned Llama 2-7B and Llama 3-8B on Alpaca-GPT4 dataset by different methods. (i) Adam with learning le-S for finetuning Llama 3-8B overfits the Alpaca-GPT dataset and achieves MT-bench scores that are even lower than that of the base model. Hence, we omit this outlier.
Some remarks and observations Table 5 are in order. 1) Using le-5 as the initial learning rate, the average MT-bench score over 3 epochs achieved by BAdam surpasses that of LoRA by magnitude of 0.42 for instruction-tuning the Llama 2-7B model. Regarding instruction- tuning the Llama 3-8B model using an initial learning rate of le-6, the average score returned by BAdam outperforms that of LoRA by magnitude of 0.38. 2) In most cases, BAdam can beat LoRA and Galore, albeit sometimes slightly, across the two learning rate settings and when evaluating checkpoints from different epochs for both the Llama 2-7B and Llama 3-8B models. This underscores the promising performance of our proposed method. 3) BAdam is on par with the performance of Adam for the Llama 2-7B model and outperforms Adam for the Llama 3-8B model, partly illustrating the power of the BCD optimization scheme in LLM finetuning. It is worth noting that BAdam is both memory and running time efficient. In terms of memory usage, requires only single RTX3090-24GB GPU for finetuning the Llama 3-8B model, while Adam needs multiple A100-80GB GPUs.

---

Model: Ilama 3-8B
Method
GSM8K
Aqua
MMLU-Math
SAT-Math
MATH
NumGLUE
Average
Base model
25.9
22.8
33.7
39.5
12.8
34.5
28.2
Adam
54.5
40.5
44.3
51.4
18.4
55.4
44.1
LOMO
32.1
28.0
40.0
39.5
13.1
37.1
31.6
LoRA
47.5
44.9
45.3
50.9
14.5
56.9
43.3
Galore
33.1
37.4
41.2
42.7
15.0
36.9
34.4
BAdam
42.5
50.5
56.8
15.7
53.0
44.4
Model: Llama 3-70B
Method
GSM8K
Aqua
MMLU-Math
SAT-Math
MATH
NumGLUE
Average
Base model
52.4
46.5
52.2
58.2
21.2
37.9
44.7
LoRA 73.3
59.5
58.3
64.1
34.2
64.8
59.0
BAdam 78.8
63.4
64.2
76.4
26.2
67.3
62.7
Table 6: Zero-shot math benchmark scores of the finetuned Llama 3-8B and Llama 3-70B on MathInstruct dataset by different methods.
Math benchmarks. We also finetune the Llama 3-8B and Llama 3-70B models on the MathInstruct dataset for 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 signifi- cantly less memory consumption compared to Adam. In particular, for the experiments on finetuning Ilama 3-8B, BAdam outperforms LoRA in out of tasks, and surpasses LOMO and Galore in all the tasks by large margin. For finetuning Llama 3-70B, BAdam beats LoRA in 5 out of 6 tasks.
3.4 Allation Study: SGD's Update Rule
1.3
Full
Adam
SGD
1.0
MT-bench
6.29
5.82
BCD
BAdam
BSGD
ACETI
520
MTThench
6.67
6.50
8500
2000 tertions
Figure 3: Ablation study for BCD variants and their full counterparts for finetuning Llama 3-8B on Alpaca-GPT4 dataset. Left and middle: Convergence behavior. Right: MT-bench scores.
In this subsection, conduct ablation study to consider SGD's update rule in 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. 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 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.

---

3.5 Additional Experiment Results
We provide more experiment results in Appendix Here are the summarized results: 1) In Ap- pendix C.I, we conduct an ablation study on the ordering scheme of the partition т in BAdam, considering random reshuffling, ascending, and descending orders. In terms 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, 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 the selection of K. 2) In Appendix C.2, we examine BAdam's ' capability in classification tasks by training RoBERTa-large on SuperGLUE benchmarks. The results show that BAdam achieve similar average scores Adam. 3) In Appendix C.3, conduct a preliminary continue pretraining (CPT) experiment. We apply BAdam to train the Llama 3.1-8B-Insii on the StarCoder-Pythor dataset [28] for about epoch. The result shows that BAdam can effectively decrease the CPT loss, making it a strong candidate for CPT tasks when GPU memory is limited. 4) We display the memory consumption and running time costs for finetuning the Llama 2-7B model in Appendix C.4, which match the results for finetuning the Llama 3-8B model presented in Section 3.1.
model
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 tifte-of-the-as memory efficient methods such as LoRA.
4 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 purameter-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.
5 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 's update rule. We finetune the Llama 3-8B and Llama 3-70B models on the A Alpaca-GPP 4PP and MathInstruct datasets by BAdam with a single RTX and 4x 1A100080GBB 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 's superior performance in instruction finetuning and math finetuning, in comparison to LOMO, LoRA, and Galore. Additionally, BAdam has 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.
GPU
Limitations. Our focus has been on applying BAdam for supervised finetuning. Extending its application to preference optimization represents another opportunity to demonstrate BAdam's 's capabilities. Moreover, our CPT experiment using BAdam is only preliminary. Exploring extensively BAdam's '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 a technical algorithmic contribution that does not yield explicit negative socictal impacts. However, it carries risk of misuse.
10

---

Acknowledgments and Disclosure of Funding
The authors thank the reviewers for their insightful comments, which have helped greatly to improve the quality and presentation of the manuscript.
Xiao Li is supported in part by the National Natural Science Foundation of China (NSFC) un- der grant 12201534, and in part by the Shenzhen Science and Technology Program under grants RCYX20021008803003000 and RCYX30210609100222903
References
[1] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. GPT-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
[2] Eugene Belilovsky, Michael Eickenberg, and Edouard Oyallon. Greedy layerwise learning can scale to imagenet. In International Conference on Machine Learning, pages 583-593, 2019.
[3] Yoshua Bengio, Pascal Lamblin, Dan Popovici, and Hugo Larochelle. Greedy layer-wise training of deep networks. Advances in Neural Information Processing Systems, 19, 2006.
[4] Dimitri P. Bertsekas and John N. Tsitsiklis. Parallel and distributed computation. Prentice-Hou Hall, 1989.
[5] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in Neural Information Processing Systems, 33:1877-1901, 2020.
[6] Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, al. Sparks of artificial general intelligence: Early experiments with GPT-4. arXiv preprint arXiv:2303. /27/2, 2023.
[7] Xufeng Cai. Chaobing Song, Stephen Wright, and Jelena Diakonikolas. Cyclic block coordinate descent with variance reduction for composite nonconvex optimization. International Conference Machine Learning, pages 3469 1469-3494. PMLR, 2023.
[8] Darshan Chakrabarti, Jelena Diakonikolas, and Christian Kroer. Block-coordinate methods and restarting for solving extensive-form games. Advances in Neural Information Processing Systems, 36, 2023.
[9] Chih-Chung Chang and Chih-Jen Lin. LIBSVM: a library for support vector machines. ACM transactions Intelligent Systems and Technology (TIST), 2(3):127,, 2011.
[10] Caihua Chen, Bingsheng He, Yinyu Ye, and Xiaoming Yuan. The direct extension of ADMM for multi-block convex minimization problems is not necessarily convergent. Mathematical Programminng 155(1):57-79, 2016.
[11] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.061 2016.
[12] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reichirr Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110. 14168. 2021.
[13] Alexandre Défossez, Leon Bottou, Francis Bach, and Nicolas Usunier. A simple convergence proof of adam and adagrad. Transactions on Machine Learning Research, 2022.
[14] Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. QLoRA: Efficient finetuning of quantized LLMs. Advances in Neural Information Processing Systems, 36, 2023,
[15] Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, al. The llama 3 herd of models. arXiv preprint ar arX.v22407.21788 2024.

---

[16] Mert Gurbuzbalaban, Asuman Ozdaglar, Pablo A Parrilo, and Nuri Vanli. When eyclic coor- dinate descent outperforms randomized coordinate descent. Advances in Neural Information Processing Systems, 30, 2017.
[17] Junxian He, Chunting Zhou, Xuezhe Ma. Taylor Berg-Kirkpatie and Graham Neubig. To- wards a unified view of parameter-efficient transfer learning. International Conference on Learning Representations, 2021.
[18] Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint w:Xi:2209.033:00 2020.
[19] Dan Hendryeks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. 2021.
arXiv preprint
[20] Geoffrey E Hinton, Simon Osindero, and Why- Teh. A fast learning algorithm for deep belief nets. Neural Computation, 18(7):12771554, 2006.
[21] Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-effiiien transfer learning for NLP. In International Conference on Machine Learning, pages 2790-2799. PMLR, 2019.
[22] Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. arXiv preprint wwx.iv:210..0990
2021.
[23] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.69800 2014.
[24] Soroush Abbasi Koohpayegani, KI. Navaneet, Parsa Nooralinejad, Soheil Kolouri, and Hamed Pirsiavash. NOLA: Networks as linear combination of low rank random basis. In The Twelfth International Conference on Learning Representations, 2024.
[25] Dawid Jan Kopiczko, Tijmen Blankevoort, and Yuki M Asano. VeRA: Vector-based random matrix adaptation. In The Twelfth International Conference on Learning Representations, 2024.
[26] Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 3045-3059 2021.
[27] Haochuan Li, Alexander Rakhlin, and Ali Jadbabaie. Convergence of Adam under relaxed assumptions. Advances in Neural Information Processing Systems, 36, 2023.
[28] Raymond Li, Loubna Ben Allal, Yangtian Zi, Niklas Muennighoff, Denis Kocetkov, Chenghao Mou, Mare Marone, Christopher Akiki, Jia Li, Jenny Chim, et al. Starcoder: may the source be with you! arXiv preprint 2rXiv:2305.06111, 2023.
[29] Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. the 59th Annual Meeting of the Association for Computational Linguistics and the I 11th International Joint Conference on Natural Language Processing, pages 4582 4597, 2021.
In Proceedings the
[30] Vladislav Lialin, Sherin Muckatira, Namrata Shivagunde, and Anna Rumshisky. ReLoRA: High-rank training through low-rank updates. In The Twelfth International Conference on Learning Representations, 2024.
[31] Wang Ling, Dani Yogatama, Chris Dyer, and Phil Blunsom. Program induction by rationale generation: Learning to solve and explain algebraic word problems. In Regina Barzilay and Min- Yen Kan, editors, Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Pupers), pages 158-167. Association for Computational Linguistics, July 2017.
12

---

[32] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer. and Veselin Stoyanov. RoBERTa: A robustly optimized BERT pretraining approach. arXiv preprint 190iv 1692, 2019.
[33] Yongkang Liu, Yiqun Zhang, Qian Li, Shi Feng, Daling Wang, Yifei Zhang, and Hinrich Schütze. HiFT: A hierarchical full parameter fine-tuning strategy. arXiv preprint arXiv:2401. 15207, 2024.
[34] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 171 2017.
[35] Zhaosong Lu and Lin Xiao. On the complexity analysis of randomized block-coor descent methods. Mathematical Programming. 152:615-642, 2015.
[36] Zhi-Quan Luo and Paul Tseng. On the convergence of the coordinate descent method for convex differentiable minimization. Journal of Optimization Theory and Applications, 72(1):7-33, 1992.
[37] Kai Lv. Yuqing Yang, Tengxiao Liu, Qinghui Gao, Qipeng Guo, and Xipeng Qiu. Full parameter fine-tuning for large language models with limited resources. arXiv preprint arXiv:2306.09778) 2023.
[38] Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Alex Damian, Jason D Lee, Danqi Chen, and Sanjeev Arora. Fine-tuning language models with just forward passes. Advances in Neural Information Processing Systems, 36, 2023.
[39] Krešimir Mihić, Mingxi Zhu, and Yinyu Ye. Managing randomization in the multi-block alter- nating direction method of multipliers for quadratic optimization. Mathematical Programming Computation, 13(2):339 413, 2021.
[40] Swaroop Mishra, Arindam Mitra, Neeraj Varshney, Bhavdeep Sachdeva, Peter Clark, Chitta Baral, and Ashwin Kalyan. NumGLUE: suite of fundamental yet challenging mathematical reasoning tasks. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pages 3505-3523, Dublin, Ireland, May 2022. Association for Computational Linguistics.
[41] Yu Nesterov. Efficiency of coordinate descent methods on huge-scale optimization problems. SIAM Journal on Optimization, 22(2):341-36 2012.
[42] Julie Nutini, Issam Laradji, and Mark Schmidt. Let's make block coordinate descent converge faster: faster greedy rules, message-passing. active-set complexity, and superlinear convergence. Journal of Machine Leurning Research, 23(131):1-74, 2022.
[43] J.M. Ortega and W.C. Rheinboldt. Iterative Solution of Nonlinear Equations in Several Variables, volume 30. SIAM, 1970.
[44] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35:27730-27744, 2022.
[45] Rui Pan, Xiang Liu, Shizhe Diao, Renjie Pi, Jipeng Zhang, Chi Han, and Tong Zhang. LISA: Layerwise importance sampling for memory-efficient large language model fine-tuning. arXiv preprint arXiv:2403.17911 2024.
[46] Baolin Peng, Chunyuan Li, Pengcheng He, Michel Galley, and Jianfeng Gao. Instruction tuning with GPT-4. urXiv preprint arXiv:2304.0322 032777 2023.
[47] Yada Pruksachatkun, Phil Yeres, Haokun Liu, Jason Phang, Phu Mon Htut, Alex Wang, Ian Tenney, and Samuel R Bowman. jiant: A software toolkit for research on general-purpose text understanding models. In 58th Annual Meeting of the Association for Computational I. Linguistics, ACL 2020, pages 109-117, 2020.
13

---

[48] Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He. Zero- infinity: Breaking the gpu memory wall for extreme scale deep learning. In Proceedings of the international conference for high performance computing, networking, storage and analysis, pages 1-14, 2021.
[49] Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He. Zero-offload: Democratizing billion-scale model training. arXiv preprint arXiv:2101.06844 2021.
[50] Peter Richtárik and Martin Takáč. Iteration complexity of randomized block-coordinate descent methods for minimizing composite function. Mathematical Programming, 144(1):1-38, 2014.
[51] Ruoyu Sun, Zhi- QuanQuan Luo, and Yinyu Ye. On the expected convergence of randomly permuted ADMM. arXiv preprint arXiv:1503.000
4(6), 2015.
[52] Ruoyu Sun and Yinyu Ye. Worst-case complexity of cyclic coordinate descent: O (n" 2) gap with randomized version. Mathematical Programming, 185:487-520, 2021.
[53] Rohan Taori, Ishaan Gulrajani. Tianyi Zhang. Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford Alpaca: An instruction-followee LLaMA model. GitHub repository, 2023.
[54] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint wwXii:2307.00988, 2023.
[55] Paul Tseng. Convergence of a block coordinate descent method for nondifferentiable minimiza- tion. Journal of Optimization Theory and Applications, 109:475 4944, 2001.
[56] Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. SuperGLUE: A stickier benchmark for general-purposs language understanding systems. Advances in Neural Information Processing Systems, 32, 2019.
[57] Bohan Wang, Jingwen Fu, Huishuai Zhang, Nanning Zheng, and Wei Chen. Closing the gap between the upper bound and lower bound of adam's iteration complexity. Advances Neural Information Processing Systems, 36, 2023.
[58] Stephen J Wright. Coordinate descent algorithms. Mathematical Programming, 151(1)33-34, 2015.
[59] Wenhan Xia, Chengwei Qin, and Elad Hazan. Chain of LoRA: Efficient fine-tuning of language models via residual learning. arXiv preprint arXiv:2401 2It.:2400.04151,, 2024.
[60] Jingfeng Yang. Hongye Jin, Ruixiang Tang, Xiaotian Han, Qizhang Feng, Haoming Jiang, Shaochen Zhong, Bing Yin, and Xia Hu. Harnessing the power of LLMs in practice: A survey on ChatGPT and beyond. ACM Transactions Knowledge Discovery from Data, 18(6):1-32, 2024.
[61] Xiang Yue, Xingwei Qu, Ge Zhang, Yao Fu, Wenhao Huang, Huan Sun, Yu Su, and Wenhu Chen. Mammoth: Building math generalist models through hybrid instruction tuning. arXiv preprint N3Xi::2309000655. 2023.
[62] Biao Zhang, Zhongtao Liu, Colin Cherry, and Orhan Firat. When scaling meets LLM finetuning: The effect of data, model and finetuning method. The Twelfth International Conference on Learning Representations, 2024.
[63] Shengyu Zhang, Linfeng Dong. Xiaoya Li, Sen Zhang. Xiaofei Sun, Shuhe Wang. Jiwei Li, Runyi Hu, Tianwei Zhang. Fei Wu, al. Instruction tuning for large language models: A survey. arXiv preprint arXiv:2308. 10792, 2023.
[64] Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Yinyu Ye, Zhi-Quan Luo, and Ruoyu Sun. Adam-mini: Use fewer learning rates to gain more. arXiv preprint arXiv:2406. 16793, 2024.
14

---

[65] Yushun Zhang, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhi- faa-Quan Luo. Adam can converge without any modification on update rules. Advances in Neural Information Processing Systems,35:28855 35:28386-288333 2022.
[66] Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficieet Ilm training by gradient low-rank projection. arXiv preprint arXiv:2403.03507, 2024.
[67] Lei Zhao, Ding Chen, Daoli Zhu, and Xiao Li. Randomized coordinate subgradient method for nonsmooth optimization. arXiv preprint arXiv:2206. /4981, 2022.
[68] Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, et al. Judging LLM-as-a-judge with MT-bench and chatbot arena. Advances in Neural Information Processing Systems, 36, 2023.
[69] Yaowei Zheng, Richong Zhang, Junhao Zhang, Yanhan Ye, Zheyan Luo, and Yongqiang Ma. LlamaFactory: Unified efficient fine-tuning 100+ language models. arXiv preprint arXiv:2403.13377 2024.
[70] Wanjun Zhong, Ruixiang Cui, Yiduo Guo, Yaobo Liang. Shuai Lu, Yanlin Wang, Amin Saied, Weizhu Chen, and Nan Duan. AGIEval: human-centric benchmark for evaluating foundation models. In Findings of the Association for Computational Linguisties: NAACI 2024, pages
2299-23 2314,, Mexico City, Mexico, June 2024. Association for Computational Linguistics.
[71] Ligeng Zhu, Lanxiang Hu, Ji Lin, and Song Han. LIFT: Efficient layer-wise fine-tuning for large model models, 2024.
[72] Shengyao Zhuang, Bing Liu, Bevan Koopman, and Guido Zuccon. Open-source large language models are strong zero-shot query likelihood models for document ranking. In The 2023 Conference on Empirical Methods Natural Language Processing, 2023.
[73] Yan Zhuang, Liu, Yuting Ning, Weizhe Huang, Rui Lv, Zhenya Huang, Guanhao Zhao, Zheng Zhang, Qingyang Mao, Shijin Wang, al. Efficiently measuring the cognitive ability of LLMs: An adaptive testing perspective. arXiv preprint arXiv:2306. 2023.
15

---

Contents
1
Introduction
2 The BAdam Method
2.1
Algorithm Description and Convergence Result
2.2
Analysis of Memory Consumption and BP Time
2.2.1
Memory Consumption Analysis
2.2.2 BP Time Analysis for Consecutive Module-base Block Partition
3
Experiment Results
3.1
Memory Consumption and Wall-clock Running Time
3.2
Optimization Capability..
3.3
Downstream Performance Evaluation
3.4
Ablation Study: SGD's Update Rule
3.5
Additional Experiment Results
10
Brief Literature Review
10
5
Conclusion and Discussions Limitations
10
More Related Works
17
B Detailed Experiment Setup and Hyperparameters
18
B.1 Experiment Setup
18
B.2 Hyperparameters
19
Additional Experiment Results
21
C.I Ablation Study on Ordering Strategy and Switching Frequency
C.2 Experiments on Natural Language Understanding
21
C.3 Continue Pretrain Llama
StarCoder Dataset
22
C.4 Memory Consumption and Running Time for Finetuning I.lama 2-7B
22
D Convergence Analysis
23
16

---

A More Related Works
We 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 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 BCD-type method.
Parameter efficient finetuning (PEFT). An effective strategy for finetuning LLMs is to train 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 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, 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-min [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].
17

---

B Detailed Experiment Setup and Hyperparameters
B.1 Experiment Setup
In this subsection, we introduce the setup including the dataset, evaluation, and training details. We present the hyperparameters in Appendix B.2.
Task setup. Our experiments mainly consist of instruction tuining, math finetuning, and natural language understanding.
1. Instruction tuning. We finetune the Llama 2-7B and Llama 3-8B models on Alpaca- GPT4 dataset [46] for 3 epochs. This dataset consists of 52k instruction-followwe generated by GPT-4, using prompts from the Alpaca dataset [53]. The finetuned model is then evaluated using the MT-bench [68] with the "gpt-4" API to test its downstream performance.
data
2. Math finetuning. We finetume the Llama 3-8B and Llama 3-70B models on MathInstruct dataset [61] for 3 epochs, which contains 260K samples from 13 math related datasets. The
model is then evaluated on in-domain benchmarks [12, 19, 31, 40] and 2 out-of-domain benchmarks [70, 18] using zero-shot prompt. Our evaluation implementation is based the released code2 of [61].
SuperGI.UE. we finetune the RoBERTa-large model [32] with 355 million parameters on the SuperGI.UE, benchmark [56], and evaluate the performance of the finetuned model using the test dataset. Since the label of the original test dataset is not revealed, we randomly split the "dev" dataset into validation and test dataset; see Table 7. We focus 6 tasks the SuperGLUE benchmark, including BoolQ, COPA, WSC, RTE, MultiRC, and WiC. Since the classification modules of RoBERTa- large randomly initialized, set these classification modules to be trainable for all methods.
Task
Train
Validation
Test
BoolQ
9427
1270
2000
COPA
400
30
70
MultiRC
5100
453
500
RTE
2500
128
150
WiC
6000
238
400
WSC
554
44
60
Table 7: Data split for the SuperGLUE experiment. The original "dev" dataset is randomly splitted into validation and test datasets.
Setup for difffrrnt finetuning methods. For BAdam, we use consecutive module-based block partition represented by transformer layers, resulting in the number of blocks D) = 32 for the Llama 2-7B and Llama 3-8B models, D) 80 for the Llama 3-70B model, and D 26 for the RoBERTa-large model. The ordering strategy in the partition æ of BAdam is random reshuffling. For Galore, LoRA, and BAdam, we train all the transformer layers while freezing the language modeling head and the embedding layers. For Adam and LOMO, we set all modules in transformer layers to be trainable. We adopt the setup in Galore's paper and apply pure BF16 and 8-bits Adam for all Galore's 's experiments. We apply pure BF16 precision training for LOMO, as it does not support mixed precision training. Since LOMO does not support gradient accumulation, than the other approaches (consequently, it runs more steps) to ensure aligned memory consumption; see Appendix B.2 for more detail.
batch size is smaller
Additional implementation details.
All the experiments in Section 3.1 are conducted using single RTX3090-24GBB We use 100-0008 GPUs to finetune the Llama 3-BOB model using LoRA and BAdam.
hattps://github..cmm// AI- AI-Lab/Ie LabHMAmmTHH
18

---

For all the experiments of finetuning Llama models, we apply the gradient checkpointing technique [11] 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 finetuning Llama models are implemented using Llama- [69].
The SuperGLUE experiments are based on the implementation of jiant [47].
B.2 Hyperparameters
We summarize the choices of hyperparameters for SuperGLUE, instruction tuning, and math finetun- ing in Table 8, Table 9, and Table 10, respectively. Some discussions follow:
BAdam introduces only one additional hyperparameter compared to Adam, namely, in Algorithm 1. One natural choice is K BD where is the number of training data points, the effective batch size, and D is the number blocks in BAdam. We round BD the nearest integer if it is fractional. Such 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 is
K == ==min min
[max
72 BD 50] 100
The suggested selection is merely a guideline. One may also choose other reasonable values for K based BD once each of the 1) blocks is appropriately trained with reasonable amount of data.
For applying LoRA to finetune the Llama 2-7B model, we set LoRA rank be 100 so that it has as many trainable parameters BAdam each iteration. However, when finetuning the Llama 3-8B model, using rank 100 results in out-of-memory error, as shown Table 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 LoRA paper [22, Section 7.2] that the LoRA rank does not essentially affect its performance. Therefore, report the performance LoRA with rank 8 for finetuning the Llama 3-8B model.
Hyperparameter
Value
Ir
le-5 linear decay (Ir min 0)
Ir scheduler
warm up ratio
0.1
bz
epoch
32
weight decay
0.01
K in BAdam
100
LoRA rank
LoRA alpha
4x rank
Table Hyperparameters of SuperGLUE tasks.
19

---

Model
Hyperparameter
Value
Ir
le-5 and le-6
Ir scheduler
cosine (Ir min 0)
bz (LOMO)
8
bz (Other methods)
grad. accu. (2) x single pass bz (8) =
epoch
3
weight decay
0.01
Llama 2-7B
K BAdam
100
LoRA rank
100
LoRA alpha
4xLoRA rank
Galore rank
256
Galore subspace change freq.
256
Galore scale factor
0.25
Ir
lc-5 and 1c-6
Ir scheduler
cosine (Ir (Ir_min= 0)
bz (LOMO)
4
(Other methods)
grad. accu. (8) single pass bz (2)
epoch
3
weight decay
0.01
Llama 3-8B
K in BAdam
100
LoRA rank
8
LoRA alpha
4x LoRA rank
Galore rank
256
Galore subspace change freq.
256
Galore scale factor
0.25
Table 9: Hyperparameters of instruction tuning.
Model
Hyperparameter
Value
Ir (BAdam, Galore, LoRA)
le-5
Ir (LOMO)
le-4
Ir (Adam)
le-6
Ir scheduler
cosine (Ir min 0)
bz (LOMO)
8
bz (Other methods)
grad. accu. (2) x single pass bz (8) 16
epoch
3
Llama 3-8B
weight decay
0.01
K in BAdam
100
LoRA rank
8
LoRA alpha
4x LoRA rank
Galore rank
256
Galore subspace change freq.
256
Galore scale factor
0.25
Ir (BAdam, LoRA)
le-5
Ir scheduler
cosine (Ir min
bz epoch
grad. accu. (8) x single pass bz (2) 16
3
weight decay
0.01
Llama 3-70B
K in BAdam
100
LoRA rank
8
LoRA alpha
4 AxLoRA rank
Galore rank
256
Galore subspace change freq.
256
Galore scale factor
0.25
Table 10: Hyperparameters math finetuning.
20

---

Additional Experiment Results
C.1 Ablation Study on Ordering Strategy and Switching Frequency
Training
Training Loss
BAdam_descendne deeccenddnng
SAdam_K200
BAdam_ast Sasdam_sscending
1.2
BAdam_Ke K100
BAdam_ran manndom
3Adam_KS0
1.3
BAdam_KD
1.2
1.0
0.9
2000
4000
5400
8000
10000
2000
6000
10000
Data pass of iterations)
Dala pass (# of iterations)
(a) Loss of different ordering strategies.
(b) Loss with different Adam steps K.
Figure 4: Effect of ordering strategies and Adam steps K.
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 finetuning Llama 3-8B Alpaca- dataset.
We compare three block ordering strategies including descending (from the output to the input module), ascending (from the input 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 refer to (7) for detailed discussion selecting K.
C.2 Experiments on Natural Language Understanding
Method
BoolQ
COPA
WSC
RTE
MultiRC
WiC
Adam 0.86
0.59
0.68
0.87
0.76
0.70
LoRA
0.81
0.56
0.62
0.79
0.69
0.59
BAdam 0.85
0.69
0.65
0.76
0.77
0.64
Table 11:
SuperGLUE benchmark scores of the finetuned RoBERTa- large using different optimization methods.
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 put Appendix B.2. We display the test results on tasks selected from the SuperGI.UR benchmark. We choose these tasks to conduct experiments since they are selected [37, 38). The results can be found in Table
can 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.
21

---

C.3 Continue Pretrain Llama 3.1-SB-Instruct on StarCoder Dataset
We conduct a preliminary continue pretraining experiment on StarCoder-Pythnn [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 le-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 epoch).
Continue Pretraining
0.94
0.93
0.92
0.91
0.90
0.0
0.2
0.4
0.6
1.0
Token seen
Figure Loss of continue pretrain Llama -B--nnuttttt on StarCoder-PPthhon dataset using BAdam.
C.4 Memory Consumption and Running Time for Finetuning Llama 2-7B
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 Section 2.2.
We also display the time costs for finetuning the Llama 2-7B model Table 13 and Table 14. One 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-leve overhead. Notably, finetuning the Llama 2-7B model requires a much longer training time compared that of the Llama 3-8B model, as the latter utilizes grouped query attention that significantly improves both inference and backward efficiency.
Method
Parameter
Gradient
Optim. states
Memory consump.
Adam LOMO
13.4GB
26.8GB
80.4GB
120.6GB+
13.4GB
0.4GB
18.8GB
LoRA-ranke 100
14.0GB
1.0GB
3.0GB
22B1GB
LoRA-rank8
13.5GB
0.1GB
0.2GB
20.1GB
BAdam
13,4GB
0,8GB
2.4GB
21.8GB
Table 12: Actual memory costs of applying mixed precision training to finetune Ilama 2-7B with gradient checkpointing using single RTX3090. Note that LOMO only supports FP16 precision training. The maximum input sequence length is 728 and the batch size is 8.
Method
Forward
Backward
Update
Backward scheme
Backward time
LOMO
1.35 hours
9.71 hours
All modules
5.180 seconds
LoRA
2.48 hours
9.45 hours
56 seconds
Input module only
3.903 seconds
BAdam
1.16 hours
5.54 hours
39 seconds
Output module only
0.053 seconds
Table 13: Time spent per epoch on forward, back- ward, and update for finetuning Llama 2-7B using a single RTX3090. The single pass batch size is 8. The results averaged over epochs.
Table 14:
Time spent on different backward schemes with batch size 8 for finetuning Llama 2-7B using single RTX3090. The re- sults are averaged over 100 backward passes.
22

---

D Convergence Analysis
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 descent method and has a complexity bound of (68 (()-). For 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 C is L-smooth. And when restricted on i-th block, is smooth.. Mathematically.
ac
ac 00;
VVC(0)) - vC((3)|I & L110 - 0
- L1 - i = I....,.. D.
00, lo
We define parameters L maxi=1..... Li as the maximum smoothness constants across all blocks and & min I4-
Assumption D.2. BAdam has bounded partial derivatives along its trajectory, i.e., there exists G >>> 1 such that
119 f.k G.
Here, we adopt the notations as specified in Algorithm 1: = O..... are block epochs, i = 1. ...... are different blocks, & = 0, K are inner iterations over data for certain block, and g. V.)( a is the partial derivative w.r.t. the i-th block t-th block epoch and k-th inner Adam step. To avoid potential confusion, in this section we denote À as the numerical stability constant used in the denominator in Adam's adaptive step sizes (Line 12 of Algorithm 1) instead of s, as the is conventional notation used to represent the target accuracy in optimization community.
Corollary D.3 (bounded adaptive step sizes). Let Hit diag (1/ (V + denote the
diagonal matrix formed by coordinate-we adaptive step sizes vecton Under Assumption D.2 and with X I, we have:
2G I x H
Proof. By definition of é, its elements (6/ squared elements (g) )3. so we have bound
exponential moving average of the history gradients'
t. (9 1 Isj
) z 11 12
a².
max ISAS
max
Therefore
1
2G wsjsdd
G
then H...k 1 I is positive semidefinite matrix. Similarly, since 1/(V 20 have H.S.h L,
clik (o + A) < 1/2, we
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.land Assumption D.2 and suppose that
0 < X < 1. If the learning rate satisfies the inequality 2. min { K: X }. then BAdam with deterministic gradients has the following descent property after one block-epoch of updates:
ak
L(0 ( 16G (1+
V2(00))1
4T2 202K12D X
+
23

---

Corollary D.5 (first-order complexity). Under the conditions in Theorem D.4 and setting the learning rate as a 4LK min KD174 6G BAdam find a 6-approximate stationary point with at most
6G T 128 - -²) max DV/AK, A
8
gradient evaluations.
Proof. With the above choice of learning rate and Theorem D.4, the descending property of one block epoch can written as
ak
C(0 L(0') s
V2(00)
32G
Sum over obtain the bound for minimum gradient norm
min||VC(0\IP 1
Σ || 32G(L(0u) - L*)
aK
Substitute the choice of and get the complexity result.
To simplify notation, in the following we abuse notation L(0.) to represent 0.
L. (of+) ,...,, 0 1.41....,0 .b). And of represents (0/+1,... "0[+,,...,0+111. 1...,,0 0
0 0%)
Lemma D.6 (approximate descent inequality for one block). Under the conditions in Theorem D.4, we have the following approximate descent property for the inner solver of Adam:
aK {1
2L,aKG X²
((0)) L(
|V,C00)112
(aG L(OK letl.
2G
1.l - -1- " as the difference between the updates of Adam and full GD scaled by coordinate wise adaptive step sizes.
where we denote |e&]|l = 11 Σk= T-g. (m,
Proof. With Assumption D.I, we have the following descent inequality:
((0)) - ((0)_1) S (V.L(0), -0,) -. 1100 0/112
Li
()
K
K
s <aa(v.c()) Ket + Σ
H.* g, + L,a²
K2 +
Σ
H
t t.l
9i"
k=0
k=1
aK
aK |VL(0,) +
v.C(O)))2
2G
4G
L,a
121
V.L(O))) + lkc
X
aK
2L,
|V, + (aKG L,a3 |et
2G
X2
where the last inequality is because Corollary D.3 and Sauchy-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/[]:
2L, X
|
19
Proof. By definition in Algorithm 1 we have the expression for inner updates:
m 1.
Bmm, 1 1 + (1 -1- -B11g,,, 1.k (1- (1-33) -B1) Σg 3/
$9 1.5
(8)
24

---

So the error term can be written as:
1 K
||f|||
H1 1. (m t.k (1 - 9² 1
=
K k=
34
1 K
1
1l f.lk (1 (1-B1)st 39
K
A(1 - //4
1 K
k
$
Σ
[[(1 - 31)
Σ
8 [0. - 9)
K
j=1
K
1- 3
Σ
11
x i
9i
KX j=2
kmj
4
1105
0% ||
KX
where the first incquality is because Corollary D.3 and the second incquality is by definition of m;" in (8).
Denote A EØ=1 10000 0. I|. now let's bound A/.
K-1
A
Σ
110. - 87:01
a Σ H l m,
1.k
Σ
(1 31
2 Ji 1.
X(1 - B4)
k=
1=1
K-1
E
1
(1 - 31)
8 11
-1:1)) + (1 - 19
X(1 84)
1=1
a
(1- 31)
B/ 11 - -0% + (1 - -35011a11
J=1 A=1 X(1
1=1
K-1 a
+ (1 110-)))100=11 ||) Σ 1
(1-
3
1=1
x
In the above, the first inequality is because Corollary D.3 and the second inequality is by Assump- tion D.I. Therefore ak² Further we have
x 110 4.1
2 (1 ol,k²) ) and since S 2L,K= we getd S 2 &2 l|.
2L,KK X2
| 1
1192 1.1 11.
Corollary D.8 (refined descent inequality for one block). The approximate descent inequality for one block in Lemma D.6 can be refined as
aK
[(09) [(0)
[V,((,)]1
8G
25

---

Proof. Substitute Lemma D.7 into Lemma D.6 and we have
aK [(0)_))
1
2L,
8LIa2K26
8L
|IV.C(99112
[(6))
2G
2
X2
X
x
aK [V,((0))
8G
where the last inequality is because S 24LKG
Now are ready to prove Theorem D.4.
Proof of Theorem D.4. Sum up (9) over i, we have
- s aK
V (1,(11
(10)
8G
In the following, we use the right hand side Σ | ((0))) ) to upper bound the whole gradient vector of one block epoch [(*)(
|V,(00)] s = - V.C(0))12
- <2]V VC(0))]
(11)
s 2 [2( -
where the second inequality is because the partial derivative vector v,L(6,) part of the gradient vector (0)) and the last inequality is by L-smoothness of L.
For the second term we have
100 - 101 12 =
Σ
100
0912
j=1
K Ke +
Σ 2
H k 1.1
Dj
j=1
k=
2a 2 K2
2a 2 K22 ||e; 121lee 12 + Σ
X2
j=1
j=1
20 K
AILa2
1) 2(0,)11
X2
2
j=1
Substitute into (11) and sum over, we have
[V Σ [V_(())) $2 1+
D
20
KIPD
(AT2a2
K3
+1
:1) D
Σ
V
X²
1
Substitute back into (10) and we get our descent inequality for a whole block epoch.
26

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The major contributions and scope are properly reflected in abstract and the "Main results" " paragraph in Section 1.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We have summarized the main limitations of this work in Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
27

---

Justification:
In Appendix D, we list all the assumptions used and provide the full proof.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We include detailed experiment setup and hyperparameters in Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
28

---

Answer: [Yes]
Justification: Our submission includes the detailed setup and command for running the experiments. We also opensource our code publicly.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.c cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: //n ctsipp.ccppubbc//uuiicc httpg/dddd///oeuummitiee
onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: In Appendix B, we clearly specify the dataset split, hyperparameters, and detailed experiment setup.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: For the downstream performance evaluation in Section 3.3, we report the MT-bench scores across different learning rate schemes and different models. For the experiments on time cost Section 3.1, we report the averaged running time over three epochs.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
29

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma enror bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We clearly report the computer resources of all experiments in the experiment section; see Section 3. We also report the memory consumption and running time in Section 3.1.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: / http:///wrrip..cppctttt/////////
Answer: [Yes]
Justification:
The paper follows the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We have summarized the broader impacts of our method in Section 5.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
30

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess 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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that 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 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 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 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe 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)?
Question:
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 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 best faith effort.
12. Licenses for existing assets
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?
Question:
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 use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include 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 use in package should be provided. For popular datasets, paperswi thcode. paperswithcode. com/ datasets
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
31

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. 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 code implementation publicly and provide detailed docu- mentation 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 anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. 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 details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve any crowdsourcing experiments 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 contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS 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.
15. 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 research with human subjects.
Depending 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.
32

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere 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.
33


Reviews:
[{"review_id": "vOTaPKVWTh", "reviewer": "Reviewer_FWuu", "summary": "The authors proposed fine-tuning of LLMs with a block coordinate descent based Adam optimizer.  \nThey presented results on convergence analysis, memory and run time profiling, and the quality of resulting fine-tuned models.", "questions": "* See above weaknesses on empirical results.\n* 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": "Yes.", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "+ The proposed idea is clearly stated.\n+ The memory usage analysis is comprehensive. \n+ 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.   \n- 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.   \n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for the constructive and helpful comments. We address the reviewer's concerns in a point-by-point manner below. *All additional experiment results (figures and tables) are put in the one page supplementary PDF of the global rebuttal.*\n\n**A. Ablation study on BCD, Adam, and SGD.** As requested by the reviewer, we have conducted experiments on Adam, SGD (LOMO), BAdam, and BSGD (BCD with SGD as inner solver) for finetuning Llama 3-8B. All optimization methods use grid-searched learning rates based on training and validation losses.\n\n*Convergence ability.* In Figure (b), it can be observed that BCD variants converge slightly slower but soon exhibit similar convergence behavior in terms of running time compared to their full counterparts. It is worth mentioning that, unlike the full counterparts, BCD variants only update one block of parameters per data batch, demonstrating the strong optimization ability of BCD for LLM finetuning.\n\n*Downstream performance.* In Table 1, we also test the MT-bench scores of the finetuned models. It is quite interesting to see that BSGD significantly outperforms SGD (almost as good as BAdam), even though they have almost the same optimization convergence behavior. The superiority of BCD variants over their full counterparts possibly stems from the fact that BCD uses each data batch to update only one block of parameters, thereby better preserving the general knowledge of the pretrained model during finetuning. These strong downstream performances of BCD further illustrate its suitability for LLM finetuning. \n\nThese ablation results partly confirm the suitability of BCD for LLM finetuning and reveal that choosing either Adam or SGD as the inner solver has similar performance in terms of our downstream performance evaluation (i.e., MT-bench score for model trained on Alpaca-GPT4 dataset). We will add BSGD to our BCD framework and leave the comprehensive study of different inner solvers as future work. \n\n**B. The suitability of BCD in LLM finetuning.** Let us elaborate on the suitability of BCD for LLM finetuning from the following three perspectives.\n\n*Performance.* Based on the results of the above ablation study, we believe that we have demonstrated the strong performance achieved by BCD compared to its full counterpart. This directly indicates that BCD is suitable for LLM finetuning in terms of both optimization ability and downstream performance.\n\n*Low memory consumption.* BCD not only achieves strong performance but is also memory efficient; it only stores the gradient and optimizer states of the active block, which is substantially lower than Adam's $18M$ cost. In particular, BSGD requires little additional memory beyond just storing the $2M$ LLM, while still exhibiting strong downstream performance. Given that memory is usually the bottleneck for finetuning LLM, the memory efficiency of BCD further demonstrates its suitability for this task.\n\n*Learnability interpreted by high rankness.* Unlike LoRA and Galore, BCD's memory efficiency is achieved without restricting its updates to a low rank space. In Figure (d), we display the cumulative explained variance of BAdam's update for the up-proj matrix, which is defined as sum of the $k$ largest squared singular values over the sum of all squared singular values. This result shows that BAdam's update has a heavy tailed singular values distribution and is far away from a low rank update. In Figure (e), we show that BAdam achieves almost the same high rank update as Adam through all modules of different transformer layers, partially ensuring BCD's learnability.\n\n**C. BCD + LoRA.** We first note that the high rank update of BAdam shown in Figures (d) and (e) partly illustrates the scientific difference between BCD and LoRA. Following the reviewer's suggestion, we also combined BCD with LoRA (B-LoRA) for LLM finetuning. Due to the advantages of BCD, B-LoRA can have a higher rank configuration under the same amount of memory budget and save computation time, compared to LoRA.  B-LoRA has a slightly worse MT-bench score than LoRA (6.12 vs. 6.28) as shown in Table 1, while achieves a better 5-shot MMLU score than LoRA as shown in Table 3.\n\n**D. Comparison with Galore.** In Figures (a)-(b) and Table 1, we report the results of Galore with rank 1024 and 8-bit Adam optimizer for finetuning Llama 3-8B. We set other hyperparameters according to their paper's suggestions. It can be seen that BAdam and BSGD outperform Galore in terms of MT-bench downstream performance, demonstrating the ability and suitability of BCD for LLM finetuning. In terms of convergence, BAdam also exhibits better convergence ability than Galore.\n\nWe have not conducted experiments on combining BCD and Galore due to time constraints, but we will add such experiments in our next manuscript version, as suggested by the reviewer.\n\n**E. Significance of advantage over LOMO.** We note that the MT-bench score difference between GPT-3.5-turbo and GPT-4 on the LMSYS leaderboard is 1.05, which is close to the gap of 0.96 presented in our manuscript between the Llama 3-8B models tuned by LOMO and BAdam (5.69 versus 6.65). Therefore, the advantage appears to be significant.\n\nWe hope that our response is satisfactory to the reviewer and that all concerns have been addressed appropriately. We will properly add the above experiments to our next manuscript version. If the reviewer has any additional concerns, please let us know during the reviewer-author discussion period.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "crZSBERcPo", "reviewer": "Reviewer_9PUh", "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.", "questions": "1. In Table 1, why LOMO update precision is only limited to Float16. Can't it be Float32? \n2. 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?\n3. 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.\n4. 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": "The author has addressed the limitations.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The paper's contributions and strengths are as follows:\n\n1. It proposes using the well-known optimization technique BCD for the task of fine-tuning large language models while being memory efficient. \n2. 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.\n3. It provides theoretical convergence analysis for the deterministic case.", "weaknesses": "The paper's weaknesses are summarised as follows:\n\n1. 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.\n2. The paper does not provide theoretical or practical insights into why BCD is effective for language model fine-tuning.\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for the constructive and helpful comments. We address the reviewer's concerns in a point-by-point manner below. *All additional experiment results (figures and tables) are put in the one page supplementary PDF of the global rebuttal.*\n\n**A. Ablation study & effectiveness of BCD for LLM finetuning.** We have conducted ablation experiments on Adam, SGD (LOMO), BAdam, and BSGD (BCD with SGD) for finetuning Llama 3-8B.  All optimization methods use grid-searched learning rates based on training and validation losses.\n\n*Convergence ability.* In Figure (b), it can be observed that BCD variants converge slightly slower but soon exhibit similar convergence behavior in terms of running time compared to their full counterparts. It is worth mentioning that, unlike the full counterparts, BCD variants only update a block of parameters per data batch, demonstrating the strong optimization ability of BCD for LLM finetuning. \n\n*Downstream performance.* In Table 1, we also test the MT-bench scores of the finetuned models. It is quite interesting to see that BSGD significantly outperforms SGD (almost as good as BAdam), even though they have almost the same optimization convergence behavior. The superiority of BCD variants over their full counterparts possibly arises from the fact that BCD uses each data batch to update only one block of parameters, thereby better preserving the general knowledge of the pretrained model during finetuning. These strong downstream performances of BCD further illustrate its suitability for LLM finetuning.\n\n*Learnability interpreted by high rankness.* Unlike LoRA and Galore, BCD's memory efficiency is achieved without restricting its updates to a low rank space. In Figure (d), we display the cumulative explained variance of BAdam's update for the up-proj matrix, which is defined as sum of the $k$ largest squared singular value over the sum of all squared singular value. This result shows that BAdam's update has a heavy tailed singular values distribution and is far away from a low rank update. In Figure (e), we show that BAdam achieves almost the same high rank update as Adam through all modules of different transformer layers, partially ensuring BCD's learnability. \n\nWe believe that these experiment results demonstrate the effectiveness of BCD for LLM finetuning. \n\n**B. Handle the extremely memory-limited setting with BAdam.** In the extremely memory-limited setting, we have the following two choices to apply BAdam.\n\n*Smaller block partition.* Our BCD optimization framework (as well as our code implementation) provides a flexible block partition. Hence, we can choose a smaller block partition, further reducing the memory consumption as we only need to store the gradient and optimizer states for a smaller block of parameters. We have tested such a smaller block setting of BAdam (each layer is further partitioned to attention and mlp blocks), and it achieves a MT-bench score of 6.50, which is just slightly worse than the score of 6.67 reported in the manuscript. \n\n*Cheap CPU update.* We can also store the optimizer states in CPU memory, and performs the cheap block update purely on CPU. In this case, only the block gradient needs to be communicated between CPU and GPU during update. We have tested this approach under the Llama 3-8B experiment setting and found that the operation only induces a 0.8-second overhead per update, resulting in roughly 16 minutes of additional time cost over 3 epochs of training. We will upload this part of new code implementation when uploading is allowed. \n\nWe remark that MeZO and LOMO also maintain the gradient of a certain block in memory when updating this block. Thus, our BCD scheme can be applied using one of the above schemes whenever MeZO and LOMO are applicable. \n\n**C. BAdam needs to store the full checkpoint.** Storing the full parameter checkpoint appears to be inevitable when using full parameter optimization methods such as Adam and SGD. While LoRA has the advantage of storing different type of checkpoints, full parameter finetuning might achieve higher performance, as we have shown. \n\n**D. LOMO update precision issue.** LOMO's update precision depends on the precision of the model weights. Consequently, updating the model using Float32 will double the memory consumption to $4M$. We follow the reported update precision (Float16) used in the LOMO paper.\n\n**E. LOMO's performance under grid searched step size.** We have conducted a grid search for SGD (LOMO)'s learning rate among (5e-2, 1e-2, 1e-3, 1e-4, 1e-5, 1e-6) based on loss, and identified 1e-2 as the optimal one for MT-bench evaluation. In Table 1, it can be observed that SGD (LOMO) with learning rate 1e-2 achieves an MT-bench score of 5.82, which aligns with the score of 5.83 obtained by using learning rate 1e-6 in our manuscript.\n\n**F. Batch size is not consistent across different methods.** Since LOMO performs update on the fly and does not store the full gradient, it does not support gradient accumulation. We choose the largest batch size under the memory constraint of RTX3090-24GB for LOMO. We also ensure each algorithm uses the same amount of data. Hence, we count twice updates of LOMO as one update in our manuscript, ensuring a fair comparison to LoRA and BAdam.\n\nWe hope that our response is satisfactory to the reviewer and that all concerns have been addressed appropriately. We will properly add the above experiments to our next manuscript version. If the reviewer has any additional concerns, please let us know during the reviewer-author discussion period.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5Hami8UyaZ", "reviewer": "Reviewer_XVND", "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.", "questions": "N/A", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- **Extensive Theoretical Proof of Convergence**: The paper provides substantial theoretical evidence to support the convergence of the proposed method.\n- **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.\n- **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.\n- **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.\n- **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.\n\n[1] https://openreview.net/forum?id=u0INlprg3U", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for the constructive and helpful comments. We address the reviewer's concerns in a point-by-point manner below. All additional experiment results (figures and tables) are put in the one page supplementary PDF of the global rebuttal.\n\n**A. More quantitative results on MMLU and math benchmarks.** Following the reviewer's suggestion, we have tested the MMLU scores of the instruction-tuned models; see Table 3. We can observe that BAdam performs as good as Adam and outperforms LoRA.\n\nIn Table 2, we have also conducted math-instruction tuning for Llama 3.1-8B on the MathInstruct dataset using BAdam and LoRA, and tested several math benchmarks including GSM8K, MATH, NumGLUE, SVAMP, MMLU-Math, and SAT-Math. The results demonstrate that BAdam clearly outperforms LoRA. We have not conducted the same experiments for Adam due to time constraints, but we will add such experiments in the next version. \n\nWe believe that these additional quantitative results, together with the MT-bench results, illustrate the capability of BAdam for LLM finetuning. \n\n**B. More block partition schemes.** As requested by the reviewer, we have tested the ratio-based block partition scheme, where one block is formed by selecting a small part of parameters from each matrix module. We can observe that ratio partition performs nearly as well as the layer partition in terms of optimization ability as shown in Figure (c). In addition to this ratio partition, we can also treat each attention / mlp module as one block, and our test shows that this scheme also has similar optimization ability to the layer partition. We will add ablation study on these additional block partition schemes in the next version.\n\n**C. Fundamental differences from LIFT.** Let us first mention that LIFT is properly cited in our Section 4. Though LIFT utilizes a layer-wise update of Adam, it differs from BAdam fundamentally in the following aspects. \n\n*Different principles in algorithm foundation.* LIFT has two types of iteration policies as stated in their Appendix D, including cyclic and grouped iteration policies. The cyclic policy updates a layer for one iteration and then offloads Adam's optimizer states to the CPU (see the last paragraph of their Section 3.2). This is not a valid BCD with Adam scheme because the gradient of the same layer in the next round of update does not match the offloaded Adam's optimizer states (since other layers' parameters have been updated). On the other hand, their grouped policy updates one layer for the total number of scheduled iterations and does not revisit it, and hence there is no BCD loop at all. This makes their grouped policy a special type of block-restricted implementation of Adam rather than a general BCD scheme.\n\nBy contrast, BAdam is built on the BCD optimization framework, resulting in a fundamentally different algorithm. BAdam includes an outer BCD loop and properly accumulates the first and second momentum for updating a block ($K$ in our algorithm description). Importantly, when BAdam revisits a block, *its first and second momentum start at 0* (no offloads), which is crucial for ensuring both theoretical and empirical convergence.\n\n*Theoretical convergence guarantee.* Thanks to BAdam's foundation in the BCD framework, we can provide a convergence guarantee, adding reliability and predictability to its performance.\n\n*Practical code implementation.* We have made significant efforts to develop a compatible and robust code implementation for BAdam, with special attention to memory management. For example, our submitted code allows efficient finetuning of Llama 3-8B using a single RTX3090-24GB GPU, whereas LIFT requires several A100 GPUs (according to their experiment description since they have not open-sourced their implementation). Indeed, implementing a direct block-restricted update of Adam could be trivial with sufficient resources, as one can simply set a certain layer to require gradients in the original Adam optimizer without needing complicated memory management and optimizer coding. Our practical implementation ensures broader applicability and benefits for practitioners.\n\n**D. Verbose experiment section.** We fully agree with the reviewer's concern. We have conducted additional experiments, including more quantitative results (MMLU and math benchmarks), ablation experiments comparing Adam, SGD, BAdam, and BSGD (BCD with SGD), comparison with Galore, and continual pretraining; see the one page supplementary PDF of the global rebuttal. We will incorporate these experiments into our next manuscript version. Specifically, our plan is to move a large part of Section 3.1.1, the convergence verification in Section 3.1.2, and the entire Section 3.2 to the Appendix. Instead, we will add more downstream evaluations (such as MMLU and math tests), create an ablation study section, and a continual pretraining section.\n\nWe hope that our response is satisfactory to the reviewer and that all concerns have been addressed appropriately. If the reviewer has any additional concerns, please let us know during the reviewer-author discussion period.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LIoIw3y94u", "reviewer": "Reviewer_JX1Q", "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.", "questions": "See my weakness part.", "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.", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n2. The theoretical convergence analysis lends credibility to the proposed method, even if limited to the deterministic case. The proof seems sound and complete.\n3. 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.\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for the constructive and helpful comments. We address the reviewer's concerns in a point-by-point manner below. *All additional experiment results (figures and tables) are put in the one page supplementary PDF of the global rebuttal.*\n\n**A. Convergence result under stochastic setting.** Our approach of obtaining complexity under the deterministic setting consists of two main steps: (1) establishing the descent property for the inner solver within one block and (2) aggregating the descent inequalities across all blocks for one block-epoch. Step (2) is a standard technique in the complexity analysis of BCD-type methods under reasonable conditions. To extend to stochastic setting, the main difficulty lies in extending step (1). The complete proof is beyond the scope of this rebuttal stage; we provide a proof sketch here. We adopt the same assumptions on the smoothness of the objective function and stochastic gradient errors as in [1].\n\n*Notation modification.* We distinguish the true gradient $g_{i}^{t,k}$ and the estimated stochastic gradient $\\tilde{g}\\_i^{t,k}$. The difference term $e_{i}^{t}$ between the true gradient $g_{i}^{t,k}$ and the scaled momentum $\\hat{m}_{i}^{t,k}$ now includes stochastic gradient errors:\n\n$$\ne_{i}^{t, k} = \\hat{m}_{i}^{t, k} - g\\_{i}^{t, k},\n$$\n\nwhere the momentum is  updated with stochastic gradients $m\\_{i}^{t,k} = \\beta\\_{1}  m\\_{i} ^{t,k-1}  + (1- \\beta\\_{1} ) \\tilde{g}\\_{i} ^{t,k} $.\n\nAssuming stochastic gradients with a uniform almost sure bound $\\sigma$, similar to [1], $||e_{i} ^{t,k}||$ can be bounded by the norms of the true gradients with probability at least $1- \\delta$:\n$$\\sum_{k=1}^{K}  ||e_{i} ^{t,k} ||^2 \\le \\Theta (\\lambda) \\sum_{k=1}^{K} ||g_i^{t,k}  || ^2 + \\mathcal{O}\\left(\\sigma ^2 \\log (\\delta ^{-1} ) ( 1/ \\beta_1 + \\beta_1 K) \\right).$$\n\n*Approximate descent inequality (Inner solver).* Substituting the above probabilistic bound of error terms into the following approximate descent inequality for each step of the inner solver:\n\n$$  \\mathcal{L}( \\theta _{i}^ {t,k}) - \\mathcal{L}(\\theta _{i} ^{t,k-1} ) \\le - \\Theta (\\alpha ) ||\\nabla _{i} \\mathcal{L} (\\theta _{i} ^{t,k-1} ) || ^2 + \\Theta (\\alpha) || e\\_{i} ^{t,k-1} || ^2.$$\n\nThen combining a telescoping argument with our technique of bounding gradient differences in the inner loop $||g_{i} ^{t,k} - g_{i} ^{t,1} ||, \\forall k \\le K $ by $||g_{i} ^{t,1} || $:\n\n$$   \\sum_{k=1}^{K}  || g _{i} ^{t,k} - g _{i} ^{t,1} || \\le \\mathcal{O} \\left(\\alpha K^2 \\right) || g _{i} ^{t,1} || + \\mathcal{O} \\left(\\alpha K^{3 /2} \\sigma \\sqrt{(1 / \\beta _{1} + \\beta _{1} K) \\log (\\delta ^{-1} ) } \\right). $$\n\nWe obtain the following approximate descent inequality for each block under the stochastic setting:\n\n$$\\mathcal{L} (\\theta\\_{i} ^{t} ) - \\mathcal{L}(\\theta\\_{i-1}^{t}  ) \\le - \\Theta ( \\alpha K)  || \\nabla_i \\mathcal{L}(\\theta_i^{t}) ||^2 +\\mathcal{O}\\left(\\alpha ^3 K ^3 \\sigma ^2 \\log (\\delta^{-1} ) ( 1/ \\beta_1 + \\beta_1 K) \\right).$$\n\n*Approximate descent inequality (BCD epoch).* By aggregating the above approximate descent inequality across different blocks, similar to our arguments in the submitted manuscript, we derive the approximate descent inequality for one block-epoch:\n\n$$\\mathcal{L}(\\theta ^{t+1} ) - \\mathcal{L}(\\theta ^{t} ) \\le - \\Theta (\\alpha K) ||\\nabla \\mathcal{L}(\\theta ^{t} )|| ^2 + \\mathcal{O}\\left(\\alpha ^3 K ^3 D\\sigma ^2 \\log (\\delta ^{-1} ) ( 1/ \\beta_1 + \\beta_1 K) \\right).$$\n\nThrough standard manipulations, we should obtain a complexity result for finding an $\\varepsilon$-stationary point.\n\n[1] Li, H., Rakhlin, A., \\& Jadbabaie, A. (2023). Convergence of Adam Under Relaxed Assumptions. NeurIPS.\n\n**B. Initial continual pretraining experiment using BAdam.** Pretraining from scratch is out of scope of this work; we have conducted an initial continual pretraining (CPT) experiment for Llama 3-8B-Instruct on the StarCoder dataset using BAdam. We evaluate its performance using the online CPT loss; see Figure (f). We are only able to finish about 0.1 epoch (1.5B tokens) training due to time constraints. The online CPT loss shows that the model is effectively learning new domain specific (code-related) knowledge. In the following, we list several possible reasons why BAdam may become a candidate optimizer for LLM continual pretraining.\n\n*BAdam exhibits as high rank update as that of Adam.* In the pretraining stage, model acquires massive knowledge over different domains by training on billions of tokens. Intuitively, new factual information (not appeared in the pretraining corpus) might not be encoded by low rank update. As shown in Figure (e), BAdam achieves almost the same high rank update as Adam through all modules of different transformer layers, partially ensuring BAdam’s learnability. \n\n*Possible ability of avoiding general knowledge forgetting.* Perhaps one of the most challenging tasks in CPT is to avoid forgetting of general knowledge. We suspect that BAdam might be good at avoding forgetting compared to Adam. This conjecture stems form the fact that the BAdam uses each data batch to update only one block of parameters, thereby better preserving the general knowledge of the model during CPT. This is partly confirmed by that Llama 3-8B-Instruct has a 67.7 MMLU score (0-shot), while the MMLU score of our CPTed checkpoint after training on 1.5B tokens only decreases to 66.7 (without using any model merging techniques). \n\nWe hope that our response is satisfactory to the reviewer and that all concerns have been addressed appropriately. We will properly add the above continual pretraining experiments (with additional comparisons with LoRA and Adam) and try to add convergence result under stochastic setting (if our sketch does not contain vital error) to our next manuscript version. If the reviewer has any additional concerns, please let us know during the reviewer-author discussion period.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Dear ACs and Reviewers,

This global response contains our one-page supplementary PDF of the rebuttal. All additional figures  and tables are included in this file. 

Best regards,

Authors.

Meta-review: 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.
Decision: Accept (poster)