VeLoRA : Entraînement efficace en mémoire grâce à des projections de sous-tokens de rang 1

VeLoRA: Memory Efficient Training using Rank-1 Sub-Token Projections

May 28, 2024

Auteurs: Roy Miles, Pradyumna Reddy, Ismail Elezi, Jiankang Deng

cs.AI

Résumé

Les grands modèles de langage (LLMs) sont récemment apparus comme des outils puissants pour aborder de nombreuses tâches de traitement du langage. Malgré leur succès, l'entraînement et le réglage fin de ces modèles restent encore trop coûteux en termes de calcul et de mémoire. Dans cet article, nous identifions et caractérisons les composants importants nécessaires à une convergence efficace des modèles en utilisant la descente de gradient. Ce faisant, nous constatons que les activations intermédiaires utilisées pour implémenter la rétropropagation peuvent être excessivement compressées sans entraîner de dégradation des performances. Ce résultat nous conduit à un algorithme peu coûteux et économe en mémoire pour le réglage fin et le pré-entraînement des LLMs. L'algorithme proposé divise simplement les tokens en sous-tokens plus petits avant de les projeter sur un sous-espace fixe à une dimension lors de la passe avant. Ces caractéristiques sont ensuite reconstruites de manière approximative lors de la passe arrière pour implémenter les règles de mise à jour. Nous confirmons l'efficacité de notre algorithme comme étant complémentaire à de nombreuses méthodes PEFT de pointe sur le benchmark de réglage fin VTAB-1k. De plus, nous surpassons QLoRA pour le réglage fin de LLaMA et montrons des performances compétitives par rapport à d'autres méthodes de pré-entraînement économes en mémoire sur le jeu de données à grande échelle C4.

English

Large language models (LLMs) have recently emerged as powerful tools for tackling many language-processing tasks. Despite their success, training and fine-tuning these models is still far too computationally and memory intensive. In this paper, we identify and characterise the important components needed for effective model convergence using gradient descent. In doing so we find that the intermediate activations used to implement backpropagation can be excessively compressed without incurring any degradation in performance. This result leads us to a cheap and memory-efficient algorithm for both fine-tuning and pre-training LLMs. The proposed algorithm simply divides the tokens up into smaller sub-tokens before projecting them onto a fixed 1-dimensional subspace during the forward pass. These features are then coarsely reconstructed during the backward pass to implement the update rules. We confirm the effectiveness of our algorithm as being complimentary to many state-of-the-art PEFT methods on the VTAB-1k fine-tuning benchmark. Furthermore, we outperform QLoRA for fine-tuning LLaMA and show competitive performance against other memory-efficient pre-training methods on the large-scale C4 dataset.