Paper: Towards Green AI in Fine-tuning Large Language Models via Adaptive Backpropagation (ICLR-2024)

Reviewer comment (Presentation):
The template is highly altered to fit in more content. For example, the author section is removed, which may violate the policy.

Author response:
Thank you for the comment. We have revised the paper to add the author section and make sure that it comply to the submission instructions.

Reviewer comment (Novelty):
The paper misses some references to existing work that uses similar techniques. For example, [a] also uses sparse update for the model and a similar optimization goal as Equation 5 to get the tensor selection scheme, which also considers the backward cost for weights and activations. Other related works include [b, c]. The authors should discuss and differentiate from existing works.

Author response:
The main goal of MCUNetV3 [a] is to reduce the memory cost of training, so that the training can be executed on MCUs with extremely low memory. This goal is orthogonal to that of GreenTrainer, which is to speed up LLM fine-tuning. The methods of importance evaluation in [a] are also different from our method in GreenTrainer. From its problem formulation (see Eq. (5) in [a]), it adopts accuracy contributions as the importance metric and hence requires enumeration on the training dataset. This is acceptable for [a]’s problem scenario where both importance evaluation and trainable component selection are done offline. In contrast, in GreenTrainer we aim to enable runtime selection of trainable tensors, and hence we designed methods of tensor importance evaluation and tensor selection with very low computing costs for efficient runtime operations.

Similarly, SPT [b] also does offline selection. It builds on top of the existing Parameter-Efficient Fine-Tuning (PEFT) methods (e.g., LoRA) and always selects the top-k most important parameters in each adapter for fine-tuning. Since existing PEFT methods typically insert adapters to every LLM block, SPT still requires propagating activation gradients throughout the LLM, which cannot guarantee that the given objective of FLOPs reduction can be met.

TinyTrain [c] handcrafts a layer importance metric to combine accuracy, parameter counts, and computation cost into a single formula. It then offline selects the top-k important layers for fine-tuning online. Although it considers both the computing cost and accuracy, simply selecting top-k important layers cannot guarantee reduction of computing cost, because its importance metric doesn’t consider the computation dependency between layers in backpropagation. In comparison, GreenTrainer rigorously models the backpropagation cost for training each tensor and use DP for tensor selection. Our proposed importance metric also ensures additivity among importances of individual tensors, to ensure proper tensor selection.

Reviewer comment (Experiments):
There are some other efficient training methods focusing more on FLOPs reduction instead of parameter reduction like [d]. The authors may also want to include them as baselines since LoRA and Prefix-tuning are not designed for computation efficiency.

Author response:
Head2Toe [d] is designed for visual classification models, by inserting untrained classification heads for classification models such as ViT and ResNet50. However, when being applied to LLMs, due to the dependency between input and output embeddings of LLMs, it would be difficult to similarly insert an untrained output embedding layer for generative LLMs.

Reviewer comment (Experiments):
Despite the effective FLOPs reduction, I would expect the method to use more GPU memory compared to methods like LoRA due to the larger optimizer states (since more weights are updated). The authors may also discuss the memory usage compared to the baseline methods, and see whether it can be combined with LoRA to achieve reduced memory and FLOPs at the same time (e.g., only add LoRA to updated tensors).

Author response:
GreenTrainer could consume more memory than LoRA, because GreenTrainer is not optimized for reducing the memory usage. However, as the reviewer suggested, GreenTrainer can build upon LoRA and be used together with LoRA. In that case, the whole set of trainable components is the LoRA adapters being inserted in the LLM blocks. GreenTrainer then can be directly applied to select which LoRA adapter to fine-tune.

The extra memory usage in GreenTrainer is mainly due to the tensor importance evaluation that requires caching all the model weights and current gradients. In Appendix A.1, we proposed techniques to further reduce such memory usage, by applying early stop strategies and excluding tensors in early layers from being involved in tensor selection. Results in Table 6 showed that we can further reduce the memory usage by 50%, approximately to be half of that in full-training.

Reviewer comment (Experiments):
The tensor importance estimation is based on the gradient information. Is it consistent throughout the training, or will the importance change?

Author response:
The importance of tensors will change over training. We hypothesize the change of tensor importance stems from the irregular trajectory of how the model parameters move on the loss surface during training. As a result, we perform tensor importance evaluation per epoch to capture such dynamics. 

In addition, our proposed tensor importance evaluation, by design, incurs only negligible computing overhead, because it only takes one extra training iteration to obtain gradients for all model parameters. To further verify such low overhead, we conducted extra experiments and the results show that, even when we repeatedly perform tensor importance evaluation in each training epoch, its overhead is only <0.2% of the entire fine-tuning cost.

