使用零阶优化微调量化神经网络

随着大型语言模型的规模呈指数级增长，GPU 内存已成为将这些模型应用于下游任务的瓶颈。在本文中，我们旨在通过在一个统一的框架内最小化模型权重、梯度和优化器状态的内存使用，从而突破内存高效训练的极限。我们的想法是使用零阶优化来消除梯度和优化器状态，零阶优化通过在前向传播过程中扰动权重来近似梯度方向。为了最小化权重上的内存使用，我们采用模型量化，例如将 bfloat16 转换为 int4。然而，直接将零阶优化应用于量化权重是不可行的，因为离散权重和连续梯度之间存在精度差距，否则需要反量化和重新量化。为了克服这一挑战，我们提出了量化零阶优化（QZO），这是一种通过扰动连续量化尺度进行梯度估计并使用方向导数裁剪方法来稳定训练的新方法。QZO 与基于标量和基于代码本的后训练量化方法正交。与 bfloat16 中的全参数微调相比，QZO 可以将 4 位 LLM 的总内存成本降低 18 倍以上，并使 Llama-2-13B 和 Stable Diffusion 3.5 Large 能够在单个 24GB GPU 内进行微调。