Sparse MeZO: Less Parameters for Better Performance in Zeroth-Order LLM Fine-Tuning

TL;DR

Sparse MeZO introduces a parameter-efficient zeroth-order optimization method that selectively applies to key parameters, significantly enhancing fine-tuning performance and speed for large language models with minimal memory use.

Contribution

It proposes a novel parameter selection scheme for zeroth-order optimization, enabling effective sparse fine-tuning of large language models with reduced memory and improved convergence.

Findings

Achieves 9% accuracy improvement on RTE task

Provides 3.5x faster convergence compared to MeZO

Enables fine-tuning LLaMA-30b on a single GPU

Abstract

While fine-tuning large language models (LLMs) for specific tasks often yields impressive results, it comes at the cost of memory inefficiency due to back-propagation in gradient-based training. Memory-efficient Zeroth-order (MeZO) optimizers, recently proposed to address this issue, only require forward passes during training, making them more memory-friendly. However, compared with exact gradients, ZO-based gradients usually exhibit an estimation error, which can significantly hurt the optimization process, leading to slower convergence and suboptimal solutions. In addition, we find that the estimation error will hurt more when adding to large weights instead of small weights. Based on this observation, this paper introduces Sparse MeZO, a novel memory-efficient zeroth-order optimization approach that applies ZO only to a carefully chosen subset of parameters. We propose a simple yet effective parameter selection scheme that yields significant performance gains with Sparse-MeZO. Additionally, we develop a memory-optimized implementation for sparse masking, ensuring the algorithm requires only inference-level memory consumption, allowing Sparse-MeZO to fine-tune LLaMA-30b on a single A100 GPU. Experimental results illustrate that Sparse-MeZO consistently improves both performance and convergence speed over MeZO without any overhead. For example, it achieves a 9\% absolute accuracy improvement and 3.5x speedup over MeZO on the RTE task. Code is available at https://github.com/NUS-HPC-AI-Lab/SparseMeZO.