Train Small, Infer Large: Geheugenefficiënte LoRA-training voor grote taalmodelle

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

February 19, 2025

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

cs.AI

Samenvatting

Grote Taalmodellen (LLMs) hebben natuurlijke taalverwerking aanzienlijk vooruitgebracht met uitzonderlijke mogelijkheden voor taakgeneralisatie. Low-Rank Adaption (LoRA) biedt een kosteneffectieve oplossing voor fine-tuning, waarbij de oorspronkelijke modelparameters worden bevroren en alleen lichtgewicht, laag-rang adaptermatrices worden getraind. Het geheugengebruik van LoRA wordt echter grotendeels bepaald door de oorspronkelijke modelparameters. Om dit te verlichten, stellen we LoRAM voor, een geheugenefficiënt LoRA-trainingsschema dat gebaseerd is op de intuïtie dat veel neuronen in overgeparameteriseerde LLMs een lage trainingsnut hebben, maar essentieel zijn voor inferentie. LoRAM introduceert een unieke draai: het traint op een gepruned (klein) model om geprunede laag-rang matrices te verkrijgen, die vervolgens worden hersteld en gebruikt met het oorspronkelijke (grote) model voor inferentie. Daarnaast minimaliseert een minimaal kostbare voortdurende pre-training, uitgevoerd door de modeluitgevers van tevoren, de kennisdiscrepantie tussen geprunede en oorspronkelijke modellen. Onze uitgebreide experimenten tonen de effectiviteit van LoRAM aan over verschillende pruningstrategieën en downstreamtaken. Voor een model met 70 miljard parameters maakt LoRAM training mogelijk op een GPU met slechts 20G HBM, waardoor een A100-80G GPU voor LoRA-training en 15 GPU's voor volledige fine-tuning worden vervangen. Specifiek reduceert QLoRAM, geïmplementeerd door gestructureerde pruning gecombineerd met 4-bit kwantisatie, voor LLaMA-3.1-70B (LLaMA-2-70B), de parameteropslagkosten die het geheugengebruik domineren in laag-rang matrixtraining met 15,81 keer (16,95 keer), terwijl het dominante prestatieverbeteringen behaalt ten opzichte van zowel de oorspronkelijke LLaMA-3.1-70B (LLaMA-2-70B) als LoRA-getrainde LLaMA-3.1-8B (LLaMA-2-13B).

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).