Simultaneous Computation and Memory Efficient Zeroth-Order Optimizer for Fine-Tuning Large Language Models

📄 arXiv: 2410.09823v1 📥 PDF

作者: Fei Wang, Li Shen, Liang Ding, Chao Xue, Ye Liu, Changxing Ding

分类: cs.LG, cs.CL

发布日期: 2024-10-13

💡 一句话要点

提出LeZO：一种计算与内存高效的零阶优化器，用于微调大型语言模型

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 零阶优化 大型语言模型 微调 内存效率 稀疏计算

📋 核心要点

  1. 大型语言模型微调面临内存消耗巨大挑战，零阶优化虽能缓解，但训练时间较长。
  2. LeZO通过层级稀疏计算，动态扰动参数子集，在SPSA和ZO-SGD中融入层级参数稀疏性。
  3. 实验表明，LeZO在SuperGLUE等任务上加速训练，且不损失ZO优化性能，速度提升显著。

📝 摘要（中文）

微调是使大型语言模型适应下游任务的有效方法，但通常会导致巨大的内存消耗。一种有前景的缓解方法是使用零阶（ZO）优化，它估计梯度以代替一阶（FO）梯度计算，但由于其随机性，训练时间更长。通过重新审视内存高效的ZO（MeZO）优化器，我们发现全参数扰动和更新过程消耗了其总体微调时间成本的50％以上。基于这些观察，我们引入了一种新颖的层级稀疏计算和内存高效的ZO优化器，名为LeZO。LeZO将层视为稀疏化的基本单元，并在每个步骤中动态扰动不同的参数子集，以实现全参数微调。LeZO在同步扰动随机逼近（SPSA）和ZO随机梯度下降（ZO-SGD）过程中结合了层级参数稀疏性。它实现了扰动和更新过程中的加速计算，而无需额外的内存开销。我们使用OPT模型系列在SuperGLUE基准和两个生成任务上进行了广泛的实验。实验表明，LeZO在不影响ZO优化性能的情况下加速了训练。具体来说，与MeZO相比，它在SST-2，BoolQ和Copa任务上实现了超过3倍的加速。

🔬 方法详解

问题定义：论文旨在解决大型语言模型微调过程中内存消耗过大以及零阶优化器训练速度慢的问题。现有MeZO优化器虽然降低了内存消耗，但其全参数扰动和更新过程耗时过长，成为性能瓶颈。

核心思路：LeZO的核心思路是引入层级稀疏计算，将模型层视为稀疏化的基本单元，并在每次迭代中只扰动部分层的参数。通过动态选择扰动的参数子集，LeZO能够在保证全参数微调的前提下，显著减少计算量和内存占用。

技术框架：LeZO的整体框架基于零阶优化，主要包含以下几个阶段：1）层级稀疏化：将模型划分为多个层，并为每一层分配一个稀疏度；2）参数扰动：根据层级稀疏度，选择部分层进行参数扰动；3）梯度估计：使用同步扰动随机逼近（SPSA）或ZO随机梯度下降（ZO-SGD）估计梯度；4）参数更新：根据估计的梯度更新模型参数。

关键创新：LeZO的关键创新在于层级稀疏计算。与传统的全参数扰动方法相比，LeZO只扰动部分层的参数，从而显著减少了计算量和内存占用。此外，LeZO动态调整每一层的稀疏度，以平衡计算效率和模型性能。

关键设计：LeZO的关键设计包括：1）层级稀疏度的选择策略：论文可能采用了某种启发式方法或学习算法来确定每一层的稀疏度，以保证模型性能；2）扰动参数的选择方法：论文可能采用了随机选择或重要性采样等方法来选择扰动的参数；3）梯度估计方法：论文采用了SPSA或ZO-SGD等零阶优化算法来估计梯度。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，LeZO在SuperGLUE基准测试和生成任务上，与MeZO相比实现了超过3倍的加速，同时保持了与MeZO相当的性能。具体而言，在SST-2、BoolQ和Copa等任务上，LeZO的训练速度显著提升，验证了其计算和内存效率。

🎯 应用场景

LeZO优化器可广泛应用于大型语言模型的微调，尤其是在资源受限的环境下，例如移动设备或边缘计算平台。该方法能够降低微调过程中的内存消耗和计算成本，加速模型部署，并促进大型语言模型在更多实际场景中的应用。未来，LeZO的层级稀疏思想可以推广到其他深度学习模型和优化算法中。

📄 摘要（原文）

Fine-tuning is powerful for adapting large language models to downstream tasks, but it often results in huge memory usages. A promising approach to mitigate this is using Zeroth-Order (ZO) optimization, which estimates gradients to replace First-Order (FO) gradient calculations, albeit with longer training time due to its stochastic nature. By revisiting the Memory-efficient ZO (MeZO) optimizer, we discover that the full-parameter perturbation and updating processes consume over 50% of its overall fine-tuning time cost. Based on these observations, we introduce a novel layer-wise sparse computation and memory efficient ZO optimizer, named LeZO. LeZO treats layers as fundamental units for sparsification and dynamically perturbs different parameter subsets in each step to achieve full-parameter fine-tuning. LeZO incorporates layer-wise parameter sparsity in the process of simultaneous perturbation stochastic approximation (SPSA) and ZO stochastic gradient descent (ZO-SGD). It achieves accelerated computation during perturbation and updating processes without additional memory overhead. We conduct extensive experiments with the OPT model family on the SuperGLUE benchmark and two generative tasks. The experiments show that LeZO accelerates training without compromising the performance of ZO optimization. Specifically, it achieves over 3x speedup compared to MeZO on the SST-2, BoolQ, and Copa tasks.