VeLoRA: Treino Eficiente em Memória Usando Projeções de Sub-Tokens de Rank-1

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

May 28, 2024

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

cs.AI

Resumo

Modelos de linguagem de grande escala (LLMs) surgiram recentemente como ferramentas poderosas para abordar diversas tarefas de processamento de linguagem. Apesar do seu sucesso, o treinamento e o ajuste fino desses modelos ainda são extremamente intensivos em termos computacionais e de memória. Neste artigo, identificamos e caracterizamos os componentes importantes necessários para a convergência eficaz do modelo usando gradiente descendente. Ao fazer isso, descobrimos que as ativações intermediárias usadas para implementar a retropropagação podem ser excessivamente comprimidas sem causar qualquer degradação no desempenho. Esse resultado nos leva a um algoritmo barato e eficiente em termos de memória para ajuste fino e pré-treinamento de LLMs. O algoritmo proposto simplesmente divide os tokens em sub-tokens menores antes de projetá-los em um subespaço unidimensional fixo durante a passagem direta. Essas características são então reconstruídas de forma grosseira durante a passagem reversa para implementar as regras de atualização. Confirmamos a eficácia do nosso algoritmo como complementar a muitos métodos PEFT de última geração no benchmark de ajuste fino VTAB-1k. Além disso, superamos o QLoRA no ajuste fino do LLaMA e mostramos desempenho competitivo contra outros métodos de pré-treinamento eficientes em memória no grande conjunto de dados 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.