Paper: LoRA-FA: Memory-efficient Low-rank Adaptation for Large Language Models Fine-tuning (ICLR-2024)

Reviewer comment (Writing):
The writing of the paper at times leans towards over-claiming or making unsuitable comparisons to bolster their method. For example, in section 1 and in other spots, the memory of LoRA-FA is compared to full fine-tuning but not LoRA. Given that LoRA-FA is fundamentally a tweak of LoRA, the natural comparison should be to LoRA, not the full fine-tuning just to make the numbers look more impressive. To the authors: the benefits of LoRA-FA are moderate but clear; there is no need to oversell the method.

Author response:
**Response**: Thank you for your insightful comment, we will refactor the paper to make it more impartial.

Reviewer comment (Experiments):
To determine if LoRA-FA underperforms LoRA (or not), there needs to be a setting where LoRA is "capacity-constrained", in order to assess if LoRA-FA has even less "capacity" than LoRA. This requires a setting where LoRA meaningfully underperforms full fine-tuning; however, for both the RoBERTa and T5 experiments, this is not the case.

Author response:
**Response**: Thank you for your kind guidance. To address this problem, we conduct the  Llama2-7B fine-tuning on Super-NaturalInstructions dataset. Super-NaturalInstructions is a challenging dataset for LLMs. Specifically, we use the task001_quoref_question_generation.json and task002_quoref_answer_generation.json as the training set, and set the sequence length to 2048, batch size to 4, train epochs to 5. We report eval_loss of 3 learning rates: [7e-5, 1e-4, 4e-4] as follows.

| Learning rate | 7e-5 | 1e-4 | 4e-4 |
|---------------|------|------|------|
| LoRA          | 0.45 | 0.53 | 4.27 |
| LoRA-FA       | 0.45 | 0.47 | 0.51 |

In the results, both LoRA and LoRA-FA can achieve similar eval_loss. When $lr=4e-4$, LoRA didn't perform well because the learning rate is too large for LoRA. Overall, LoRA-FA shows strong capability similar to LoRA when fine-tuning on generalization tasks.

Reviewer comment (Evaluation):
The LLaMA-7B experiments on Alpaca/FLAN -> MMLU come the closest to showing performance differences, but the margin is still too small to tell (and essentially no gap at all in the case of Alpaca). More broadly, MMLU is not a fine measure of LM performance.

Author response:
**Response**: Thank you for your kind guidance. To address this problem, we conduct the  Llama2-7B fine-tuning on Super-NaturalInstructions dataset. Super-NaturalInstructions is a challenging dataset for LLMs. Specifically, we use the task001_quoref_question_generation.json and task002_quoref_answer_generation.json as the training set, and set the sequence length to 2048, batch size to 4, train epochs to 5. We report eval_loss of 3 learning rates: [7e-5, 1e-4, 4e-4] as follows.

| Learning rate | 7e-5 | 1e-4 | 4e-4 |
|---------------|------|------|------|
| LoRA          | 0.45 | 0.53 | 4.27 |
| LoRA-FA       | 0.45 | 0.47 | 0.51 |

In the results, both LoRA and LoRA-FA can achieve similar eval_loss. When $lr=4e-4$, LoRA didn't perform well because the learning rate is too large for LoRA. Overall, LoRA-FA shows strong capability similar to LoRA when fine-tuning on generalization tasks.

Reviewer comment (Experiments):
It is recommended that the authors run a set of experiments on something like Super-NaturalInstructions, where there is likely to be a bigger performance gap, since the evaluation is performed on generated sequences rather than simple multiple-choice knowledge questions.

Author response:
**Response**: Thank you for your kind guidance. To address this problem, we conduct the  Llama2-7B fine-tuning on Super-NaturalInstructions dataset. Super-NaturalInstructions is a challenging dataset for LLMs. Specifically, we use the task001_quoref_question_generation.json and task002_quoref_answer_generation.json as the training set, and set the sequence length to 2048, batch size to 4, train epochs to 5. We report eval_loss of 3 learning rates: [7e-5, 1e-4, 4e-4] as follows.

