Fine-tuning Quantized Neural Networks with Zeroth-order Optimization
Authors: name: Sifeng Shang, name: Jiayi Zhou, name: Chenyu Lin, name: Minxian Li, name: Kaiyang Zhou

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 novel 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 bfloat16, QZO can reduce the total memory cost by
more than 18times for 4-bit LLMs, and enables fine-tuning Llama-2-13B and
Stable Diffusion 3.5 Large within a single 24GB GPU.