Ottimizzazione Fine-tuning di Reti Neurali Quantizzate con Ottimizzazione di Ordine Zero

Fine-tuning Quantized Neural Networks with Zeroth-order Optimization

May 19, 2025

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

cs.AI

Abstract

Con l'aumento esponenziale delle dimensioni dei modelli linguistici di grandi dimensioni, la memoria delle GPU è diventata un collo di bottiglia per l'adattamento di questi modelli a compiti downstream. In questo articolo, miriamo a spingere i limiti dell'addestramento efficiente in termini di memoria minimizzando l'utilizzo della memoria per i pesi del modello, i gradienti e gli stati dell'ottimizzatore, all'interno di un framework unificato. La nostra idea è eliminare sia i gradienti che gli stati dell'ottimizzatore utilizzando l'ottimizzazione di ordine zero, che approssima i gradienti perturbando i pesi durante le passate in avanti per identificare le direzioni del gradiente. Per minimizzare l'utilizzo della memoria sui pesi, impieghiamo la quantizzazione del modello, ad esempio convertendo da bfloat16 a int4. Tuttavia, applicare direttamente l'ottimizzazione di ordine zero ai pesi quantizzati è infattibile a causa del divario di precisione tra i pesi discreti e i gradienti continui, che altrimenti richiederebbero de-quantizzazione e ri-quantizzazione. Per superare questa sfida, proponiamo l'Ottimizzazione di Ordine Zero Quantizzata (QZO), un approccio innovativo che perturba la scala di quantizzazione continua per la stima del gradiente e utilizza un metodo di clipping della derivata direzionale per stabilizzare l'addestramento. QZO è ortogonale sia ai metodi di quantizzazione post-addestramento basati su scalari che su codebook. Rispetto al fine-tuning completo dei parametri in bfloat16, QZO può ridurre il costo totale della memoria di oltre 18 volte per LLM a 4 bit e consente il fine-tuning di Llama-2-13B e Stable Diffusion 3.5 Large all'interno di una singola GPU da 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.