| Learning rate | 7e-5 | 1e-4 | 4e-4 |
|---------------|------|------|------|
| LoRA          | 0.45 | 0.53 | 4.27 |
| LoRA-FA       | 0.45 | 0.47 | 0.51 |

In the results, both LoRA and LoRA-FA can achieve similar eval_loss. When $lr=4e-4$, LoRA didn't perform well because the learning rate is too large for LoRA. Overall, LoRA-FA shows strong capability similar to LoRA when fine-tuning on generalization tasks.

Reviewer comment (Evaluation):
The discussion on the benefits of reduced activation memory through LoRA-FA is lacking. In the finetuning phase, unlike the inference phase, both sequence length and batch size are longer, resulting in high GPU utilization. Therefore, it is not considered practically significant to make LoRA-FA more memory efficient than the existing LoRA.

Author response:
Thank you for pointing out the issue! We have corrected the results in Table 5, it was because activation recomputation was turned on when conducting memory benchmark experiments. We ensure all of the configures are right and re-conducted the experiment. Specifically, we fine-tuning Llama2-7B with sequence length = 1024. In the new results, we have several findings:

1. LoRA-FA reduces ~13GB memory compared to LoRA in fine-tuning Llama2-7B when batch size is 16.
2. In practice, when fine-tuning Llama2-7B on various hardware, LoRA-FA has better capability than LoRA. For example, on RTX4090-24G, LoRA can only fine-tuning Llama2-7B when $batch size \leq 2$, while LoRA-FA can enlarge the batch size to 4; on A100-40G, LoRA works with batch size = 4, while LoRA-FA works with batch size = 8.

We also give a runtime memory breakdown analysis of fine-tuning Llama2-7B as follows when batch size is 16, to figure out the memory usage percentage of different modules, such as optimizer state and activation memory.

|                    | LoRA        |                | LoRA-FA     |                |
|:------------------:|:-----------:|:--------------:|:-----------:|:--------------:|
|                    | Memory (GB) | Percentage (%) | Memory (GB) | Percentage (%) |
| Hardware researved | 0.7         | 0.9            | 0.7         | 1.1            |
| Pytorch researved  | 0.5         | 0.6            | 0.5         | 0.8            |
| Model in bf16      | 12.6        | 17             | 12.6        | 20.3           |
| LoRA module        | 0.8         | 1.1            | 0.8         | 1.4            |
| Cache              | 5.1         | 6.7            | 4.8         | 7.7            |
| Optimizer state    | 0.7         | 0.9            | 0.4         | 0.4            |
| Activation         | 55.6        | 73.1           | 42.3        | 68.2           |

This result shows that LoRA-FA can significantly reduce the activation memory usage compared to LoRA when fine-tuning Llama2-7B (from 55.6 GB to 42.3 GB). This advantage can either help fine-tuning on consumer level GPUs, or can achieve higher throughput by enlarging batch size.

Reviewer comment (Experiments):
Is there any performance difference based on the initialization method for the LoRA adaptation layer A?

Author response:
We discussed the performance difference of three different initialization methods for the LoRA adaptation layer A in Appendix B.3.1.

Reviewer comment (Experiments):
When integrating the proposed LoRA-FA into QLoRA, how does it impact the performance of CSR or MMLU in models such as LLaMA, RoBERTa, and T5? While Section 4 of the experiment results demonstrates the approach of freezing adaption layer A in various LLMs' LoRA, including QLoRA, there seems to be a mention of potential application without concrete results on the performance after actual finetuning.

Author response:
The fine-tuning with QLoRA is still on-going, we will release the result ASAP.

Reviewer comment (Evaluation):
When looking at Figure 4 in the Appendix, it appears that LoRA and LoRA-FA exhibit a similar trend in TFLOPS. What are the benefits gained during the Finetuning phase by actually reducing activation memory?

