Entraîner petit, inférer grand : un entraînement LoRA économe en mémoire pour les grands modèles de langage

Train Small, Infer Large: Memory-Efficient LoRA Training for Large Language Models

February 19, 2025

Auteurs: Jun Zhang, Jue Wang, Huan Li, Lidan Shou, Ke Chen, Yang You, Guiming Xie, Xuejian Gong, Kunlong Zhou

cs.AI

Résumé

Les modèles de langage de grande taille (LLMs) ont considérablement fait progresser le traitement du langage naturel grâce à leurs capacités exceptionnelles de généralisation des tâches. L'adaptation à faible rang (LoRA) offre une solution de fine-tuning économique, en gelant les paramètres originaux du modèle et en entraînant uniquement des matrices d'adaptation légères et de faible rang. Cependant, l'empreinte mémoire de LoRA est largement dominée par les paramètres originaux du modèle. Pour atténuer ce problème, nous proposons LoRAM, un schéma d'entraînement LoRA économe en mémoire, fondé sur l'intuition que de nombreux neurones dans les LLMs sur-paramétrés ont une utilité d'entraînement faible mais sont essentiels pour l'inférence. LoRAM introduit une approche unique : il s'entraîne sur un modèle élagué (de petite taille) pour obtenir des matrices de faible rang élaguées, qui sont ensuite récupérées et utilisées avec le modèle original (de grande taille) pour l'inférence. De plus, un pré-entraînement continu à coût minimal, effectué à l'avance par les éditeurs du modèle, aligne les écarts de connaissances entre les modèles élagués et originaux. Nos expériences approfondies démontrent l'efficacité de LoRAM à travers diverses stratégies d'élagage et tâches en aval. Pour un modèle de 70 milliards de paramètres, LoRAM permet l'entraînement sur un GPU avec seulement 20 Go de HBM, remplaçant un GPU A100-80G pour l'entraînement LoRA et 15 GPU pour le fine-tuning complet. Plus précisément, QLoRAM, mis en œuvre par élagage structuré combiné à une quantification sur 4 bits, pour LLaMA-3.1-70B (LLaMA-2-70B), réduit le coût de stockage des paramètres qui domine l'utilisation de la mémoire dans l'entraînement des matrices de faible rang par un facteur de 15,81 (16,95), tout en obtenant des gains de performance significatifs par rapport à la fois au LLaMA-3.1-70B (LLaMA-2-70B) original et au LLaMA-3.1-8B (LLaMA-2-13B) entraîné avec LoRA.

English

Large Language Models (LLMs) have significantly advanced natural language processing with exceptional task generalization capabilities. Low-Rank Adaption (LoRA) offers a cost-effective fine-tuning solution, freezing the original model parameters and training only lightweight, low-rank adapter matrices. However, the memory footprint of LoRA is largely dominated by the original model parameters. To mitigate this, we propose LoRAM, a memory-efficient LoRA training scheme founded on the intuition that many neurons in over-parameterized LLMs have low training utility but are essential for inference. LoRAM presents a unique twist: it trains on a pruned (small) model to obtain pruned low-rank matrices, which are then recovered and utilized with the original (large) model for inference. Additionally, minimal-cost continual pre-training, performed by the model publishers in advance, aligns the knowledge discrepancy between pruned and original models. Our extensive experiments demonstrate the efficacy of LoRAM across various pruning strategies and downstream tasks. For a model with 70 billion parameters, LoRAM enables training on a GPU with only 20G HBM, replacing an A100-80G GPU for LoRA training and 15 GPUs for full fine-tuning. Specifically, QLoRAM implemented by structured pruning combined with 4-bit quantization, for LLaMA-3.1-70B (LLaMA-2-70B), reduces the parameter storage cost that dominates the memory usage in low-rank matrix training by 15.81times (16.95times), while achieving dominant performance gains over both the original LLaMA-3.1-70B (LLaMA-2-70B) and LoRA-trained LLaMA-3.1-8B (LLaMA-2-13B).