arXiv:2410.09823cs.LGcs.CL2024-10被引 7

提出LeZO优化器,让大模型微调更快更省内存。

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

  • 按层稀疏化设计,动态扰动参数子集加速计算
  • 在SST-2等任务上比MeZO快3倍以上,性能相当
  • 适合资源受限场景下高效微调大模型

微调大语言模型虽有效,但内存开销巨大。零阶(ZO)优化通过估计梯度替代一阶梯度计算,可缓解内存压力,但训练时间更长。我们重新分析了内存高效的零阶优化器MeZO,发现全参数扰动和更新过程占总耗时超50%。为此提出新型分层稀疏计算与内存高效的零阶优化器LeZO。LeZO以层为单位进行稀疏化,在同时扰动随机近似(SPSA)和零阶随机梯度下降(ZO-SGD)中引入层级参数稀疏性,实现扰动与更新过程加速且无额外内存开销。在OPT模型族上的SuperGLUE基准及两个生成任务上的实验表明,LeZO在不降低零阶优化性能的前提下显著加速训练:在SST-2、BoolQ和Copa任务上相比MeZO提速超3倍。

原文摘要 · Abstract (English)

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.

大模型微调零阶优化稀疏计算

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。