Reviewer comment (Novelty):
The novelty of the proposed method is rather limited. In particular, the method for evaluating the tensor importance has already been proposed for neural network pruning. And the authors did not cite any relevant literature here.

Author response:
We agree that some existing metrics (e.g., SNIP and Fisher information as suggested by Reviewer B8AP), could be adopted into GreenTrainer for tensor importance evaluation. However, these metrics, as designed for neural network pruning, are used to quantify the importance of model weights at their current values, so the less important weights can be set to zeros during inference. In contrast, our proposed tensor importance evaluation is to quantify the importance of performing weight updates on a tensor to the reduction of training loss. In other words, we measure the importance of “tensor update” instead of the importance of the tensor’s current value itself. Therefore, most existing importance metrics being used in neural network pruning would be ineffective in our training scenario.

In addition, additivity is an essential requirement on the importance metric in our problem formulation. However, existing metrics, such as SNIP ($|\partial L/\partial w \cdot w|$) [1] and Fisher information ($|\partial L/\partial w \cdot w|^2$) [2], take absolute or square values in the metric, and they are hence not additive, because updating parameters can be either constructive or destructive to the loss reduction. In contrast, our proposed importance metric is directly derived from the first-order approximation after Taylor expansion of the training loss. Therefore, it retains a moderate level of additivity in low learning rate ($10^{-5}$) regimes for fine-tuning [3].

[1] https://arxiv.org/abs/1810.02340
[2] https://arxiv.org/abs/2108.00708
[3] https://dl.acm.org/doi/10.1145/3581791.3596852

We are also aware of other choices of designing such importance metric while retaining additivity, and we have provided experiment results in Table 4 to demonstrate the performance of GT using different tensor importance metrics, such as the magnitude-based metric ($\Delta_w$) and gradients-only metric ($\partial L/\partial w$).

Reviewer comment (Writing):
The writing of the paper is sometimes confusing. For example, in section 3.1, it is not clear how such profiling is done for different layers in LLM. The authors should give some concrete examples of the FLOPs for these layers to better demonstrate the idea of tensor FLOPs profiling.

Author response:
We provided technical details about how to obtain tensor FLOPs for different types of layers (e.g., MHA, LayerNorm, and FFN) in section 3.1. In the revised paper draft, we also added an example about practical profiling in our experiments in Appendix A.3 to help understanding. Please check.

Reviewer comment (Theory):
Since the proposed dynamic programming can only find approximate solutions, there is no discussions on the accuracy of the found solution using the proposed approach.

Author response:
DP algorithm can find the optimal solution of tensor selection. However, in order to reduce its computing cost, we downscale $T_{full}$ to $T_q$  in Appendix A.2, which introduces such sub-optimality. We studied the impact of such downscaling in Appendix A.2, and results show that we can always find an value of $T_q$ that achieves on-par model accuracy but restrains the computing cost of DP within 1% of the entire fine-tuning cost.

Reviewer comment (Experiments):
In the experiments, the authors only fine-tune the model for 5 epochs, then it is possible for some fine-tuning methods (eg. FT-Top2) to have the underfitting issue, how about fine-tuning for more epochs?

Author response:
We agree that it is possible for FT-Top2 and other baselines to further have higher accuracy if being fine-tuned for more than 5 epochs. However,  these baseline methods generally involve more computations in each epoch but achieve similar or lower model accuracy compared to GreenTrainer. For example, in Table 2, when fine-tuning a BLOOMZ on DialogSum dataset for 5 epochs, GT-0.5 takes 3.1 hours to reach R1=24.9% while FT-Top2 takes 4.6 hours to reach R1=22.1%. This demonstrates that GreenTrainer has better accuracy-compute efficiency. We also did more experiments by fine-tuning the OPT-2.7B model for 10 epochs on the webquestions dataset. The model accuracy for the generative QA tasks is listed below: 

|         | Accuracy after 5 epochs | Accuracy after 10 epochs | validation loss in the first 5 epochs |
|---------|-------------------------|--------------------------|---------------------------------------|
| FT-Top2 | 6.3%                    | 7.1%                     | [2.907, 2.805, 2.667, 2.626, 2.616]   |
| LoRA    | 19.6%                   | 21.9%                    | [1.765, 1.672, 1.646, 1.624, 1.620]   |

