Afinamiento de Redes Neuronales Cuantizadas con Optimización de Orden Cero

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

Resumen

A medida que el tamaño de los modelos de lenguaje grandes crece exponencialmente, la memoria de las GPU se ha convertido en un cuello de botella para adaptar estos modelos a tareas específicas. En este artículo, buscamos superar los límites del entrenamiento eficiente en memoria minimizando el uso de memoria en los pesos del modelo, los gradientes y los estados del optimizador, dentro de un marco unificado. Nuestra idea es eliminar tanto los gradientes como los estados del optimizador utilizando optimización de orden cero, que aproxima los gradientes perturbando los pesos durante las pasadas hacia adelante para identificar las direcciones del gradiente. Para minimizar el uso de memoria en los pesos, empleamos cuantización del modelo, por ejemplo, convirtiendo de bfloat16 a int4. Sin embargo, aplicar directamente la optimización de orden cero a pesos cuantizados no es factible debido a la brecha de precisión entre los pesos discretos y los gradientes continuos, lo que requeriría des-cuantización y re-cuantización. Para superar este desafío, proponemos Optimización de Orden Cero Cuantizada (QZO), un enfoque novedoso que perturba la escala de cuantización continua para la estimación del gradiente y utiliza un método de recorte de derivadas direccionales para estabilizar el entrenamiento. QZO es ortogonal tanto a los métodos de cuantización post-entrenamiento basados en escalares como en libros de códigos. En comparación con el ajuste fino de parámetros completos en bfloat16, QZO puede reducir el costo total de memoria en más de 18 veces para modelos de lenguaje grandes de 4 bits, y permite ajustar Llama-2-13B y Stable Diffusion 3.5 Large en una sola GPU de 24 GB.

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.