Réglage fin des réseaux de neurones quantifiés par optimisation d'ordre zéro

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

Résumé

À mesure que la taille des grands modèles de langage croît de manière exponentielle, la mémoire des GPU est devenue un goulot d'étranglement pour l'adaptation de ces modèles à des tâches en aval. Dans cet article, nous visons à repousser les limites de l'entraînement efficace en mémoire en minimisant l'utilisation de la mémoire pour les poids du modèle, les gradients et les états de l'optimiseur, dans un cadre unifié. Notre idée est d'éliminer à la fois les gradients et les états de l'optimiseur en utilisant l'optimisation d'ordre zéro, qui approxime les gradients en perturbant les poids lors des passes avant pour identifier les directions des gradients. Pour minimiser l'utilisation de la mémoire sur les poids, nous employons la quantification du modèle, par exemple en convertissant de bfloat16 à int4. Cependant, l'application directe de l'optimisation d'ordre zéro à des poids quantifiés est impossible en raison de l'écart de précision entre les poids discrets et les gradients continus, ce qui nécessiterait sinon une dé-quantification et une re-quantification. Pour surmonter ce défi, nous proposons l'Optimisation d'Ordre Zéro Quantifiée (QZO), une approche novatrice qui perturbe l'échelle de quantification continue pour l'estimation des gradients et utilise une méthode de découpage des dérivées directionnelles pour stabiliser l'entraînement. QZO est orthogonale aux méthodes de quantification post-entraînement basées sur des scalaires et des codebooks. Par rapport à l'ajustement fin des paramètres complets en bfloat16, QZO peut réduire le coût total de la mémoire de plus de 18 fois pour les modèles de langage à 4 bits, et permet l'ajustement fin de Llama-2-13B et Stable Diffusion 3.5 Large sur un seul GPU de 24 Go.

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.