In comparison, GT-0.5 has an accuracy of 28.7% after 5 epochs of fine-tuning. Since the model accuracy in baselines increases slowly with more epochs, we believe that such advantage of GreenTrainer will hold with more training epochs.

Reviewer comment (Reproducibility):
How the running time is measured? There is a lack of hardware specifications.

Author response:
We measured the end-to-end fine-tuning time on a Lambda Lab cloud instance with 24 vCPUs and a single H100 80GB GPU. Note that, in all experiments, the measured fine-tuning time covers the entire fine-tuning procedure, include the overhead of our importance evaluation and dynamic programming.

Reviewer comment (Evaluation):
The main concern is the vagueness on the efficiency or costs of identifying important tensors during finetuning. This makes it difficult to judge whether or not this is a practical approach. The tensor selection is simply claimed to be O(N), but this is insufficient to understand the runtime costs.

Author response:
In general, the end-to-end computing time of standard LLM fine-tuning consists of the following components: 1) doing the forward pass to compute training loss; 2) backpropagating the loss gradients to the trainable tensors and 3) compute the updates to the tensors. In GreenTrainer, we add two extra components: 1) evaluation of tensor importance; 2) dynamic programming to find the optimal tensor selection. Our major contribution in this paper is that we are the first to explicitly incorporate the backpropagation cost, i.e., the components 2) and 3) listed above that are the dominant factor of the fine-tuning cost, into the selection of trainable tensors. In this way, we can ensure that the selected set of tensors can guarantee the required reduction of fine-tuning cost with the minimum impact on the trained model accuracy. In contrast, most existing LLM fine-tuning works adopt indirect objectives, such as the parameter count and per-layer forward pass FLOPs, to guide tensor selection. Without a proper backpropagation cost model, the tensor selection of existing works can lead to either insufficient FLOPs reduction or low accuracy of the fine-tuned model. 

In all experiments, we measured the end-to-end computing time of LLM fine-tuning, by including all the components described above for standard LLM fine-tuning and the two extra components we added. 

Our experiment results showed that our tensor importance evaluation and DP are both very lightweight and introduce little extra computing overhead. For tensor importance evaluation, it by design incurs only negligible computing overhead, because it only takes one extra training iteration to obtain gradients for all model parameters, as described in Section 3.1 of the paper. To further verify such low overhead, we conducted extra experiments and the results show that, even when we repeatedly perform such tensor importance evaluation in each training epoch at runtime to capture the possible variations of tensor importance at runtime, the overhead of importance evaluation is only <0.2% on SciTLDR dataset and <0.01% on DialogSum dataset, with respect to the entire fine-tuning cost. In Appendix A.1, we also showed that, by involving a smaller set of tensors into training, we can further reduce the amount of memory usage by tensor importance evaluation by another 50%.

Similarly, in the Appendix A.2, we also showed that the computing cost of DP is very low, and such computing costs can be further reduced by exploiting different tradeoffs between the model accuracy and cost. For example, by reducing the resolution of solving the DP, we can further reduce the wall-clock time of running DP by another 100x, with negligible lost on the model accuracy.

Reviewer comment (Evaluation):
What is the e2e finetuning time of GreenTrainer compared to baselines.

Author response:
In general, the end-to-end computing time of standard LLM fine-tuning consists of the following components: 1) doing the forward pass to compute training loss; 2) backpropagating the loss gradients to the trainable tensors and 3) compute the updates to the tensors. In GreenTrainer, we add two extra components: 1) evaluation of tensor importance; 2) dynamic programming to find the optimal tensor selection. Our major contribution in this paper is that we are the first to explicitly incorporate the backpropagation cost, i.e., the components 2) and 3) listed above that are the dominant factor of the fine-tuning cost, into the selection of trainable tensors. In this way, we can ensure that the selected set of tensors can guarantee the required reduction of fine-tuning cost with the minimum impact on the trained model accuracy. In contrast, most existing LLM fine-tuning works adopt indirect objectives, such as the parameter count and per-layer forward pass FLOPs, to guide tensor selection. Without a proper backpropagation cost model, the tensor selection of existing works can lead to either insufficient FLOPs reduction or low accuracy of the fine-tuned model. 

In all experiments, we measured the end-to-end computing time of LLM fine-tuning, by including all the components described above for standard LLM fine-tuning and the two extra components we added. 

