面向大语言模型微调的同步计算与内存高效的零阶优化器
Simultaneous Computation and Memory Efficient Zeroth-Order Optimizer for Fine-Tuning Large Language Models
AI总结:
针对大语言模型微调中零阶优化器全参数扰动与更新耗时过长的问题,提出逐层稀疏的LeZO优化器,在无额外内存开销下实现超过3倍加速且不损失性能。
摘要:
微调在将大语言模型适配到下游任务方面非常强大,但通常会导致巨大的内存消耗。缓解这一问题的有效方法是使用零阶(ZO)优化,它通过估计梯度来替代一阶(FO)梯度计算,但由于其随机性,会导致训练时间更长。通过重新审视内存高效的ZO(MeZO)优化器,我们发现全参数扰动和更新过程消耗了其整体微调时间成本的50%以上。基于这些观察,我们引入了一种新颖的逐层稀疏计算与内存高效的ZO优化器,命名为LeZO。LeZO将层作为稀疏化的基本单元,并在每一步动态扰动不同的参数子集,以实现全参数微调。LeZO在同步扰动随机近似(SPSA)和ZO随机梯度下降(ZO-SGD)过程中引入了逐层参数稀疏性。它在扰动和更新过程中实现了加速计算,且没有额外的内存开销。我们在SuperGLUE基准和两个生成任务上使用OPT模型家族进行了大量实验。实验表明,LeZO在不影响ZO优化性能的情况下加速了训练。具体而言,与MeZO相比,它在SST-2、BoolQ和Copa任务上实现了超过3倍的加速。
英文摘要:
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.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。