Memory Efficient Block Coordinate Descent Method for Forward-Only Second-Order Finetuning of LLM Models

Dear Reviewer Nphh,

We sincerely appreciate your feedback. In addressing the comments from all reviewers, we have made efforts to clarify and improve each point raised. Below, we provide a concise summary of the key changes made in response to your concerns and recommendations.

Weaknesses:

1. Technical contribution: At the time of our first submission, no prior work had combined second-order Newton methods, zeroth-order optimization, and block coordinate descent. Through extensive parameter searches and experiments, we demonstrate that activating partial layers with an ascending order as the block coordinate descent form succeeded and match MeZO's performance as a second-order integrated method.
2. Memory Efficiency: We argue that as a second-order informed method, there must be a memory-convergence tradeoff, and we managed to reduce the impact to the minimum. New experiment results on OPT-30B demonstrate that our method improves speed and accuracy compared to MeZO, showing breakthroughs in efficiency (Refer to our responses to Reviewers h68x and CqZr ).

3. Performance Gains: We can reach 1.83x wall-clock time speedup than MeZO and 2.77x speedup than HiZOO baseline, training opt-30B on SST2. As shown in our results on larger models (30B) as above, we achieve better accuracy and computational efficiency than MeZO. We will explore block selection strategies and perform thorough parameter searches in future work to support better baselines.
4. Theoretical Justification: For the BCD bound, we proposed a layer selection strategy based on adjusting bandit probabilities via block gradient norms in response to Reviewer K578. As a starting point, this may shed light on an efficient block selection method.
5. First Order Baselines: In response to Reviewer h68x, we show experiments comparing memory requirements and explain why first-order methods (even with BCD) are less efficient than zeroth-order methods. We attach the table as follows. Note that in low-end scenarios, first-order methods remain unsuitable for comparison with zeroth-order methods because they still require significant memory overhead. Even if optimized for consumer-level 24GB GPUs, their memory usage approaches hardware limits, easily leading to unstable training (OOM errors) and small usable batch sizes (1, 2, 4). See our baseline results in response to h68x, whereas our zeroth-order method could run far below memory limits (<8GB), or support much larger batchsizes (128).

Note: These results illustrate significant differences in memory requirements between first-order and zeroth-order methods. First-order approaches, even though including memory-efficient or BCD variants, still incur memory overhead from storing states such as activations. When optimizing the first layer in a first-order BCD framework, peak memory usage remains high due to the need to retain activations for gradient computation in subsequent layers.

6. Scope of Experiments: After extending the parameter scale, our SST-2 result reaches 92.8% is close to the accuracy reported by HiZOO for 66B models (93.6%). We added an additional experiment with 8 GPUs (16 batch size per GPU), achieving 93.6% accuracy with the fastest speed. This could prove the generalizability of our approach to larger-scale LLMs. We attach this result to the table "Memory Efficiency".

Thank you for your help in improving our work. We look forward to your feedback.

Dear Reviewer CqZr,

Thank you for your thoughtful questions and insights. We appreciate your feedback and will address it as follows:

Claims and Evidences:

Training and test loss for figure 3, train loss are smoothed as in the figure:

Stability analysis: To address your request, we test under extreme conditions with batch size = 1 (without extensive tuning), observing a accuracy drop: SST2 acc=0.8337 (baseline 0.917). Since using batch size = 1 introduces significant variance, in practice, we have avoided this by employing batch size $\ge$ 16, which also easily fits within memory limits for zeroth-order methods (see also our response to Reviewer h68x, Experimental Design).

Theoretical Claims: We direct you to the randomized BCD framework discussed in our response to Reviewer K578.

Missing Reference:

We appreciate your suggestion to include ReLIZO (Wang et al., NeurIPS 2024). We note that this work improves computation efficiency by modeling the gradient estimation into a QCLP problem, and applying query-reuse strategy in zeroth-order optimization, achieving impressive results across tasks (especially, 93.4% accuracy on SST-2). We will cite it as a representative example of computationally efficient zeroth-order techniques.

Weaknesses:

1. As detailed in our response to Reviewer h68x, we expanded our experiments training OPT-30B on SST2 task, using 2×A100 nodes. B-PDF achieves 92.89% accuracy in 7.5 hours, outperforming HiZOO (90.3%, 19.8 hours) and MeZO (90.6%, 13.7 hours). This result validates the scalability and efficiency of our method when scaling up.

2. We tested some additional strategies: We evaluated following block selection strategies on OPT-1.3B, SST-2 (We note that arameters are not fully searched in these experiments, so there is room for improvement):

Note: In practice, we found counting gradient norm introduces heavy computational overhead, making it less efficient than the natural ascending order or random sampling. This matches the results in BAdam [4]. As for the odd-even staged strategy, we perform [1, 3, 5, 7], followed by [2, 4, 6, 8], ... , as integrating our active block sets with your suggestion.

[4] Luo, Qi, Hengxu Yu and Xiao Li. “BAdam: A Memory Efficient Full Parameter Optimization Method for Large Language Models.” Neural Information Processing Systems (2024).

We will add this analysis to appendix and extend it in future work. Currently, we see that the ascending order remains optimal for balancing performance and simplicity.

3. Regarding the application prospects of combining low-precision formats with memory-efficient optimizers, we fully agree with your point of view. In our current work, we conduct experiments based on MeZO framework with FP16 precision, and in the future we could explore combining our method and lower precision training. In addition, we observed challenges in utilizing Hessian information matrices under low-precision conditions during our experiments, and have applied clamping methods to stabilize numerical computations. We are following advancements in recent work and advanced GPU architectures supporting FP8 quantization, which we believe that quantization methods will significantly enhance training efficiency in the future.

Questions:

1: Quantities in Table 3: The quantities represent test accuracy.

2: Due to low-end system instability on our consumer-grade hardware (RTX 4090 on an old and low-bandwidth motherboard), BAdam caused overheating and crashes.

For comparison, we have tested Adamw-HF(HuggingFace version) accuracy on SST-2, opt-1.3b: 93.70% (with BS=8 LR=5e-6, peak allocated memory=10450MB, 10k training steps).

3: We conduct an experiment under a MeZO setting, which shows that the accuracy slightly improves (91.7% to 91.86%) when using single perturbations.

(MODEL=opt-1.3b TASK=SST2 BS=16 MODE=ft LR=3e-7 EPS=1e-3 STEPS=20000, running time is 2h57min on an A6000 node)

We sincerely express our gratitude for your guidance in improving our work.

Dear Reviewer h68x,

Thanks for your feedback and valuable suggestions. Below, we address each of your concerns to improve our manuscript.

Experimental Designs and Question 1:

Baseline Comparisons with GaLore and Other Methods: We have carefully considered your comments and those from Reviewers CqZr and Nphh, regarding comparisons with first-order memory-efficient baselines such as BAdam, GaLore, and LOMO.

Below, new experiments demonstrate the high memory consumption of these methods, conducted on a single RTX A6000 (48GB) with the LLaMA-3-8B model loaded in FP32 format for the SST-2 task (table shown in 5.First Order Baselines to Reviewer Nphh) :

• GaLore-AdamW: Encountered out-of-memory (OOM) errors even with a batch size of 1. When loaded in BF16 format, it could run with a maximum batch size of 8, consuming 42,132 MiB of memory. However, the initial training steps were significantly slower due to its additional computational overhead.
• LOMO: Achieved a maximum batch size of 4 with 39,910 MiB memory usage.
• MeZO: Supported a maximum batch size of 128 with 35,636 MiB memory consumption.
• BAdam: Required 23.5 GB of memory with a batch size of 2, as reported in the paper [1]. We note that while the official implementation utilizes a paged optimizer to reduce memory pressure, we observed high I/O costs between CPU and GPU memory in our test environment (consumer-grade motherboard and RAM), leading to system instability. This limitation is less pronounced in data center environments with optimized cooling and hardware support, but it highlights challenges for low-end systems.

[1] Luo, Qi, Hengxu Yu and Xiao Li. “BAdam: A Memory Efficient Full Parameter Optimization Method for Large Language Models.” Neural Information Processing Systems (2024).

