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

🔼 This figure compares the LoRA and LoRAM training and inference processes. LoRA uses the same original model for both training and inference, while LoRAM uses a pruned model for training and recovers the weights to use the original model for inference. The figure highlights the key stages: (a) LoRA training, updating low-rank matrices with original weights frozen; (b) LoRAM training, updating low-rank matrices on the pruned model; (c) LoRA inference, using the updated low-rank matrices with original weights; (d) LoRAM inference, using recovered low-rank matrices to integrate with original weights; (e) offline processing on the full-rank matrix; and (f) online generation of the low-rank matrix.

read the caption

Figure 2: Comparison of LoRAM and LoRA: Training (subfigures a and b) and Inference (c and d). Key stages include the offline process of the frozen full-rank matrix 𝐖0∗superscriptsubscript𝐖0\mathbf{W}_{0}^{*}bold_W start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT ∗ end_POSTSUPERSCRIPT (subfigure e) and the online generation of the learnable low-rank matrix 𝐖Δ∗superscriptsubscript𝐖Δ\mathbf{W}_{\Delta}^{*}bold_W start_POSTSUBSCRIPT roman_Δ end_POSTSUBSCRIPT start_POSTSUPERSCRIPT ∗ end_POSTSUPERSCRIPT (f) during LoRAM training (b) and inference (d).

🔼 This figure displays the test perplexity results for fine-tuning two different sizes of LLAMA language models (13B and 70B parameters) using the OpenHermes dataset. The perplexity is tracked over various training iterations for different training methods, including standard LoRA and several variations of the proposed LORAM method incorporating different pruning strategies (random, structured, semi-structured, and unstructured). This allows comparison of the training convergence speeds and the resulting test perplexity for these different training approaches, demonstrating the effectiveness of LORAM in improving efficiency without significant loss of performance.

read the caption

Figure 3: The test perplexity of training LLaMA-2-13B & LLaMA-2-70B on OpenHermes.

🔼 This figure displays the test perplexity results for training two large language models, LLaMA-2-13B and LLaMA-2-70B, on the OpenOrca dataset. The graphs show the perplexity over training iterations for each model, comparing standard LORA fine-tuning with several variants of the proposed memory-efficient LORA training scheme (LORAM) using different pruning strategies: LORAM-RAND, LORAM-STRU, LORAM-SEMI, and LORAM-UNST. The figure also includes results for a smaller LLaMA model (7B) trained with LORA to provide a performance baseline for comparison. This comparison allows assessment of the tradeoffs between model size, memory efficiency, and performance.

read the caption

Figure 4: The test perplexity of training LLaMA-2-13B & LLaMA-2-70B on OpenOrca.

🔼 This figure displays the results of training the LLaMA-3.1-70B model on the OpenHermes dataset using different training methods. It shows the test perplexity during training (how well the model predicts the next word in a sequence, lower is better) on the Alpaca dataset (an out-of-domain dataset; that is, a dataset different than the one used for training). Additionally, the figure presents downstream task performance results across various tasks, showing the effectiveness of different training methods on various downstream applications after training.

read the caption

Figure 5: The test perplexity & downstream performance of training LLaMA-3.1-70B on OpenHermes.

🔼 This figure demonstrates the impact of the recovery and alignment steps in the LORAM model training process on LLaMA-2-13B. It compares the performance of four different pruning strategies (Rand, Stru, Semi, Unst) with and without the recovery and alignment steps. The plots show the test perplexity on the Alpaca dataset for each strategy and configuration. This illustrates that the recovery and alignment steps significantly improve performance, especially under aggressive pruning rates. The results highlight that simply pruning and then fine-tuning is not enough to achieve high performance; the recovery and alignment steps are crucial in bridging the gap between the pruned model used for training and the full original model used for inference.

read the caption

Figure 6: Necessity of Recovery & Alignment across different pruning strategies on LLaMA-2-13B.

🔼 This figure demonstrates the impact of varying parameter reduction ratios on the performance of LoRA and QLORAM-STRU models. The x-axis represents the parameter reduction ratio, ranging from approximately 10x to nearly 30x. The y-axis shows the test perplexity on the Alpaca dataset. LoRA-trained and QLORAM-STRU-trained models show a trend of improved performance with increasing reduction ratios, but excessive pruning (larger reduction ratios) negatively impacts the model’s performance. A comparison with a naive pruning method demonstrates that LoRAM’s method leads to far more robust performance as the parameter reduction increases.

read the caption

Figure 7: Effect of scaling parameter reduction ratio.

🔼 This figure presents the performance of various downstream tasks (GSM8K, MathQA, CSR, HumanEval) for different parameter reduction ratios achieved by QLORAM-STRU on the LLaMA-2-70B model. It shows how the model’s performance on these tasks changes as the parameter reduction ratio increases from approximately 10x to nearly 30x. The graph illustrates a trade-off: increasing the reduction ratio initially improves performance, reaching an optimal point, before performance starts to decrease again at the most aggressive pruning ratios. This demonstrates the effectiveness of QLORAM-STRU in balancing memory efficiency and task performance up to a certain level of model compression, beyond which the level of pruning negatively affects performance.

read the caption

Figure 8: Performance of downstream tasks across different parameter reduction ratios.

