Тонкая настройка квантованных нейронных сетей с использованием оптимизации нулевого порядка

Fine-tuning Quantized Neural Networks with Zeroth-order Optimization

May 19, 2025

Авторы: Sifeng Shang, Jiayi Zhou, Chenyu Lin, Minxian Li, Kaiyang Zhou

cs.AI

Аннотация

По мере экспоненциального роста размеров крупных языковых моделей память GPU становится узким местом для адаптации этих моделей к последующим задачам. В данной работе мы стремимся расширить границы эффективного использования памяти при обучении, минимизируя потребление памяти на веса модели, градиенты и состояния оптимизатора в рамках единой структуры. Наша идея заключается в устранении как градиентов, так и состояний оптимизатора с использованием оптимизации нулевого порядка, которая аппроксимирует градиенты путем возмущения весов во время прямых проходов для определения направлений градиентов. Чтобы минимизировать использование памяти на веса, мы применяем квантование модели, например, преобразование из bfloat16 в int4. Однако прямое применение оптимизации нулевого порядка к квантованным весам невозможно из-за разрыва точности между дискретными весами и непрерывными градиентами, что потребовало бы де-квантования и повторного квантования. Для преодоления этой проблемы мы предлагаем Квантованную Оптимизацию Нулевого Порядка (QZO), новый подход, который возмущает непрерывную шкалу квантования для оценки градиентов и использует метод ограничения направленной производной для стабилизации обучения. QZO является ортогональным как к скалярным, так и к кодексным методам пост-тренировочного квантования. По сравнению с полной тонкой настройкой в bfloat16, QZO может снизить общую стоимость памяти более чем в 18 раз для 4-битных LLM и позволяет выполнять тонкую настройку Llama-2-13B и Stable Diffusion 3.5 Large на одном GPU с 24 ГБ памяти.

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.