Q-GaLore : Quantification GaLore avec Projection INT4 et Gradients de Bas Rang Adaptatifs par Couche

Q-GaLore: Quantized GaLore with INT4 Projection and Layer-Adaptive Low-Rank Gradients

July 11, 2024

Auteurs: Zhenyu Zhang, Ajay Jaiswal, Lu Yin, Shiwei Liu, Jiawei Zhao, Yuandong Tian, Zhangyang Wang

cs.AI

Résumé

L'entraînement des grands modèles de langage (LLMs) est gourmand en mémoire en raison du grand nombre de paramètres et des états d'optimisation associés. GaLore, une méthode récente, réduit l'utilisation de la mémoire en projetant les gradients des poids dans un sous-espace de faible rang sans compromettre les performances. Cependant, GaLore repose sur des opérations de décomposition en valeurs singulières (SVD) chronophages pour identifier le sous-espace, et les mises à jour fréquentes du sous-espace entraînent un surcoût significatif en temps d'entraînement. De plus, GaLore offre des améliorations minimales en précision et efficacité par rapport à LoRA dans des scénarios de fine-tuning plus accessibles. Pour pallier ces limitations, nous introduisons Q-Galore, une approche novatrice qui réduit considérablement l'utilisation de la mémoire en combinant la quantification et la projection en faible rang, surpassant les avantages de GaLore. Notre méthode repose sur deux observations clés : (i) le sous-espace des gradients présente des propriétés variées, certaines couches convergeant tôt lors de l'entraînement tandis que d'autres sont sujettes à des changements fréquents ; (ii) les matrices de projection sont très résilientes à la quantification en faible précision. En exploitant ces insights, Q-Galore met à jour de manière adaptative le sous-espace des gradients en fonction de ses statistiques de convergence, atteignant des performances comparables tout en réduisant significativement le nombre d'opérations SVD. Nous maintenons les matrices de projection au format INT4 et les poids au format INT8, en incorporant un arrondi stochastique pour capturer les informations cumulées des gradients. Cette approche permet une trajectoire d'entraînement de haute précision en utilisant uniquement des poids de faible précision. Nous démontrons que Q-Galore atteint des performances très compétitives avec une efficacité mémoire exceptionnelle. En pré-entraînement, Q-Galore facilite l'entraînement d'un modèle LLaMA-7B à partir de zéro sur une seule carte NVIDIA RTX 4060 Ti avec seulement 16 Go de mémoire. En fine-tuning, il réduit la consommation mémoire jusqu'à 50 % par rapport à LoRA et GaLore, tout en surpassant systématiquement QLoRA à un coût mémoire équivalent.

English

Training Large Language Models (LLMs) is memory-intensive due to the large number of parameters and associated optimization states. GaLore, a recent method, reduces memory usage by projecting weight gradients into a low-rank subspace without compromising performance. However, GaLore relies on time-consuming Singular Value Decomposition (SVD) operations to identify the subspace, and the frequent subspace updates lead to significant training time overhead. Moreover, GaLore offers minimal improvements in accuracy and efficiency compared to LoRA in more accessible fine-tuning scenarios. To address these limitations, we introduce Q-Galore, a novel approach that substantially reduces memory usage by combining quantization and low-rank projection, surpassing the benefits of GaLore. Our method is based on two key observations: (i) the gradient subspace exhibits diverse properties, with some layers converging early in training while others are subject to frequent changes; (ii) the projection matrices are highly resilient to low-bit quantization. Leveraging these insights, Q-GaLore adaptively updates the gradient subspace based on its convergence statistics, achieving comparable performance while significantly reducing the number of SVD operations. We maintain the projection matrices in INT4 format and weights in INT8 format, incorporating stochastic rounding to capture accumulated gradient information. This approach enables a high-precision training trajectory using only low-precision weights. We demonstrate that Q-GaLore achieves highly competitive performance with exceptional memory efficiency. At pre-training, Q-GaLore facilitates training a LLaMA-7B model from scratch on a single NVIDIA RTX 4060 Ti with only 16 GB memory. At fine-tuning, it reduces memory consumption by up to 50% compared to LoRA and GaLore, while consistently outperforming QLoRA at the same memory cost.