As the size of large language models grows exponentially, GPU memory has become a bottleneck for adapting these models to downstream tasks. In this paper, we aim to push the limits of memory-efficient training by minimizing memory usage on model weights, gradients, and optimizer states, within a unified framework. Our idea is to eliminate both gradients and optimizer states using zeroth-order optimization, which approximates gradients by perturbing weights during forward passes to identify gradient directions. To minimize memory usage on weights, we employ model quantization, e.g., converting from bfloat16 to int4. However, directly applying zeroth-order optimization to quantized weights is infeasible due to the precision gap between discrete weights and continuous gradients, which would otherwise require de-quantization and re-quantization. To overcome this challenge, we propose Quantized Zeroth-order Optimization (QZO), a simple yet effective approach that perturbs the continuous quantization scale for gradient estimation and uses a directional derivative clipping method to stabilize training. QZO is orthogonal to both scalar-based and codebook-based post-training quantization methods. Compared to full-parameter fine-tuning in 16 bits, QZO can reduce the total memory cost by more than 18$\times$ for 4-bit LLMs, and enables fine-tuning Llama-2-13B within a single 24GB GPU.

翻译：随着大语言模型规模的指数级增长，GPU内存已成为将这些模型适配至下游任务的瓶颈。本文旨在通过最小化模型权重、梯度和优化器状态的内存占用，在统一框架内突破内存高效训练的极限。我们的核心思路是利用零阶优化同时消除梯度和优化器状态，该方法通过在前向传播中扰动权重来估计梯度方向。为最小化权重的内存占用，我们采用模型量化技术，例如将bfloat16转换为int4。然而，由于离散权重与连续梯度之间的精度鸿沟，直接将零阶优化应用于量化权重并不可行——否则需要反复执行反量化和再量化操作。为攻克此难题，我们提出量化零阶优化（QZO），这是一种简洁而高效的方法：通过扰动连续量化尺度进行梯度估计，并采用方向导数裁剪技术以稳定训练过程。QZO与基于标量的和基于码本的后训练量化方法均保持正交性。相较于16位全参数微调，QZO可将4位大语言模型的总内存开销降低超过18倍，并能在单块24GB GPU上完成Llama-2-13B的微调。