Fine-tuning Quantized Neural Networks with Zeroth-order Optimization

随着大型语言模型的规模呈指数级增长，GPU 内存已成为将这些模型适配到下游任务时的主要瓶颈。在本文中，我们致力于通过最小化模型权重、梯度和优化器状态的内存使用，在一个统一的框架内推动高效内存训练的极限。我们的思路是利用零阶优化（zeroth-order optimization）来同时消除梯度和优化器状态，这种方法通过在前向传播过程中扰动权重以识别梯度方向，从而近似梯度。为了进一步减少权重的内存占用，我们采用了模型量化技术，例如从 bfloat16 转换为 int4。然而，直接将零阶优化应用于量化后的权重是不可行的，因为离散权重与连续梯度之间的精度差距会导致需要反量化和重新量化的过程。为了解决这一问题，我们提出了一种新颖的方法——量化零阶优化（Quantized Zeroth-order Optimization, QZO），该方法通过对连续的量化尺度进行扰动来进行梯度估计，并采用方向导数截断法（directional derivative clipping）来稳定训练过程。QZO 与基于标量的后训练量化方法以及基于码本的量化方法相互独立且可结合使用。相比于使用 bfloat16 进行全参数微调，QZO 可以将 4 位 LLM 的总内存成本降低超过 18 倍，并且能够在单个 24GB 的 GPU 上实现对 Llama-2-13B 和 Stable Diffusion 3.5 Large 的微调。

该论文试图解决大型语言模型（LLM）在下游任务微调过程中GPU内存不足的问题。随着模型规模的快速增长，存储权重、梯度和优化器状态所需的内存成为瓶颈。这是一个现有问题，但随着模型尺寸的增加，其重要性日益凸显。

论文提出了一种名为Quantized Zeroth-order Optimization（QZO）的新方法，通过零阶优化消除梯度和优化器状态的存储需求，并结合模型量化技术（如从bfloat16转换为int4）来减少权重存储需求。为了克服量化权重与连续梯度之间的精度差距，QZO引入了对连续量化尺度的扰动以估计梯度，并使用方向导数裁剪方法稳定训练。相比传统的全参数微调方法，QZO显著降低了内存消耗。

实验表明，QZO可以将4位LLM的总内存成本降低超过18倍，并成功在单个24GB GPU上实现Llama-2-13B和Stable Diffusion 3.5 Large的微调。论文还提供了详细的实验设计，包括不同模型和任务的对比结果。此外，作者提到代码可能会开源，这为后续研究提供了便利。未来值得深入研究的方向包括进一步优化QZO算法性能以及将其扩展到其他类型的深度学习模型中。

最近的相关研究包括：1)《Low-Rank Adaptation of Large Language Models》探讨了低秩分解在节省内存方面的应用；2)《Paged Optimizer: A Memory-Efficient Optimizer for Large-Scale Training》提出了分页优化器以缓解内存压力；3)《GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers》专注于生成式预训练Transformer的后训练量化技术；4)《AdaBits: Adaptive Precision Training for Transformer Models》研究了自适应精度训练方法。这些工作都围绕如何提高大规模模型训练的效率展开，而QZO提供了一个全新的视角，即结合零阶优化和量化技术来实现极低内存开销的微调。