Entraînement LLM Économe en Mémoire avec Descente de Sous-Espace en Ligne

Memory-Efficient LLM Training with Online Subspace Descent

August 23, 2024

Auteurs: Kaizhao Liang, Bo Liu, Lizhang Chen, Qiang Liu

cs.AI

Résumé

Récemment, une large gamme d'algorithmes d'entraînement LLM efficaces en mémoire a gagné une popularité substantielle. Ces méthodes exploitent la structure de bas rang des gradients pour projeter les états de l'optimiseur dans un sous-espace en utilisant une matrice de projection trouvée par décomposition en valeurs singulières (SVD). Cependant, la convergence de ces algorithmes dépend fortement des règles de mise à jour de leur matrice de projection. Dans ce travail, nous fournissons la première garantie de convergence pour des règles de mise à jour arbitraires de la matrice de projection. Cette garantie s'applique généralement aux optimiseurs pouvant être analysés avec la Descente Hamiltonienne, y compris les plus courants tels que LION, Adam. Inspirés par notre compréhension théorique, nous proposons la Descente en Sous-espace en Ligne, une nouvelle famille d'optimiseurs de descente en sous-espace sans SVD. Au lieu de mettre à jour la matrice de projection avec des vecteurs propres, la Descente en Sous-espace en Ligne met à jour la matrice de projection avec une PCA en ligne. La Descente en Sous-espace en Ligne est flexible et n'introduit qu'un minimum de surcharge à l'entraînement. Nous montrons que pour la tâche de pré-entraînement des modèles LLaMA allant de 60M à 7B de paramètres sur l'ensemble de données C4, la Descente en Sous-espace en Ligne atteint une perplexité plus faible et de meilleures performances dans les tâches ultérieures que les méthodes d'entraînement de bas rang de pointe dans différents contextes, réduisant ainsi l'écart avec les références de rang complet.

English

Recently, a wide range of memory-efficient LLM training algorithms have gained substantial popularity. These methods leverage the low-rank structure of gradients to project optimizer states into a subspace using projection matrix found by singular value decomposition (SVD). However, convergence of these algorithms is highly dependent on the update rules of their projection matrix. In this work, we provide the first convergence guarantee for arbitrary update rules of projection matrix. This guarantee is generally applicable to optimizers that can be analyzed with Hamiltonian Descent, including most common ones, such as LION, Adam. Inspired by our theoretical understanding, we propose Online Subspace Descent, a new family of subspace descent optimizer without SVD. Instead of updating the projection matrix with eigenvectors, Online Subspace Descent updates the projection matrix with online PCA. Online Subspace Descent is flexible and introduces only minimum overhead to training. We show that for the task of pretraining LLaMA models ranging from 60M to 7B parameters on the C4 dataset, Online Subspace Descent achieves lower perplexity and better downstream tasks performance than state-of-the-art low-rank training methods across different settings and narrows the gap with full-rank baselines.