Train Small, Infer Large: Memory-Efficient LoRA Training for Large
  Language Models
Authors: name: Jun Zhang, name: Jue Wang, name: Huan Li, name: Lidan Shou, name: Ke Chen, name: Yang You, name: Guiming Xie, name: Xuejian Gong, name: Kunlong Zhou

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