These results illustrate significant differences in memory requirements between first-order and zeroth-order methods. First-order approaches, including memory-efficient or BCD variants, still incur memory overhead from storing activations for backpropagation. For example, when optimizing the first layer in a first-order BCD framework, peak memory usage remains high due to the need to retain activations for gradient computation in subsequent layers. Thus, our original experiments focused on zeroth-order baselines. We will attach the comparison in the supplementary material to enable a more comprehensive understanding of the memory efficiency of first-order methods.

Regarding ZO-Adam, while it demonstrates higher accuracy than MeZO in certain scenarios, its memory footprint is substantially larger. As reported by Zhang et al. [2], fine-tuning the full OPT-13B model on the MultiRC dataset with a batch size of 4 requires 64 GB for ZO-SGD and 158 GB for ZO-Adam, exceeding the capacity of typical low-end devices.

[2] Zhang, Yihua, Pingzhi Li, Junyuan Hong, Jiaxiang Li, Yimeng Zhang, Wenqing Zheng, Pin-Yu Chen, Jason D. Lee, Wotao Yin, Mingyi Hong, Zhangyang Wang, Sijia Liu and Tianlong Chen. “Revisiting Zeroth-Order Optimization for Memory-Efficient LLM Fine-Tuning: A Benchmark.” ArXiv abs/2402.11592 (2024).

To ensure comprehensive analysis, we will include detailed comparisons with these methods in the supplementary material. This addition will clarify the trade-offs between memory efficiency and computational requirements across different optimization paradigms.

Weaknesses:

Scalability to Models $>$ 7B Parameters:

While our initial experiments focused on OPT-1.3B and LLaMA-2-7B due to hardware limitations, now we have conducted an additional experiment to address scalability concerns and further validate our method. For OPT-30B fine-tuned on the SST-2 dataset with a batch size of 16 (2xA100 80GB nodes),

B-PDF achieved 92.89% accuracy in 7.5 hours , significantly outperforming HiZOO (90.3% accuracy [3], 20.8 hours) and MeZO (90.6% accuracy [3], 13.7 hours).

The improved accuracy and reduced runtime highlight the efficiency of our proposed method for both low-end and high-end environments. While zeroth-order methods traditionally trade accuracy for memory savings, and second order integration brings significant computing overhead, our method mitigates this trade-off, enabling competitive performance even on larger models.

[3] Zhao, Yanjun, Sizhe Dang, Haishan Ye, Guang Dai, Yi Qian and Ivor Wai-Hung Tsang. “Second-Order Fine-Tuning without Pain for LLMs: A Hessian Informed Zeroth-Order Optimizer.” ArXiv abs/2402.15173 (2024).

Question 2:

Impact of Block Selection Strategies:

Due to length limits, please refer to our answer to Reviewer CqZr for a detailed discussion on the impact of block selection strategies. We will include this analysis in the supplementary material.

We appreciate your guidance in helping us strengthen this work.

Dear Reviewer K578,

We sincerely appreciate your thoughtful review and apologize for the confusion caused by the typos and unclear claims. We are grateful for your careful reading and will address each point thoroughly.

Theoretical Claims:

For Equation 2 and the EMA, the first term $\Sigma_{t+1}^{-1}$ represents the stored Hessian. We will add a hat to indicate the $\hat\Sigma_t$ is the estimated Hessian, as described in line 311 of Algorithm 1, and the use of the absolute value sign to ensure non-negativity. We will remove the term diag() in line 312 since it is already in diagonal form. This revision will address your concern about the unaligned EMA process. (point 5)

Regarding the confusion in Algorithm 1:

• point 1, indices: we will explain them more carefully as below.
• point 4, the index $\pi_b$ means that in all we have $b$ active blocks for the step. We find that selecting more than one blocks for efficiency works well, and this could be a hyperparameter to search for. In practice, using 4 blocks works effectively.
• The parameter $\mu$ represents the in-place perturbation, and we perform forward passes for $\theta$, $\theta+\mu\Sigma^{\frac{1}{2}}$, and $\theta-\mu\Sigma^{\frac{1}{2}}$ with $\mu_i$ in the directions 0, +$\mu$, and -2$\mu$.
• point 3, The index $s$ refers to the seed, and we use it to sample the same random direction $z$ for both perturbation and updates, which is a key idea in MeZO.
• point 2, we admit that our implementation follows HiZOO's design in computing Hessian first. In our practice, block switching is frequent, and this implementation sometimes converges faster since the Hessian initialization is an identity matrix. We will try to ablate its impact in revisions.

