Entrena Pequeño, Infiere Grande: Entrenamiento Eficiente en Memoria con LoRA para Modelos de Lenguaje a Gran Escala
Train Small, Infer Large: Memory-Efficient LoRA Training for Large Language Models
February 19, 2025
Autores: Jun Zhang, Jue Wang, Huan Li, Lidan Shou, Ke Chen, Yang You, Guiming Xie, Xuejian Gong, Kunlong Zhou
cs.AI
Resumen
Los Modelos de Lenguaje de Gran Escala (LLMs, por sus siglas en inglés) han avanzado significativamente en el procesamiento del lenguaje natural, demostrando capacidades excepcionales de generalización de tareas. La Adaptación de Bajo Rango (LoRA, por sus siglas en inglés) ofrece una solución de ajuste fino rentable, congelando los parámetros originales del modelo y entrenando únicamente matrices adaptadoras ligeras de bajo rango. Sin embargo, la huella de memoria de LoRA está dominada en gran medida por los parámetros del modelo original. Para mitigar esto, proponemos LoRAM, un esquema de entrenamiento de LoRA eficiente en memoria basado en la intuición de que muchas neuronas en LLMs sobreparametrizados tienen una utilidad de entrenamiento baja pero son esenciales para la inferencia. LoRAM presenta un enfoque único: entrena en un modelo podado (pequeño) para obtener matrices de bajo rango podadas, las cuales luego se recuperan y utilizan con el modelo original (grande) para la inferencia. Además, un preentrenamiento continuo de bajo costo, realizado de antemano por los publicadores del modelo, alinea la discrepancia de conocimiento entre los modelos podados y originales. Nuestros extensos experimentos demuestran la eficacia de LoRAM en diversas estrategias de poda y tareas posteriores. Para un modelo con 70 mil millones de parámetros, LoRAM permite el entrenamiento en una GPU con solo 20G de HBM, reemplazando una GPU A100-80G para el entrenamiento de LoRA y 15 GPUs para el ajuste fino completo. Específicamente, QLoRAM implementado mediante poda estructurada combinada con cuantización de 4 bits, para LLaMA-3.1-70B (LLaMA-2-70B), reduce el costo de almacenamiento de parámetros que domina el uso de memoria en el entrenamiento de matrices de bajo rango en 15.81 veces (16.95 veces), logrando además ganancias de rendimiento dominantes sobre tanto el LLaMA-3.1-70B (LLaMA-2-70B) original como el LLaMA-3.1-8B (LLaMA-2-13B) entrenado con LoRA.
English
Large Language Models (LLMs) have significantly advanced natural language
processing with exceptional task generalization capabilities. Low-Rank Adaption
(LoRA) offers a cost-effective fine-tuning solution, freezing the original
model parameters and training only lightweight, low-rank adapter matrices.
However, the memory footprint of LoRA is largely dominated by the original
model parameters. To mitigate this, we propose LoRAM, a memory-efficient LoRA
training scheme founded on the intuition that many neurons in
over-parameterized LLMs have low training utility but are essential for
inference. LoRAM presents a unique twist: it trains on a pruned (small) model
to obtain pruned low-rank matrices, which are then recovered and utilized with
the original (large) model for inference. Additionally, minimal-cost continual
pre-training, performed by the model publishers in advance, aligns the
knowledge discrepancy between pruned and original models. Our extensive
experiments demonstrate the efficacy of LoRAM across various pruning strategies
and downstream tasks. For a model with 70 billion parameters, LoRAM enables
training on a GPU with only 20G HBM, replacing an A100-80G GPU for LoRA
training and 15 GPUs for full fine-tuning. Specifically, QLoRAM implemented by
structured pruning combined with 4-bit quantization, for LLaMA-3.1-70B
(LLaMA-2-70B), reduces the parameter storage cost that dominates the memory
usage in low-rank matrix training by 15.81times (16.95times), while
achieving dominant performance gains over both the original LLaMA-3.1-70B
(LLaMA-2-70B) and LoRA-trained LLaMA-3.1-8B (LLaMA-2-13B).