GaLore : Entraînement efficace en mémoire des LLM par projection de gradient en bas rang
GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection
March 6, 2024
Auteurs: Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, Yuandong Tian
cs.AI
Résumé
L'entraînement des grands modèles de langage (LLMs) présente des défis mémoriels importants, principalement en raison de la taille croissante des poids et des états de l'optimiseur. Les approches courantes de réduction de la mémoire, telles que l'adaptation de bas rang (LoRA), ajoutent une matrice de bas rang entraînable aux poids pré-entraînés gelés dans chaque couche, réduisant ainsi les paramètres entraînables et les états de l'optimiseur. Cependant, ces approches sous-performent généralement l'entraînement avec des poids de rang complet, tant lors du pré-entraînement que du fine-tuning, car elles limitent la recherche des paramètres à un sous-espace de bas rang et altèrent la dynamique d'entraînement. De plus, elles peuvent nécessiter un démarrage à chaud avec des poids de rang complet. Dans ce travail, nous proposons la Projection de Gradient de Bas Rang (GaLore), une stratégie d'entraînement qui permet un apprentissage à paramètres complets tout en étant plus économe en mémoire que les méthodes courantes d'adaptation de bas rang comme LoRA. Notre approche réduit l'utilisation de la mémoire jusqu'à 65,5 % dans les états de l'optimiseur tout en maintenant à la fois l'efficacité et les performances pour le pré-entraînement sur les architectures LLaMA 1B et 7B avec le jeu de données C4 contenant jusqu'à 19,7 milliards de tokens, ainsi que pour le fine-tuning de RoBERTa sur les tâches GLUE. Notre version 8 bits de GaLore réduit encore la mémoire de l'optimiseur jusqu'à 82,5 % et la mémoire totale d'entraînement de 63,3 %, par rapport à une référence BF16. Notamment, nous démontrons, pour la première fois, la faisabilité du pré-entraînement d'un modèle de 7B sur des GPU grand public avec 24 Go de mémoire (par exemple, NVIDIA RTX 4090) sans recourir à des stratégies de parallélisation de modèle, de point de contrôle ou de déchargement.
English
Training Large Language Models (LLMs) presents significant memory challenges,
predominantly due to the growing size of weights and optimizer states. Common
memory-reduction approaches, such as low-rank adaptation (LoRA), add a
trainable low-rank matrix to the frozen pre-trained weight in each layer,
reducing trainable parameters and optimizer states. However, such approaches
typically underperform training with full-rank weights in both pre-training and
fine-tuning stages since they limit the parameter search to a low-rank subspace
and alter the training dynamics, and further, may require full-rank warm start.
In this work, we propose Gradient Low-Rank Projection (GaLore), a training
strategy that allows full-parameter learning but is more memory-efficient than
common low-rank adaptation methods such as LoRA. Our approach reduces memory
usage by up to 65.5% in optimizer states while maintaining both efficiency and
performance for pre-training on LLaMA 1B and 7B architectures with C4 dataset
with up to 19.7B tokens, and on fine-tuning RoBERTa on GLUE tasks. Our 8-bit
GaLore further reduces optimizer memory by up to 82.5% and total training
memory by 63.3%, compared to a BF16 baseline. Notably, we demonstrate, for the
first time, the feasibility of pre-training a 7B model on consumer GPUs with
24GB memory (e.g., NVIDIA RTX 4090) without model parallel, checkpointing, or
offloading strategies.