Treinamento eficiente em memória com Descida de Subespaço Online para LLM.

Memory-Efficient LLM Training with Online Subspace Descent

August 23, 2024

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

cs.AI

Resumo

Recentemente, uma ampla variedade de algoritmos de treinamento eficientes em memória para LLM ganharam substancial popularidade. Esses métodos aproveitam a estrutura de baixa classificação dos gradientes para projetar os estados do otimizador em um subespaço usando uma matriz de projeção encontrada pela decomposição em valores singulares (SVD). No entanto, a convergência desses algoritmos depende fortemente das regras de atualização de sua matriz de projeção. Neste trabalho, fornecemos a primeira garantia de convergência para regras de atualização arbitrárias da matriz de projeção. Essa garantia é geralmente aplicável a otimizadores que podem ser analisados com Descida Hamiltoniana, incluindo os mais comuns, como LION, Adam. Inspirados em nosso entendimento teórico, propomos o Descent Subespaço Online, uma nova família de otimizadores de descida de subespaço sem SVD. Em vez de atualizar a matriz de projeção com autovetores, o Descent Subespaço Online atualiza a matriz de projeção com PCA online. O Descent Subespaço Online é flexível e introduz apenas um mínimo de sobrecarga ao treinamento. Mostramos que, para a tarefa de pré-treinamento de modelos LLaMA com parâmetros variando de 60M a 7B no conjunto de dados C4, o Descent Subespaço Online alcança menor perplexidade e melhor desempenho em tarefas subsequentes do que os métodos de treinamento de baixa classificação de última geração em diferentes configurações e reduz a diferença com as linhas de base de classificação completa.

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.