Author response:
Thank you for pointing out the issue! We have corrected the results in Table 5, it was because activation recomputation was turned on when conducting memory benchmark experiments. We ensure all of the configures are right and re-conducted the experiment. Specifically, we fine-tuning Llama2-7B with sequence length = 1024. In the new results, we have several findings:

1. LoRA-FA reduces ~13GB memory compared to LoRA in fine-tuning Llama2-7B when batch size is 16.
2. In practice, when fine-tuning Llama2-7B on various hardware, LoRA-FA has better capability than LoRA. For example, on RTX4090-24G, LoRA can only fine-tuning Llama2-7B when $batch size \leq 2$, while LoRA-FA can enlarge the batch size to 4; on A100-40G, LoRA works with batch size = 4, while LoRA-FA works with batch size = 8.

We also give a runtime memory breakdown analysis of fine-tuning Llama2-7B as follows when batch size is 16, to figure out the memory usage percentage of different modules, such as optimizer state and activation memory.

|                    | LoRA        |                | LoRA-FA     |                |
|:------------------:|:-----------:|:--------------:|:-----------:|:--------------:|
|                    | Memory (GB) | Percentage (%) | Memory (GB) | Percentage (%) |
| Hardware researved | 0.7         | 0.9            | 0.7         | 1.1            |
| Pytorch researved  | 0.5         | 0.6            | 0.5         | 0.8            |
| Model in bf16      | 12.6        | 17             | 12.6        | 20.3           |
| LoRA module        | 0.8         | 1.1            | 0.8         | 1.4            |
| Cache              | 5.1         | 6.7            | 4.8         | 7.7            |
| Optimizer state    | 0.7         | 0.9            | 0.4         | 0.4            |
| Activation         | 55.6        | 73.1           | 42.3        | 68.2           |

This result shows that LoRA-FA can significantly reduce the activation memory usage compared to LoRA when fine-tuning Llama2-7B (from 55.6 GB to 42.3 GB). This advantage can either help fine-tuning on consumer level GPUs, or can achieve higher throughput by enlarging batch size.

Reviewer comment (Novelty):
In the context of "reducing the number of trainable parameters", as mentioned in the LoRA paper, the previously proposed PEFT method significantly reduced the number of trainable parameters. This led to a drastic reduction in the memory usage of the optimizer state, bringing it down to megabytes. Thus, using fewer trainable parameters than LoRA does not yield a significant difference.

Author response:
Thank you for your insightful comment. It is true that using fewer trainable parameters than LoRA does not often yield a significant difference in LoRA-like methods. Our intention was not to imply that using fewer trainable parameters than LoRA would reduce more GPU memory. Rather, our emphasis was to show that LoRA-FA reduces great activation memory by freezing the update of low rank adapter A, which can also reduce the trainable parameters. We will ensure to clarify this in our revised manuscript to avoid any misunderstanding.

Reviewer comment (Evaluation):
Regarding "reducing the activation memory", while LLaMA uses a max sequence length of 2k, the LLaMA2 model employs a 4k length. It's evident that models are gravitating towards longer sequence lengths. As per Table 1, LoRA-FA utilizes '2bsr' of memory (compared to LoRA's '2bsd+2bsr'), but the advantage in memory savings during training becomes less pronounced as the model's sequence length increases relative to LoRA.

Author response:
This is a very detailed comment. Regarding to Table 1, in a linear layer, LoRA-FA saves `2bsd` usage compared to LoRA in activation memory, which is linearly related to the coefficient sequence length. This means the memory saving has a linear-like scaling efficiency regarding to sequence length. We further conduct an experiment with longer sequence length in fine-tuning Llama2-13B. Specifically, we set the sequence length from 1024 to 2048, and the batch size to 2. We report the peak GPU memory usage as follows.

| seq_length | 1024  | 2048  |
|------------|-------|-------|
| LoRA       | 50 GB | 71 GB |
| LoRA-FA    | 45 GB | 62 GB |

