Train Small, Infer Large: Speichereffizientes LoRA-Training für große Sprachmodelle

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

February 19, 2025

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

cs.AI

Zusammenfassung

Große Sprachmodelle (LLMs) haben die natürliche Sprachverarbeitung durch außergewöhnliche Fähigkeiten zur Aufgabenverallgemeinerung erheblich vorangetrieben. Low-Rank Adaption (LoRA) bietet eine kosteneffiziente Lösung für das Feinabstimmen, indem die ursprünglichen Modellparameter eingefroren und nur leichte, niedrigrangige Adaptermatrizen trainiert werden. Allerdings wird der Speicherbedarf von LoRA weitgehend von den ursprünglichen Modellparametern dominiert. Um dies zu mildern, schlagen wir LoRAM vor, ein speichereffizientes LoRA-Trainingsschema, das auf der Intuition basiert, dass viele Neuronen in überparametrisierten LLMs einen geringen Trainingsnutzen haben, aber für die Inferenz unerlässlich sind. LoRAM bietet eine einzigartige Wendung: Es trainiert auf einem beschnittenen (kleinen) Modell, um beschnittene niedrigrangige Matrizen zu erhalten, die dann wiederhergestellt und mit dem ursprünglichen (großen) Modell für die Inferenz genutzt werden. Zusätzlich wird eine minimal kostspielige kontinuierliche Vorabschulung, die von den Modellherausgebern im Voraus durchgeführt wird, verwendet, um die Wissensdiskrepanz zwischen beschnittenen und ursprünglichen Modellen auszugleichen. Unsere umfangreichen Experimente demonstrieren die Wirksamkeit von LoRAM über verschiedene Beschnittstrategien und nachgelagerte Aufgaben hinweg. Für ein Modell mit 70 Milliarden Parametern ermöglicht LoRAM das Training auf einer GPU mit nur 20G HBM, wodurch eine A100-80G GPU für das LoRA-Training und 15 GPUs für das vollständige Feinabstimmen ersetzt werden. Insbesondere reduziert QLoRAM, das durch strukturiertes Beschneiden in Kombination mit 4-Bit-Quantisierung implementiert wird, für LLaMA-3.1-70B (LLaMA-2-70B) die Speicherkosten für die Parameter, die den Speicherverbrauch beim Training niedrigrangiger Matrizen dominieren, um das 15,81-fache (16,95-fache), während es dominante Leistungssteigerungen gegenüber dem ursprünglichen LLaMA-3.1-70B (LLaMA-2-70B) und dem LoRA-trainierten LLaMA-3.1-8B (LLaMA-2-13B) erzielt.

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