Treine Pequeno, Infira Grande: Treinamento Eficiente em Memória com LoRA para Modelos de Linguagem de Grande 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

Resumo

Os Modelos de Linguagem de Grande Escala (LLMs) avançaram significativamente o processamento de linguagem natural com capacidades excepcionais de generalização de tarefas. A Adaptação de Baixa Classificação (LoRA) oferece uma solução de ajuste fino econômica, congelando os parâmetros originais do modelo e treinando apenas matrizes adaptadoras leves e de baixa classificação. No entanto, a pegada de memória do LoRA é amplamente dominada pelos parâmetros originais do modelo. Para mitigar isso, propomos o LoRAM, um esquema de treinamento LoRA eficiente em memória, baseado na intuição de que muitos neurônios em LLMs superparametrizados têm baixa utilidade de treinamento, mas são essenciais para inferência. O LoRAM apresenta uma abordagem única: ele treina em um modelo podado (pequeno) para obter matrizes de baixa classificação podadas, que são então recuperadas e utilizadas com o modelo original (grande) para inferência. Além disso, um pré-treinamento contínuo de custo mínimo, realizado antecipadamente pelos publicadores do modelo, alinha a discrepância de conhecimento entre os modelos podados e originais. Nossos extensos experimentos demonstram a eficácia do LoRAM em várias estratégias de poda e tarefas subsequentes. Para um modelo com 70 bilhões de parâmetros, o LoRAM permite o treinamento em uma GPU com apenas 20G de HBM, substituindo uma GPU A100-80G para treinamento LoRA e 15 GPUs para ajuste fino completo. Especificamente, o QLoRAM, implementado por poda estruturada combinada com quantização de 4 bits, para o LLaMA-3.1-70B (LLaMA-2-70B), reduz o custo de armazenamento de parâmetros que domina o uso de memória no treinamento de matrizes de baixa classificação em 15,81 vezes (16,95 vezes), enquanto alcança ganhos de desempenho dominantes sobre o LLaMA-3.1-70B (LLaMA-2-70B) original e o LLaMA-3.1-8B (LLaMA-2-13B) treinado com 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).