🔼 This table presents the average accuracy scores achieved by different LLaMA-2 models on six common sense reasoning tasks. The models compared include the original, untrained LLaMA-2-13B, a smaller LoRA-trained LLaMA-2-7B, and various versions of the memory-efficient LORAM model with different pruning techniques (random, structured, semi-structured, and unstructured). The results show the average accuracy across the six tasks, along with the parameter reduction ratio for each LORAM variant compared to the original LLaMA-2-13B model. Appendix E provides a detailed breakdown of the performance on each individual common sense reasoning sub-task.

read the caption

Table 2: Average accuracy (%) of the CSR in the common sense reasoning domain (1-shot) under the LLaMA-2. Baseline results for each subtask of CSR are detailed in Appendix E.

🔼 This table presents the performance of different LLaMA-2 models on the HumanEval code generation benchmark. The HumanEval benchmark evaluates the ability of a language model to generate correct code given a function signature, a docstring (description of the function), and unit tests. The models tested include the original untuned LLaMA-2 model, smaller LoRA-tuned models, and various LoRAM-tuned models which use different pruning techniques. The results reported are PASS@1 (the percentage of tasks for which the top-ranked generated solution is correct) and PASS@10 (the percentage of tasks for which at least one of the top-10 generated solutions is correct). The best results from a range of temperature settings (controlling randomness in the model’s output) are shown for each model.

read the caption

Table 3: Pass@1(%) and Pass@10(%) of HumanEval in the code generation domain under LLaMA-2. The best results for all baselines are reported, selected from temperature settings in {0.0, 0.2, 0.4, 0.6, 0.8} with toppsubscripttopp\textsc{top}_{\textsc{p}}top start_POSTSUBSCRIPT p end_POSTSUBSCRIPT fixed at 0.95.

🔼 This table presents a comparison of four different pruning methods used in the LoRAM (Low-Rank Adaptation with Memory efficiency) training scheme. Specifically, it shows the impact of each pruning method on LLaMA-2-13B. The table lists the original number of parameters, the pruning ratio applied, the resulting number of pruned parameters, the parameter reduction ratio achieved (expressed as a multiple), and the amount of High Bandwidth Memory (HBM) required to store the pruned parameters. Note that the low-rank matrix overhead isn’t included in this HBM calculation.

read the caption

Table 4: LoRAM configures on LLaMA-2-13B. Comparison of different pruning methods in terms of parameter reduction ratio (Reduction) and HBM footprint (GB) of pruned parameters (HBM), ignoring low-rank matrix overhead.

🔼 This table presents the configurations used for the LoRAM experiments on the LLaMA-2-70B and LLaMA-3.1-70B models. It shows how different pruning ratios affect the number of pruned parameters, the reduction ratio compared to the original number of parameters, and the amount of GPU High Bandwidth Memory (HBM) required for training. The table is crucial for understanding the impact of different pruning strategies on memory efficiency and the trade-off between memory savings and performance.

read the caption

Table 5: LoRAM configures on LLaMA-2-70B and LLaMA-3.1-70B with different pruning ratios.

🔼 This table shows different configurations for QLoRAM (Quantized LoRAM) on two large language models, LLaMA-2-70B and LLaMA-3.1-70B. It details the original number of parameters, the pruning ratio used, the number of parameters after pruning, the reduction ratio achieved (showing how much smaller the model became), and the amount of High Bandwidth Memory (HBM) used by the pruned parameters. The table highlights how QLoRAM, by combining LORAM with quantization, achieves very aggressive parameter compression, significantly reducing memory requirements during training.

read the caption

Table 6: QLoRAM configures on LLaMA-2-70B and LLaMA-3.1-70B with , demonstrating more aggressive parameter compression.

🔼 This table presents the results of evaluating LoRAM’s performance on the GSM8K dataset, a benchmark specifically designed for mathematical reasoning. It shows the accuracy achieved by different model configurations, including the original untuned model (70B w/o Fine-Tuning), a smaller model fine-tuned with LoRA (8B LORA), and several LoRAM variants (70B QLORAM-Stru) with varying degrees of parameter reduction. The accuracy is shown as a percentage, and the parameter reduction ratio indicates how much smaller the LoRAM models are compared to the 70B parameter model. The results demonstrate LoRAM’s performance on a domain-specific task and how it compares to alternative fine-tuning methods, highlighting its effectiveness in scenarios with limited resources.

read the caption

Table 7: Evaluation of LoRAM on the GSM8K dataset for domain-specific fine-tuning. Results show accuracy (%) and parameter reduction ratios for different configurations.

🔼 This table compares the peak memory usage, latency, and throughput of different models during the online training phase. The models compared are LoRA-trained 7B parameter model and LoRAM-trained 13B parameter model. The training workload consists of 1024 samples with a batch size of 128, micro-batch size of 4, and a sequence length of 512. The results show that LoRAM achieves comparable performance to the smaller 7B parameter LoRA model in terms of memory, latency and throughput, despite using a larger model.

read the caption

Table 8: Comparison of peak memory (MiB), latency (s), and throughput (samples/s) during the online training phase for LoRAM and LoRA models. Results are based on a workload of 1024 samples (batch size 128, micro-batch size 4, sequence length 512).