For the Convergence Proof for BCD, as appropriately noted, the process different from full-parameter optimization can raise theoretical concerns. To address this, we provide the analysis to a randomized BCD framework by incorporating a probabilistic block selection mechanism. Below is the proof sketch:

For brevity, consider the objective function $\mathcal{L}(\theta)$, where the parameters $\theta = [\theta_1, \dots, \theta_D]$ are partitioned into $D$ blocks. In each iteration, a block $i \in \\{1,\dots,D\\}$ is randomly selected with probability $p_i$ to be updated. At each iteration, the gradient block $\hat{\nabla} \mathcal{L}\_{t,i}$ is retained with probability $p_{t,i}$ and dropped otherwise. This sparsification is formalized as: $\hat{\nabla} \mathcal{L}(\theta_t) = \sum_{i=1}^D \frac{\hat{\nabla} \mathcal{L}\_{t,i}}{p_{t,i}} Z_{t,i},$ where $Z_{t,i} \sim \text{Bernoulli}(p_{t,i})$. The sparsified gradient $\hat{\nabla} \mathcal{L}(\theta_t)$ is unbiased. Under following assumptions:

• block L-smoothness $||\nabla_i \mathcal{L}(\theta) - \nabla_i \mathcal{L}(\theta')|| \leq L_i ||\theta_i - \theta_i'||$,
• bounded gradient estimation variance $\mathbb{E}[|| \hat{\nabla}\_i \mathcal{L}(\theta) - {\nabla}\_i \mathcal{L}(\theta)||^2]\le \sigma_i^2$,
• Hessian preconditioning constraints $0 < \beta_\ell \leq \lambda_{\min}(\Sigma_i) \leq \lambda_{\max}(\Sigma_i) \leq \beta_u$,
• Bregman divergence bound: $B(\theta, \theta') \leq R^2$ for all $\theta, \theta'$. Define $\Sigma = \text{diag}(p_1 \Sigma_1, \dots, p_D \Sigma_D),$

we can finally prove that,

$\mathbb{E} \sum_{t=1}^T \sum_{i=1}^D\left[||\nabla \mathcal{L}(\theta_d^t)||^2_{\Sigma}\right] \le \frac{2 R^2}{\eta}+16\eta L_\infty (\text{tr}(\Sigma_\infty)+\beta_u) \sum_{t=1}^T \sum_{i=1}^D \left( \frac{||\nabla_i \mathcal{L}(\theta^t)||^2_{\Sigma_i}}{p_{i,t}}+\sigma_i^2 \right)$

Fortunately, to the best of our knowledge, existing approach can solve $p_{t,i} = \frac{||\hat{\nabla} \mathcal{L}\_{t,i}||}{\sum_{d=1}^D ||\hat{\nabla} \mathcal{L}\_{t,d}||}$ with a bandit trick, and it does not affect the $\mathcal{O}(1/\sqrt{T})$ bound[1].

[1] Communication-efficient Distributed Learning for Large Batch Optimization. Rui Liu, Barzan Mozafari, PMLR 162:13925-13946, 2022.

Thus, we extend the convergence analysis to a block coordinate descent setting, demonstrating the applicability of our method. This also matches with that, while BCD may prevent convergence in non-convex settings, our experiments demonstrate stable convergence in practice. Due to response length limits, we can discuss this further in the final response period and we welcome your suggestions.

Broader Scientific Literature:

Regarding performance: Initial experiments were limited by hardware. However, with additional computational resources, we have benchmarked on OPT-30B to validate scalability, which achieves better accuracy with less training time for larger models. Please, refer to our response to Reviewers h68x and Nphh. We leave hyperparameter tuning and BCD schemes for future work.

Thank you for your constructive feedback on strengthening both our theory and practical validation.