Обучай мало, выводи много: Эффективное по памяти обучение LoRA для больших языковых моделей

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

February 19, 2025

Авторы: Jun Zhang, Jue Wang, Huan Li, Lidan Shou, Ke Chen, Yang You, Guiming Xie, Xuejian Gong, Kunlong Zhou

cs.AI

Аннотация

Крупные языковые модели (LLM) значительно продвинули обработку естественного языка благодаря исключительной способности к обобщению задач. Низкоранговая адаптация (LoRA) предлагает экономичное решение для тонкой настройки, замораживая параметры исходной модели и обучая только легковесные низкоранговые адаптерные матрицы. Однако объем памяти, занимаемый LoRA, в основном определяется параметрами исходной модели. Чтобы смягчить эту проблему, мы предлагаем LoRAM — эффективную с точки зрения памяти схему обучения LoRA, основанную на интуиции, что многие нейроны в перепараметризованных LLM имеют низкую полезность для обучения, но необходимы для вывода. LoRAM предлагает уникальный подход: обучение проводится на урезанной (маленькой) модели для получения урезанных низкоранговых матриц, которые затем восстанавливаются и используются с исходной (большой) моделью для вывода. Дополнительно, минимально затратное непрерывное предварительное обучение, выполняемое издателями моделей заранее, устраняет расхождение в знаниях между урезанной и исходной моделями. Наши обширные эксперименты демонстрируют эффективность LoRAM при различных стратегиях урезания и последующих задачах. Для модели с 70 миллиардами параметров LoRAM позволяет проводить обучение на GPU с объемом памяти всего 20 ГБ HBM, заменяя GPU A100-80G для обучения LoRA и 15 GPU для полной тонкой настройки. В частности, QLoRAM, реализованный с использованием структурированного урезания в сочетании с 4-битной квантизацией для LLaMA-3.1-70B (LLaMA-2-70B), снижает затраты на хранение параметров, которые доминируют в использовании памяти при обучении низкоранговых матриц, в 15,81 раза (16,95 раза), при этом достигая значительного прироста производительности по сравнению как с исходной LLaMA-3.1-70B (LLaMA-2-70B), так и с LLaMA-3.1-8B (LLaMA-2-13B), обученной с использованием 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).