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.

翻译：尽管针对特定任务微调大语言模型（LLMs）通常能取得显著效果，但由于基于梯度的训练中反向传播的存在，这种方法往往以内存效率低下为代价。近期提出的内存高效零阶（MeZO）优化器通过仅需前向传播的训练方式有效缓解了这一问题，使其更具内存友好性。然而，与精确梯度相比，基于零阶（ZO）的梯度通常存在估计误差，这种误差可能严重干扰优化过程，导致收敛速度减缓和解质量下降。此外，我们发现估计误差对较大权重的负面影响远比对小权重的更显著。基于这一观察，本文提出了稀疏MeZO——一种新型的内存高效零阶优化方法，该方法仅对精心筛选的参数子集应用零阶优化。我们提出了一种简洁而有效的参数选择方案，该方案使稀疏MeZO获得了显著的性能提升。同时，我们开发了针对稀疏掩码的内存优化实现，确保算法仅需推理级别的内存消耗，从而使稀疏MeZO能够在单张A100 GPU上完成LLaMA-30b的微调。实验结果表明，稀疏MeZO在无额外开销的情况下，持续在性能和收敛速度上超越原始MeZO。例如，在RTE任务中，其绝对准确率较MeZO提升9%，且加速比达到3.5倍。代码已发布于 https://github.com/NUS-HPC-AI-Lab/SparseMeZO。