Fine-tuning van gekwantiseerde neurale netwerken met zeroth-order optimalisatie

Fine-tuning Quantized Neural Networks with Zeroth-order Optimization

May 19, 2025

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

cs.AI

Samenvatting

Naarmate de omvang van grote taalmodellen exponentieel toeneemt, is GPU-geheugen een knelpunt geworden bij het aanpassen van deze modellen voor downstream-taken. In dit artikel streven we ernaar de grenzen van geheugenefficiënte training te verleggen door het geheugengebruik voor modelgewichten, gradiënten en optimizerstatussen te minimaliseren binnen een uniform raamwerk. Ons idee is om zowel gradiënten als optimizerstatussen te elimineren door gebruik te maken van zeroth-order optimalisatie, waarbij gradiënten worden benaderd door gewichten te verstoren tijdens forward passes om gradiëntrichtingen te identificeren. Om het geheugengebruik voor gewichten te minimaliseren, passen we modelkwantisatie toe, bijvoorbeeld door conversie van bfloat16 naar int4. Het direct toepassen van zeroth-order optimalisatie op gekwantiseerde gewichten is echter niet haalbaar vanwege het precisieverschil tussen discrete gewichten en continue gradiënten, wat anders de-quantisatie en her-kwantisatie zou vereisen. Om deze uitdaging te overwinnen, stellen we Quantized Zeroth-order Optimization (QZO) voor, een nieuwe aanpak die de continue kwantisatieschaal verstoort voor gradiëntschatting en een directional derivative clipping-methode gebruikt om de training te stabiliseren. QZO staat los van zowel scalar-based als codebook-based post-training kwantisatiemethoden. Vergeleken met full-parameter fine-tuning in bfloat16 kan QZO de totale geheugenkosten met meer dan 18 keer verminderen voor 4-bit LLM's en maakt het fine-tuning van Llama-2-13B en Stable Diffusion 3.5 Large mogelijk binnen een enkele 24GB GPU.

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.