GaLore: Entrenamiento Eficiente en Memoria de Modelos de Lenguaje mediante Proyección de Gradientes de Bajo Rango

Resumen

El entrenamiento de modelos de lenguaje grandes (LLMs, por sus siglas en inglés) presenta desafíos significativos de memoria, principalmente debido al creciente tamaño de los pesos y los estados del optimizador. Los enfoques comunes para reducir la memoria, como la adaptación de bajo rango (LoRA), añaden una matriz de bajo rango entrenable a los pesos preentrenados congelados en cada capa, reduciendo los parámetros entrenables y los estados del optimizador. Sin embargo, estos enfoques suelen tener un rendimiento inferior al entrenamiento con pesos de rango completo tanto en las etapas de preentrenamiento como de ajuste fino, ya que limitan la búsqueda de parámetros a un subespacio de bajo rango y alteran la dinámica del entrenamiento, además de que pueden requerir un inicio cálido con rango completo. En este trabajo, proponemos Proyección de Bajo Rango de Gradiente (GaLore), una estrategia de entrenamiento que permite el aprendizaje de parámetros completos pero es más eficiente en memoria que los métodos comunes de adaptación de bajo rango como LoRA. Nuestro enfoque reduce el uso de memoria hasta un 65.5% en los estados del optimizador, manteniendo tanto la eficiencia como el rendimiento en el preentrenamiento de las arquitecturas LLaMA 1B y 7B con el conjunto de datos C4 con hasta 19.7 mil millones de tokens, y en el ajuste fino de RoBERTa en tareas de GLUE. Nuestra versión de 8 bits de GaLore reduce aún más la memoria del optimizador hasta un 82.5% y la memoria total de entrenamiento en un 63.3%, en comparación con una línea base BF16. Notablemente, demostramos, por primera vez, la viabilidad de preentrenar un modelo de 7B en GPUs de consumo con 24GB de memoria (por ejemplo, NVIDIA RTX 4090) sin estrategias de paralelismo de modelos, checkpointing o descarga.

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.

GaLore: Entrenamiento Eficiente en Memoria de Modelos de Lenguaje mediante Proyección de Gradientes de Bajo Rango

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

Resumen

Support