Addestramento Piccolo, Inferenza Grande: Addestramento Efficiente in Memoria con LoRA per Modelli Linguistici di Grandi Dimensioni

Train Small, Infer Large: Memory-Efficient LoRA Training for Large Language Models

February 19, 2025

Autori: Jun Zhang, Jue Wang, Huan Li, Lidan Shou, Ke Chen, Yang You, Guiming Xie, Xuejian Gong, Kunlong Zhou

cs.AI

Abstract

I Large Language Model (LLM) hanno fatto avanzare significativamente l'elaborazione del linguaggio naturale con capacità eccezionali di generalizzazione dei compiti. La Low-Rank Adaption (LoRA) offre una soluzione di fine-tuning economicamente vantaggiosa, congelando i parametri originali del modello e addestrando solo matrici adattatrici leggere e a basso rango. Tuttavia, l'impronta di memoria di LoRA è largamente dominata dai parametri originali del modello. Per mitigare questo problema, proponiamo LoRAM, uno schema di addestramento LoRAM efficiente in termini di memoria, basato sull'intuizione che molti neuroni nei LLM sovra-parametrizzati hanno una bassa utilità durante l'addestramento ma sono essenziali per l'inferenza. LoRAM introduce un approccio unico: addestra su un modello potato (di dimensioni ridotte) per ottenere matrici a basso rango potate, che vengono poi recuperate e utilizzate con il modello originale (di grandi dimensioni) per l'inferenza. Inoltre, un pre-addestramento continuo a costo minimo, eseguito in anticipo dagli editori del modello, allinea la discrepanza di conoscenza tra i modelli potati e quelli originali. I nostri ampi esperimenti dimostrano l'efficacia di LoRAM attraverso varie strategie di pruning e compiti downstream. Per un modello con 70 miliardi di parametri, LoRAM consente l'addestramento su una GPU con solo 20G di HBM, sostituendo una GPU A100-80G per l'addestramento LoRA e 15 GPU per il fine-tuning completo. In particolare, QLoRAM implementato tramite pruning strutturato combinato con quantizzazione a 4 bit, per LLaMA-3.1-70B (LLaMA-2-70B), riduce il costo di archiviazione dei parametri che domina l'uso della memoria nell'addestramento di matrici a basso rango di 15,81 volte (16,95 volte), ottenendo al contempo guadagni di prestazioni dominanti rispetto sia al LLaMA-3.1-70B (LLaMA-2-70B) originale che al LLaMA-3.1-8B (LLaMA-2-13B) addestrato 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).