Our experiment results showed that our tensor importance evaluation and DP are both very lightweight and introduce little extra computing overhead. For tensor importance evaluation, it by design incurs only negligible computing overhead, because it only takes one extra training iteration to obtain gradients for all model parameters, as described in Section 3.1 of the paper. To further verify such low overhead, we conducted extra experiments and the results show that, even when we repeatedly perform such tensor importance evaluation in each training epoch at runtime to capture the possible variations of tensor importance at runtime, the overhead of importance evaluation is only <0.2% on SciTLDR dataset and <0.01% on DialogSum dataset, with respect to the entire fine-tuning cost. In Appendix A.1, we also showed that, by involving a smaller set of tensors into training, we can further reduce the amount of memory usage by tensor importance evaluation by another 50%.

Similarly, in the Appendix A.2, we also showed that the computing cost of DP is very low, and such computing costs can be further reduced by exploiting different tradeoffs between the model accuracy and cost. For example, by reducing the resolution of solving the DP, we can further reduce the wall-clock time of running DP by another 100x, with negligible lost on the model accuracy.

Reviewer comment (Reproducibility):
How frequently are important tensors identified? Is it per iteration, per epoch (as suggested by Figure 1), or some other intervals? What is the frequency for the experimental results?

Author response:
In all experiments, we evaluated tensor importance per epoch. As discussed above, our tensor importance evaluation is very lightweight, and even doing so per epoch produces very little extra computing overhead at runtime.

We evaluate the importances of tensors at the beginning of every epoch. Yes, we take the same approach of importance evaluation for different epochs, and so in every epoch, we evaluate the tensor importance at the beginning of the epoch. In this way, we make sure that each epoch's tensor selection is made based on the up-to-date information of tensor importances.

Following the previous question, we would like to clarify that we only evaluate the tensor importance once in every epoch, more specifically, in the beginning of every epoch. The reasons are two-fold. First, frequently checking the number of iterations in each epoch to decide whether to re-evaluate tensor importance is very computationally expensive and will significantly slow down the training in real implementations. Second, in practical LLM fine-tuning, we usually use a small batch size (e.g., 4 in our implementation and evaluations) but a large number of iterations in each epoch. As a result, we expect only very small changes of tensor importances over different iterations in the same epoch. 

Due to these reasons, we only evaluate tensor importance once per epoch. Our experiment results show that such runtime evaluation incurs very little extra overhead (<0.2% of the entire fine-tuning cost).

Regarding your question about "whether more frequent tensor importance evaluation would bring accuracy benefits", we conducted some extra experiments by applying more frequent tensor importance evaluation within each epoch, using OPT-2.7B model on the webquestions dataset and generative QA task. The results are as follows:

| Frequency of tensor importance evaluation | Accuracy | Wall-clock training time |
|-------------------------------------------|----------|--------------------------|
| Every 945 iterations (once per epoch)     | 28.4%    | 0.50 h                   |
| Every 600 iterations                      | 28.5%    | 0.54h                    |
| Every 400 iterations                      | 28.2%    | 0.56h                    |
| Every 200 iterations                      | 27.5%    | 0.64h                    |

The results show that: 

1) more frequent tensor importance evaluation brings only very small improvement on the task accuracy. Considering the randomness in different training trials, we believe that such accuracy improvement is negligible, and the accuracy could even drop down by 1% when the frequency of evaluation is very high (every 200 iterations). We believe that this is due to accumulation of tensor importance evaluation and tensor selection errors, which stem from the first-order approximation in the tensor importance metric and the approximate solution in DP. Another possible reason is that the tensor importances are calculated over the training dataset, and too frequent tensor importance evaluation may make the training process overfit to the training dataset.

2) The training cost steadily increases with the frequency of tensor importance evaluation. When the interval of evaluation reduces from 945 iterations to 200 iterations, the training time increases by 28%.

In summary, performing more frequent tensor importance evaluation within each epoch brings little improvement on the task accuracy but noticeably increase the training cost. We hope that you find our experiment results and analysis reasonable and useful. Thanks again!

Reviewer comment (Reproducibility):
What is the `m` binary vector, or at least the number of important tensors, for the experimental results?

Author response:
The binary vector $\mathbf{m}$ represents the selection of tensors being involved in fine-tuning: the $i$-th component of the vector $\mathbf{m}$ is valued in [0, 1], as 1 indicates that tensor i is selected to be fine-tuned and 0 indicates otherwise. In our approach, we first evaluate the importances of different tensors at runtime, and then solve $\mathbf{m}$ using DP based on our problem formulation in Eq. (5). Note that the solution to $\mathbf{m}$ is not constant throughout the training procedure, because the importances of tensors may change as the training progresses, and we run DP once per epoch to make sure that the tensor selection (represented by $\mathbf{m}$) always reflect the tensors’ current contribution to training.