GaLore: Treinamento Eficiente em Memória de LLMs por Projeção de Gradiente de Baixa Dimensionalidade

GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection

March 6, 2024

Autores: Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, Yuandong Tian

cs.AI

Resumo

O treinamento de Modelos de Linguagem de Grande Escala (LLMs) apresenta desafios significativos de memória, principalmente devido ao tamanho crescente dos pesos e estados do otimizador. Abordagens comuns para redução de memória, como a adaptação de baixo posto (LoRA), adicionam uma matriz de baixo posto treinável aos pesos pré-treinados congelados em cada camada, reduzindo os parâmetros treináveis e os estados do otimizador. No entanto, tais abordagens geralmente têm desempenho inferior ao treinamento com pesos de posto completo tanto na fase de pré-treinamento quanto na de ajuste fino, pois limitam a busca de parâmetros a um subespaço de baixo posto e alteram a dinâmica do treinamento, além de poderem exigir um aquecimento inicial com posto completo. Neste trabalho, propomos a Projeção de Baixo Posto de Gradiente (GaLore), uma estratégia de treinamento que permite o aprendizado de parâmetros completos, mas é mais eficiente em termos de memória do que métodos comuns de adaptação de baixo posto, como o LoRA. Nossa abordagem reduz o uso de memória em até 65,5% nos estados do otimizador, mantendo tanto a eficiência quanto o desempenho no pré-treinamento das arquiteturas LLaMA 1B e 7B com o conjunto de dados C4 com até 19,7 bilhões de tokens, e no ajuste fino do RoBERTa em tarefas do GLUE. Nossa versão de 8 bits do GaLore reduz ainda mais a memória do otimizador em até 82,5% e a memória total de treinamento em 63,3%, em comparação com uma linha de base BF16. Notavelmente, demonstramos, pela primeira vez, a viabilidade de pré-treinar um modelo de 7B em GPUs de consumo com 24GB de memória (por exemplo, NVIDIA RTX 4090) sem estratégias de paralelismo de modelo, checkpointing ou offloading.

English

Training Large Language Models (LLMs) presents significant memory challenges, predominantly due to the growing size of weights and optimizer states. Common memory-reduction approaches, such as low-rank adaptation (LoRA), add a trainable low-rank matrix to the frozen pre-trained weight in each layer, reducing trainable parameters and optimizer states. However, such approaches typically underperform training with full-rank weights in both pre-training and fine-tuning stages since they limit the parameter search to a low-rank subspace and alter the training dynamics, and further, may require full-rank warm start. In this work, we propose Gradient Low-Rank Projection (GaLore), a training strategy that allows full-parameter learning but is more memory-efficient than common low-rank adaptation methods such as LoRA. Our approach reduces memory usage by up to 65.5% in optimizer states while maintaining both efficiency and performance for pre-training on LLaMA 1B and 7B architectures with C4 dataset with up to 19.7B tokens, and on fine-tuning RoBERTa on GLUE tasks. Our 8-bit GaLore further reduces optimizer memory by up to 82.5% and total training memory by 63.3%, compared to a BF16 baseline. Notably, we demonstrate, for the first time, the feasibility of pre-training a 7B model on consumer GPUs with 24GB memory (e.g., NVIDIA RTX 4090) without model parallel, checkpointing, or offloading strategies.