Ajuste Fino de Redes Neurais Quantizadas com Otimização de Ordem Zero

Fine-tuning Quantized Neural Networks with Zeroth-order Optimization

May 19, 2025

Autores: Sifeng Shang, Jiayi Zhou, Chenyu Lin, Minxian Li, Kaiyang Zhou

cs.AI

Resumo

À medida que o tamanho dos grandes modelos de linguagem cresce exponencialmente, a memória da GPU tornou-se um gargalo para adaptar esses modelos a tarefas subsequentes. Neste artigo, buscamos expandir os limites do treinamento eficiente em memória, minimizando o uso de memória em pesos do modelo, gradientes e estados do otimizador, dentro de uma estrutura unificada. Nossa ideia é eliminar tanto os gradientes quanto os estados do otimizador usando otimização de ordem zero, que aproxima os gradientes perturbando os pesos durante as passagens diretas para identificar as direções dos gradientes. Para minimizar o uso de memória nos pesos, empregamos quantização de modelos, por exemplo, convertendo de bfloat16 para int4. No entanto, aplicar diretamente a otimização de ordem zero a pesos quantizados é inviável devido à diferença de precisão entre pesos discretos e gradientes contínuos, o que exigiria desquantização e requantização. Para superar esse desafio, propomos a Otimização de Ordem Zero Quantizada (QZO), uma abordagem inovadora que perturba a escala de quantização contínua para estimar gradientes e utiliza um método de corte de derivada direcional para estabilizar o treinamento. O QZO é ortogonal tanto aos métodos de quantização pós-treinamento baseados em escalares quanto em codebooks. Em comparação com o ajuste fino de parâmetros completos em bfloat16, o QZO pode reduzir o custo total de memória em mais de 18 vezes para LLMs de 4 bits, e permite o ajuste fino de Llama-2-13B e Stable Diffusion 3.5 Large em uma única GPU de 24GB.

English

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.