This result shows that LoRA-FA saves more GPU memory when using longer sequence length.

Reviewer comment (Evaluation):
The proposed methodology primarily targets the reduction of memory usage during training. In section 4.2, Table 5, the variance in memory peak between LoRA and LoRA-FA is minimal, especially for generative models like LLaMA-7B.

Author response:
Thank you for pointing out the issue! We have corrected the results in Table 5, it was because activation recomputation was turned on when conducting memory benchmark experiments. We ensure all of the configures are right and re-conducted the experiment. Specifically, we fine-tuning Llama2-7B with sequence length = 1024. In the new results, we have several findings:

1. LoRA-FA reduces ~13GB memory compared to LoRA in fine-tuning Llama2-7B when batch size is 16.
2. In practice, when fine-tuning Llama2-7B on various hardware, LoRA-FA has better capability than LoRA. For example, on RTX4090-24G, LoRA can only fine-tuning Llama2-7B when $batch size \leq 2$, while LoRA-FA can enlarge the batch size to 4; on A100-40G, LoRA works with batch size = 4, while LoRA-FA works with batch size = 8.

We also give a runtime memory breakdown analysis of fine-tuning Llama2-7B as follows when batch size is 16, to figure out the memory usage percentage of different modules, such as optimizer state and activation memory.

|                    | LoRA        |                | LoRA-FA     |                |
|:------------------:|:-----------:|:--------------:|:-----------:|:--------------:|
|                    | Memory (GB) | Percentage (%) | Memory (GB) | Percentage (%) |
| Hardware researved | 0.7         | 0.9            | 0.7         | 1.1            |
| Pytorch researved  | 0.5         | 0.6            | 0.5         | 0.8            |
| Model in bf16      | 12.6        | 17             | 12.6        | 20.3           |
| LoRA module        | 0.8         | 1.1            | 0.8         | 1.4            |
| Cache              | 5.1         | 6.7            | 4.8         | 7.7            |
| Optimizer state    | 0.7         | 0.9            | 0.4         | 0.4            |
| Activation         | 55.6        | 73.1           | 42.3        | 68.2           |

This result shows that LoRA-FA can significantly reduce the activation memory usage compared to LoRA when fine-tuning Llama2-7B (from 55.6 GB to 42.3 GB). This advantage can either help fine-tuning on consumer level GPUs, or can achieve higher throughput by enlarging batch size.

Reviewer comment (Experiments):
As stated in the LoRA paper, amplifying the trainable parameters doesn't notably affect accuracy (refer to Figure 2). For a more precise evaluation, the reviewer suggests the following comparative experiments with LoRA:
- Match LoRA and LoRA-FA at the same level of trainable parameters and then compare their accuracy and memory peaks.
- Adjust both LoRA and LoRA-FA to similar peak memory levels and compare their MMLU accuracy.
- Examine the peak memory usage of both LoRA and LoRA-FA when using generative models larger than LLaMA-7B to determine if the gap increases as the model size grows.

Author response:
- We set the rank of LoRA to 64, and the rank of LoRA-FA to 128, to maintain the same level of trainable parameters. We fine-tune Llama2-7B on Super-NaturalInstructions dataset with sequence length = 1024. We report the eval_loss of both approaches. From our result, both LoRA and LoRA-FA can achieve a similar performance (eval_loss = 0.45). LoRA-FA saves 13GB of peak GPU memory compared to LoRA. This means that enlarging rank doesn't yield a great difference in memory consumption.
- Since enlarging rank doesn't yield a great difference in memory consumption, we choose to adjust the batch size of both methods to maintain a similar peak memory levels. From our result, both LoRA and LoRA-FA achieves a similar MMLU accuracy when fine-tuning Llama2-7B on Super-NaturalInstructions dataset (accuracy = ~46.5).
- We have checked the peak memory usage when fine-tuning Llama2-7B and Llama2-13B using batch size 4, the results didn't show much increase in GPU memory gap as the model size grows.