Entrenamiento Eficiente en Memoria de LLM con Descenso de Subespacio en Línea

Memory-Efficient LLM Training with Online Subspace Descent

August 23, 2024

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

cs.AI

Resumen

Recientemente, una amplia variedad de algoritmos de entrenamiento LLM eficientes en memoria han ganado una considerable popularidad. Estos métodos aprovechan la estructura de bajo rango de los gradientes para proyectar los estados del optimizador en un subespacio utilizando una matriz de proyección encontrada mediante descomposición en valores singulares (SVD). Sin embargo, la convergencia de estos algoritmos depende en gran medida de las reglas de actualización de su matriz de proyección. En este trabajo, proporcionamos la primera garantía de convergencia para reglas de actualización arbitrarias de la matriz de proyección. Esta garantía es generalmente aplicable a optimizadores que pueden ser analizados con Descenso Hamiltoniano, incluyendo los más comunes, como LION, Adam. Inspirados por nuestra comprensión teórica, proponemos Descenso de Subespacio en Línea, una nueva familia de optimizadores de descenso de subespacio sin SVD. En lugar de actualizar la matriz de proyección con autovectores, Descenso de Subespacio en Línea actualiza la matriz de proyección con PCA en línea. Descenso de Subespacio en Línea es flexible e introduce solo un mínimo sobrecosto al entrenamiento. Mostramos que para la tarea de preentrenamiento de modelos LLaMA que van desde 60M hasta 7B parámetros en el conjunto de datos C4, Descenso de Subespacio en Línea logra una menor perplejidad y un mejor rendimiento en tareas posteriores que los métodos de entrenamiento de bajo rango de vanguardia en diferentes configuraciones, y reduce la brecha con los baselines de rango completo.

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.